# 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)


"""Pack hi/lo columns together and fuse their sum into the GEMM epilogue."""

@triton.jit
def _paired_router(X, W, P, 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)
    pairs = tl.program_id(1) * (BN * 2) + tl.arange(0, BN * 2)
    kk = tl.arange(0, BK)
    acc = tl.full((BM, BN * 2), 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)
        w = tl.load(W + k[:, None] * (2 * N) + pairs[None, :],
                    (k[:, None] < K) & (pairs[None, :] < 2 * N), other=0)
        acc += tl.dot(x, w)
    # Adjacent columns are the high and low products of one original weight.
    # They stay separate FP32 accumulators until this final FP32 addition.
    result = tl.sum(tl.reshape(acc, (BM, BN, 2)), axis=2)
    cols = tl.program_id(1) * BN + tl.arange(0, BN)
    tl.store(P + tl.program_id(2) * M * N + rows[:, None] * N + cols[None, :],
             result, (rows[:, None] < M) & (cols[None, :] < N))


_paired_base_prepare = prepare
_paired_base_kernel = kernel_fn


def _paired_run(x, state, out, partial, config):
    bm, bn, bk, split = config
    m, k = x.shape
    n = state["router"].shape[0]
    _paired_router[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, state["router_interleaved"], out if split == 1 else partial,
        m, n, k, bm, bn, bk, split, num_warps=4)
    if split > 1:
        _sum_splits[(triton.cdiv(m * n, 256),)](partial, out, m * n, split, 256)


def _paired_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 = _paired_base_prepare(weights)
    n, k = weights["router"].shape
    state["router_interleaved"] = torch.stack(
        (state["router_hi_t"], state["router_lo_t"]), dim=-1).reshape(k, 2 * n).contiguous()
    state["paired_choices"] = {}
    device = weights["router"].device
    flush = torch.empty(128 * 1024 * 1024, device=device, dtype=torch.uint8)
    for m in (1, 8, 16, 32, 64, 128, 256):
        x = torch.ones((m, k), device=device, dtype=torch.bfloat16)
        out = torch.empty((m, n), device=device, dtype=torch.float32)
        partial = torch.empty((16, m, n), device=device, dtype=torch.float32)
        data = dict(x=x, config=dict(layer="router", m=m, n=n, k=k))
        prior = _paired_measure(lambda: _paired_base_kernel(data, state), flush)
        best_ms, choice = prior * 0.97, None
        for bm, bn, bk in ((16, 16, 256), (16, 32, 128), (32, 32, 256)):
            tiles = triton.cdiv(m, bm) * triton.cdiv(n, bn)
            split = min(16, triton.next_power_of_2(triton.cdiv(128, tiles)))
            config = (bm, bn, bk, split)
            ms = _paired_measure(lambda: _paired_run(x, state, out, partial, config), flush)
            if ms < best_ms:
                best_ms, choice = ms, config
        if choice is not None:
            state["paired_choices"][m] = choice
        print("PAIRED_PLAN", 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"]
    m, n = x.shape[0], cfg["n"]
    choice = state["paired_choices"].get(m) if cfg["layer"] == "router" else None
    if choice is None:
        return _paired_base_kernel(data, state)
    out = torch.empty((m, n), device=x.device, dtype=torch.float32)
    split = choice[3]
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32) if split > 1 else out
    _paired_run(x, state, out, partial, choice)
    return out