# okc: problem=muse_bf16_proj
"""Small-batch split router; packed hi/lo columns for larger router GEMMs."""

from typing import TypedDict

import torch


class Weights(TypedDict):
    qkv_gate: torch.Tensor  # bfloat16 [2304, 8192]
    o: torch.Tensor         # bfloat16 [8192, 1024]
    pooler: torch.Tensor    # bfloat16 [4096, 8192]
    unpooler: torch.Tensor  # bfloat16 [8192, 4096]
    router: torch.Tensor    # float32 [256, 8192]


class Config(TypedDict):
    m: int
    n: int      # output width; the unpooler's is 1024 up to 8 tokens (this GPU's slice), else 8192
    k: int
    layer: str  # which weight: "qkv_gate", "o", "pooler", "unpooler" or "router"


class Inputs(TypedDict):
    x: torch.Tensor  # bfloat16 [m, k]
    config: Config


import triton
import triton.language as tl
from sglang.srt.layers.quantization.unquant import _bf16_gemm_dispatch_impl


@triton.jit
def _gemv(X, W, Y, N: tl.constexpr, K: tl.constexpr,
          BN: tl.constexpr, BK: tl.constexpr):
    # One CTA owns several output rows and reduces the complete input width.
    rows = tl.program_id(0) * BN + tl.arange(0, BN)
    cols = tl.arange(0, BK)
    x = tl.load(X + cols, cols < K, other=0).to(tl.float32)
    w = tl.load(W + rows[:, None] * K + cols[None, :],
                (rows[:, None] < N) & (cols[None, :] < K), other=0).to(tl.float32)
    out = tl.sum(w * x[None, :], axis=1)
    tl.store(Y + rows, out, rows < N)


@triton.jit
def _router(X, HI, LO, Y, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
            BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, SPLIT: tl.constexpr):
    rows = tl.program_id(0) * BM + tl.arange(0, BM)
    cols = tl.program_id(1) * BN + tl.arange(0, BN)
    kk = tl.arange(0, BK)
    hi_acc = tl.full((BM, BN), 0, tl.float32)
    lo_acc = tl.full((BM, BN), 0, tl.float32)
    for start in range(tl.cdiv(K, BK * SPLIT)):
        k = (start * SPLIT + tl.program_id(2)) * BK + kk
        x = tl.load(X + rows[:, None] * K + k[None, :],
                    (rows[:, None] < M) & (k[None, :] < K), other=0)
        offsets = k[:, None] * N + cols[None, :]
        mask = (k[:, None] < K) & (cols[None, :] < N)
        hi = tl.load(HI + offsets, mask, other=0)
        lo = tl.load(LO + offsets, mask, other=0)
        hi_acc += tl.dot(x, hi)
        lo_acc += tl.dot(x, lo)
    tl.store(Y + tl.program_id(2) * M * N + rows[:, None] * N + cols[None, :], hi_acc + lo_acc,
             (rows[:, None] < M) & (cols[None, :] < N))


@triton.jit
def _sum_splits(P, Y, SIZE: tl.constexpr, SPLIT: tl.constexpr, BLOCK: tl.constexpr):
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    total = tl.full((BLOCK,), 0, tl.float32)
    for split in range(SPLIT):
        total += tl.load(P + split * SIZE + offsets, offsets < SIZE, other=0)
    tl.store(Y + offsets, total, offsets < SIZE)


def prepare(weights: Weights):
    state = {name: w.contiguous() for name, w in weights.items()}
    hi = state["router"].to(torch.bfloat16)
    lo = (state["router"] - hi.float()).to(torch.bfloat16)
    state["router_hi_t"] = hi.t().contiguous()
    state["router_lo_t"] = lo.t().contiguous()
    state["router_pair_t"] = torch.cat((state["router_hi_t"], state["router_lo_t"]), dim=1)
    return state


