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


"""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"):
        if not _search_budget_open():
            break
        full_n, k = weights[layer].shape
        # Reuse benefits tiny batches. Larger batches retain tensor-core GEMMs.
        for m in (1, 8):
            if not _search_budget_open():
                break
            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)):
                if not _search_budget_open():
                    break
                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


"""Tune exact hi/lo router GEMMs on synthetic activations during preparation."""

@triton.jit
def _router_tunable(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, K_CONTIG: 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)
    high = tl.full((BM, BN), 0, tl.float32)
    low = 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)
        if K_CONTIG:
            offsets = cols[None, :] * K + k[:, None]
        else:
            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)
        high += tl.dot(x, hi)
        low += tl.dot(x, lo)
    tl.store(Y + tl.program_id(2) * M * N + rows[:, None] * N + cols[None, :],
             high + low, (rows[:, None] < M) & (cols[None, :] < N))


_router_base_prepare = prepare
_router_base_kernel = kernel_fn


def _router_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 _router_run(x, state, out, partial, config):
    bm, bn, bk, split, k_contig = config
    m, k = x.shape
    n = state["router"].shape[0]
    high = state["router_hi_k"] if k_contig else state["router_hi_t"]
    low = state["router_lo_k"] if k_contig else state["router_lo_t"]
    _router_tunable[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, high, low, out if split == 1 else partial,
        m, n, k, bm, bn, bk, split, k_contig, num_warps=4)
    if split > 1:
        _sum_splits[(triton.cdiv(m * n, 256),)](partial, out, m * n, split, 256)


def prepare(weights: Weights):
    state = _router_base_prepare(weights)
    state["router_hi_k"] = state["router_hi_t"].t().contiguous()
    state["router_lo_k"] = state["router_lo_t"].t().contiguous()
    state["router_choices"] = {}
    device = weights["router"].device
    n, k = weights["router"].shape
    flush = torch.empty(128 * 1024 * 1024, device=device, dtype=torch.uint8)
    for m in (16, 32, 64, 128, 256):
        if not _search_budget_open():
            break
        x = torch.ones((m, k), dtype=torch.bfloat16, device=device)
        out = torch.empty((m, n), dtype=torch.float32, device=device)
        partial = torch.empty((8, m, n), dtype=torch.float32, device=device)
        data = dict(x=x, config=dict(m=m, n=n, k=k, layer="router"))
        prior = _router_measure(lambda: _router_base_kernel(data, state), flush)
        best_ms, choice = prior * 0.97, None
        # Each trial changes only tiling, layout and the FP32 reduction partition.
        tiles = ((16, 32, 128, 4), (16, 32, 256, 4), (16, 32, 256, 8),
                 (16, 16, 256, 1), (32, 32, 128, 2), (32, 32, 256, 2),
                 (32, 32, 256, 4), (32, 64, 128, 2))
        for tile in tiles:
            if not _search_budget_open():
                break
            for k_contig in (False, True):
                if not _search_budget_open():
                    break
                cfg = (*tile, k_contig)
                ms = _router_measure(lambda: _router_run(x, state, out, partial, cfg), flush)
                if ms < best_ms:
                    best_ms, choice = ms, cfg
        if choice is not None:
            state["router_choices"][m] = choice
        print("ROUTER_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["router_choices"].get(m) if cfg["layer"] == "router" else None
    if choice is None:
        return _router_base_kernel(data, state)
    out = torch.empty((m, n), dtype=torch.float32, device=x.device)
    split = choice[3]
    partial = torch.empty((split, m, n), dtype=torch.float32, device=x.device) if split > 1 else out
    _router_run(x, state, out, partial, choice)
    return out


"""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):
        if not _search_budget_open():
            break
        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)):
            if not _search_budget_open():
                break
            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


import time as _preparation_clock
_preparation_deadline = float("inf")


def _search_budget_open():
    return _preparation_clock.monotonic() < _preparation_deadline


_unbounded_prepare = prepare


def prepare(weights: Weights):
    global _preparation_deadline
    # Time only setup/search; leave 90 seconds of the official preparation cap
    # for a final in-flight compilation and remaining lossless weight packing.
    _preparation_deadline = _preparation_clock.monotonic() + 210.0
    return _unbounded_prepare(weights)