# okc: problem=muse_bf16_proj
"""Validated router plus measured dense-kernel choices across broad batch-size ranges."""

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)


@triton.jit
def _dense_split(X, W, P, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                 BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, SPLIT: tl.constexpr):
    rm = tl.program_id(0)*BM + tl.arange(0, BM)
    rn = tl.program_id(1)*BN + tl.arange(0, BN)
    kk = tl.arange(0, BK)
    acc = tl.full((BM, BN), 0, tl.float32)
    for it in range(tl.cdiv(K, BK*SPLIT)):
        rk = (it*SPLIT + tl.program_id(2))*BK + kk
        a = tl.load(X + rm[:, None]*K + rk[None, :], (rm[:, None]<M)&(rk[None, :]<K), 0)
        b = tl.load(W + rn[None, :]*K + rk[:, None], (rn[None, :]<N)&(rk[:, None]<K), 0)
        acc += tl.dot(a,b)
    tl.store(P + tl.program_id(2)*M*N + rm[:, None]*N + rn[None, :], acc, (rm[:, None]<M)&(rn[None, :]<N))


DENSE_CHOICES = {'qkv_gate|single': 'hopper', 'pooler|small': 'transposed', 'pooler|large': 'transposed'}

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)
    if "transposed" in DENSE_CHOICES.values():
        state["dense_t"] = {name: state[name].t().contiguous()
                            for name in ("qkv_gate", "o", "pooler", "unpooler")}
    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"]
    family = "single" if m == 1 else ("small" if m <= 32 else "large")
    method = DENSE_CHOICES.get(layer + "|" + family, "base")
    if layer != "router":
        if method == "hopper":
            from sglang.kernels.ops.gemm.hopper_bf16_gemv import hopper_bf16_gemv, use_hopper_bf16_gemv
            if use_hopper_bf16_gemv(m, n, k):
                return hopper_bf16_gemv(x, state[layer][:n])
        elif method == "transposed":
            return torch.mm(x, state["dense_t"][layer][:, :n])
        elif method == "wide_k" and 2 <= m <= 64:
            split, bm, bn, bk = 4, (16 if m <= 16 else 32), 64, 256
            partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32)
            out = torch.empty((m, n), device=x.device, dtype=torch.bfloat16)
            _dense_split[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
                x, state[layer], partial, m, n, k, bm, bn, bk, split,
                num_warps=4, num_stages=3)
            _sum_splits[(triton.cdiv(m*n, 256),)](partial, out, m*n, split, 256)
            return out
        elif method == "quack_large" and m >= 64:
            from quack.gemm import gemm as quack_gemm
            from quack.gemm_config import SplitKMode
            bm, bn = 64, 128
            tiles = triton.cdiv(m, bm) * triton.cdiv(n, bn)
            split = min(8, triton.next_power_of_2(max(1, triton.cdiv(132, tiles))))
            out = torch.empty((m, n), device=x.device, dtype=torch.bfloat16)
            quack_gemm(x.unsqueeze(0), state[layer][:n].unsqueeze(0), out.unsqueeze(0), None, None,
                       bm, bn, 1, 1, pingpong=False, persistent=True,
                       is_dynamic_persistent=False, split_k=split,
                       split_k_mode=SplitKMode.SEPARATE)
            return out
        elif method == "quack_small" and 2 <= m <= 32:
            from quack.gemm import gemm as quack_gemm
            from quack.gemm_config import SplitKMode
            bm, bn = 64, (16 if m <= 16 else 32)
            split = min(8, triton.next_power_of_2(max(1, triton.cdiv(132, triton.cdiv(n, bm)))))
            out = torch.empty((m, n), device=x.device, dtype=torch.bfloat16)
            quack_gemm(state[layer][:n].unsqueeze(0), x.unsqueeze(0), out.t().unsqueeze(0), None, None,
                       bm, bn, 1, 1, pingpong=False, persistent=True,
                       is_dynamic_persistent=False, split_k=split,
                       split_k_mode=SplitKMode.SEPARATE)
            return out
    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)