def kernel_fn(data: Inputs, state) -> torch.Tensor:
    x, cfg = data["x"], data["config"]
    m, n, k = x.shape[0], cfg["n"], cfg["k"]
    layer = cfg["layer"]
    if m == 1 and layer in ("o", "unpooler", "pooler"):
        out = torch.empty((m, n), device=x.device,
                          dtype=torch.float32 if layer == "router" else torch.bfloat16)
        _gemv[(triton.cdiv(n, 4),)](x, state[layer], out, n, k, 4,
                                     triton.next_power_of_2(k), num_warps=8)
        return out
    if layer == "router" and m > 32:
        # One BF16 GEMM computes both FP32 products, then combines the pair.
        pair = torch.mm(x, state["router_pair_t"], out_dtype=torch.float32)
        return pair[:, :n] + pair[:, n:]
    if layer == "router":
        out = torch.empty((m, n), device=x.device, dtype=torch.float32)
        split = 16 if m <= 32 else 8
        partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32)
        bm = 16 if m <= 32 else 32
        _router[(triton.cdiv(m, bm), triton.cdiv(n, 32), split)](
            x, state["router_hi_t"], state["router_lo_t"], partial,
            m, n, k, bm, 32, 64, split, num_warps=4)
        _sum_splits[(triton.cdiv(m * n, 256),)](partial, out, m * n, split, 256)
        return out
    return _bf16_gemm_dispatch_impl(x, state[layer][:n], None)


"""Synthetic-input search for small-batch, direct FP32 vector reductions."""

@triton.jit
def _vector_rows(X, W, Y, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                 BN: tl.constexpr, BK: tl.constexpr):
    rows = tl.program_id(0) * BN + tl.arange(0, BN)
    cols = tl.arange(0, BK)
    # Load the original weight exactly once, reuse it for several input tokens.
    # Router weights remain FP32; all other weights remain exact BF16.
    w = tl.load(W + rows[:, None] * K + cols[None, :],
                (rows[:, None] < N) & (cols[None, :] < K), other=0).to(tl.float32)
    for token in range(M):
        x = tl.load(X + token * K + cols, cols < K, other=0).to(tl.float32)
        dot = tl.sum(w * x[None, :], axis=1)
        tl.store(Y + token * N + rows, dot, rows < N)


_vector_base_prepare = prepare
_vector_base_kernel = kernel_fn


def _vector_measure(fn, flush):
    fn()
    fn()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        fn()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    elapsed = []
    for _ in range(7):
        flush.zero_()
        start.record()
        graph.replay()
        end.record()
        end.synchronize()
        elapsed.append(start.elapsed_time(end))
    return sorted(elapsed)[3]


def prepare(weights: Weights):
    state = _vector_base_prepare(weights)
    device = weights["qkv_gate"].device
    flush = torch.empty(128 * 1024 * 1024, device=device, dtype=torch.uint8)
    state["vector_choices"] = {}
    for layer in ("qkv_gate", "o", "pooler", "unpooler", "router"):
        full_n, k = weights[layer].shape
        # Reuse benefits tiny batches. Larger batches retain tensor-core GEMMs.
        for m in (1, 8):
            n = 1024 if layer == "unpooler" and m <= 8 else full_n
            dtype = torch.float32 if layer == "router" else torch.bfloat16
            x = torch.ones((m, k), device=device, dtype=torch.bfloat16)
            out = torch.empty((m, n), device=device, dtype=dtype)
            data = dict(x=x, config=dict(layer=layer, m=m, n=n, k=k))
            prior = _vector_measure(lambda: _vector_base_kernel(data, state), flush)
            best_ms, choice = prior * 0.97, None
            for bn, warps in ((1, 4), (1, 8), (2, 4), (2, 8), (4, 8), (8, 8)):
                def op():
                    _vector_rows[(triton.cdiv(n, bn),)](
                        x, weights[layer], out, m, n, k, bn,
                        triton.next_power_of_2(k), num_warps=warps)
                ms = _vector_measure(op, flush)
                if ms < best_ms:
                    best_ms, choice = ms, (bn, warps)
            if choice is not None:
                state["vector_choices"][(layer, m, n, k)] = choice
            print("VECTOR_PLAN", layer, m, choice, "synthetic_ms", best_ms,
                  "prior_ms", prior, flush=True)
    return state


def kernel_fn(data: Inputs, state) -> torch.Tensor:
    x, cfg = data["x"], data["config"]
    layer, m, n, k = cfg["layer"], x.shape[0], cfg["n"], cfg["k"]
    choice = state["vector_choices"].get((layer, m, n, k))
    if choice is None:
        return _vector_base_kernel(data, state)
    bn, warps = choice
    out = torch.empty((m, n), device=x.device,
                      dtype=torch.float32 if layer == "router" else torch.bfloat16)
    _vector_rows[(triton.cdiv(n, bn),)](x, state[layer], out, m, n, k, bn,
                                      triton.next_power_of_2(k), num_warps=warps)
    return out