# okc: problem=muse_bf16_proj
"""Per-shape choice among the projection methods that passed the judge (004-027).

Each layer has a plan: token-count ranges, each mapped to one method and its launch
parameters, chosen by timing every method on an H200 (gpu_lab/msweep.py). Dispatch
depends only on the layer and the token count. Weights stay at their own precision: the
dense BF16 weights are used as given or as a lossless transpose, and the router's FP32
weights directly or as the BF16 hi + lo pair with FP32 accumulation.
"""

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
from sglang.kernels.ops.gemm.hopper_bf16_gemv import hopper_bf16_gemv, use_hopper_bf16_gemv

DENSE = ("qkv_gate", "o", "pooler", "unpooler")
# layer -> [(largest token count, method, params)], first match wins; beyond the last
# entry the layer uses its final method. Filled from msweep.py's measurements.
PLAN = {
    "qkv_gate": [(1, "vector", (1, 4)), (16, "widek", (8, 16, 64, 128, 4, 3)), (1 << 30, "cublas", ())],
    "o": [(1, "vector", (2, 4)), (1 << 30, "cublas", ())],
    "pooler": [(1, "vector", (1, 8)), (16, "splitf", (16, 32, 512, 1, 2, 4)),
               (32, "splitf", (32, 32, 512, 1, 2, 3)), (128, "trans", ()), (1 << 30, "cublas", ())],
    "unpooler": [(1, "vector", (1, 8)), (8, "splitf", (16, 32, 512, 4, 2, 4)), (1 << 30, "trans", ())],
    "router": [(1, "rvector", (1, 8)), (8, "rpair", (16, 32, 128, 16, 4, 4)), (16, "rpair", (16, 64, 64, 32, 4, 4)),
               (32, "rpair", (64, 32, 128, 16)), (128, "rpair", (64, 32, 128, 8, 4, 4)),
               (1 << 30, "rpair", (128, 32, 128, 8, 4, 4))],
}


@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 _vector_rows(X, W, Y, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                 BN: tl.constexpr, BK: tl.constexpr):
    # Load the weight rows once and reuse them for every token, in FP32.
    rows = tl.program_id(0) * BN + tl.arange(0, BN)
    cols = tl.arange(0, BK)
    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)


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


@triton.jit
def _split_fused(X, W, P, Y, C, M, N, K: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr,
                 BK: tl.constexpr, SPLIT: tl.constexpr, STAGES: tl.constexpr):
    # Split s reduces the contiguous K-chunk [s K/SPLIT, (s+1) K/SPLIT) of a BM x BN tile in
    # FP32. With several splits, the last to arrive sums all partials in a fixed order,
    # writes the BF16 tile and resets the tile's counter for the next call.
    pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    kk = tl.arange(0, BK)
    chunk: tl.constexpr = K // SPLIT
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for start in tl.range(0, chunk, BK, num_stages=STAGES):
        rk = pid_k * chunk + start + kk
        a = tl.load(X + rm[:, None] * K + rk[None, :], mask=rm[:, None] < M, other=0)
        b = tl.load(W + rn[None, :] * K + rk[:, None], mask=rn[None, :] < N, other=0)
        acc += tl.dot(a, b)
    out_mask = (rm[:, None] < M) & (rn[None, :] < N)
    if SPLIT == 1:
        tl.store(Y + rm[:, None] * N + rn[None, :], acc.to(tl.bfloat16), mask=out_mask)
    else:
        tl.store(P + pid_k * M * N + rm[:, None] * N + rn[None, :], acc, mask=out_mask)
        tile = pid_m * tl.num_programs(1) + pid_n
        if tl.atomic_add(C + tile, 1, sem="acq_rel") == SPLIT - 1:
            total = tl.zeros((BM, BN), dtype=tl.float32)
            for s in range(SPLIT):
                total += tl.load(P + s * M * N + rm[:, None] * N + rn[None, :], mask=out_mask,
                                 other=0, cache_modifier=".cg")
            tl.store(Y + rm[:, None] * N + rn[None, :], total.to(tl.bfloat16), mask=out_mask)
            tl.atomic_xchg(C + tile, 0)


@triton.jit
def _split_fused_t(X, WT, P, Y, C, M, N, LDW, K: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr,
                   BK: tl.constexpr, SPLIT: tl.constexpr, STAGES: tl.constexpr):
    # As _split_fused, reading the lossless transpose WT [K, N] (row stride LDW): B tiles
    # arrive row-major for tl.dot with N contiguous.
    pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    kk = tl.arange(0, BK)
    chunk: tl.constexpr = K // SPLIT
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for start in tl.range(0, chunk, BK, num_stages=STAGES):
        rk = pid_k * chunk + start + kk
        a = tl.load(X + rm[:, None] * K + rk[None, :], mask=rm[:, None] < M, other=0)
        b = tl.load(WT + rk[:, None] * LDW + rn[None, :], mask=rn[None, :] < N, other=0)
        acc += tl.dot(a, b)
    out_mask = (rm[:, None] < M) & (rn[None, :] < N)
    if SPLIT == 1:
        tl.store(Y + rm[:, None] * N + rn[None, :], acc.to(tl.bfloat16), mask=out_mask)
    else:
        tl.store(P + pid_k * M * N + rm[:, None] * N + rn[None, :], acc, mask=out_mask)
        tile = pid_m * tl.num_programs(1) + pid_n
        if tl.atomic_add(C + tile, 1, sem="acq_rel") == SPLIT - 1:
            total = tl.zeros((BM, BN), dtype=tl.float32)
            for s in range(SPLIT):
                total += tl.load(P + s * M * N + rm[:, None] * N + rn[None, :], mask=out_mask,
                                 other=0, cache_modifier=".cg")
            tl.store(Y + rm[:, None] * N + rn[None, :], total.to(tl.bfloat16), mask=out_mask)
            tl.atomic_xchg(C + tile, 0)


@triton.jit
def _streamk(X, W, Y, P, C, M, N, K: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
             TILES_N, UNITS, STAGES: tl.constexpr):
    # Stream-K: the UNITS = tiles x K-steps of work are split into equal contiguous ranges,
    # one per program, so every SM gets the same share whatever the tile count. A tile
    # covered by one program is stored directly; a shared tile's programs leave FP32
    # partials (slot 0 for a program's first tile, 1 for its last), and the last to finish
    # sums them in program order, stores the BF16 tile and resets the tile's counter.
    KI: tl.constexpr = K // BK
    pid = tl.program_id(0)
    progs = tl.num_programs(0)
    unit = (pid * UNITS) // progs
    end = ((pid + 1) * UNITS) // progs
    first_tile = unit // KI
    while unit < end:
        tile = unit // KI
        seg_end = tl.minimum(end, (tile + 1) * KI)
        pid_m = tile // TILES_N
        pid_n = tile % TILES_N
        rm = pid_m * BM + tl.arange(0, BM)
        rn = pid_n * BN + tl.arange(0, BN)
        acc = tl.zeros((BM, BN), dtype=tl.float32)
        for u in tl.range(unit, seg_end, num_stages=STAGES):
            rk = (u - tile * KI) * BK + tl.arange(0, BK)
            a = tl.load(X + rm[:, None] * K + rk[None, :], mask=rm[:, None] < M, other=0)
            b = tl.load(W + rn[None, :] * K + rk[:, None], mask=rn[None, :] < N, other=0)
            acc += tl.dot(a, b)
        out_mask = (rm[:, None] < M) & (rn[None, :] < N)
        if (unit == tile * KI) & (seg_end == (tile + 1) * KI):
            tl.store(Y + rm[:, None] * N + rn[None, :], acc.to(tl.bfloat16), mask=out_mask)
        else:
            slot = tl.where(tile == first_tile, 0, 1)
            tl.store(P + ((pid * 2 + slot) * BM + tl.arange(0, BM)[:, None]) * BN + tl.arange(0, BN)[None, :], acc)
            done = tl.atomic_add(C + tile, seg_end - unit, sem="acq_rel") + (seg_end - unit)
            if done == KI:
                q_lo = ((tile * KI + 1) * progs + UNITS - 1) // UNITS - 1
                q_hi = (((tile + 1) * KI) * progs + UNITS - 1) // UNITS - 1
                total = tl.zeros((BM, BN), dtype=tl.float32)
                for q in range(q_lo, q_hi + 1):
                    q_slot = tl.where(((q * UNITS) // progs) // KI == tile, 0, 1)
                    total += tl.load(P + ((q * 2 + q_slot) * BM + tl.arange(0, BM)[:, None]) * BN
                                     + tl.arange(0, BN)[None, :], cache_modifier=".cg")
                tl.store(Y + rm[:, None] * N + rn[None, :], total.to(tl.bfloat16), mask=out_mask)
                tl.atomic_xchg(C + tile, 0)
        unit = seg_end


@triton.jit
def _router_small(X, W, Y, M, N: tl.constexpr, K: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
                  MP: tl.constexpr, STAGES: tl.constexpr):
    # FP32 weights and FP32 accumulation for every token at once: one pass over the rows,
    # one reduction per K-chunk, no split partials.
    rows = tl.program_id(0) * BN + tl.arange(0, BN)
    toks = tl.arange(0, MP)
    acc = tl.zeros((BN, MP), dtype=tl.float32)
    for k0 in tl.range(0, K, BK, num_stages=STAGES):
        cols = k0 + tl.arange(0, BK)
        w = tl.load(W + rows[:, None] * K + cols[None, :])
        x = tl.load(X + toks[None, :] * K + cols[:, None], mask=toks[None, :] < M, other=0).to(tl.float32)
        acc += tl.sum(w[:, :, None] * x[None, :, :], axis=1)
    tl.store(Y + toks[None, :] * N + rows[:, None], acc, mask=toks[None, :] < M)


@triton.jit
def _router_split(X, W, P, Y, C, M, N: tl.constexpr, K: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
                  MP: tl.constexpr, SPLIT: tl.constexpr, STAGES: tl.constexpr):
    # FP32 router weights, FP32 products and sums, on CUDA cores. Program (n, s) reduces rows
    # [n BN, (n+1) BN) over the s-th K-chunk with the whole chunk's loads in flight at once
    # (STAGES = chunk / BK), then the last split to finish sums the partials in order.
    pid_n, pid_k = tl.program_id(0), tl.program_id(1)
    rows = pid_n * BN + tl.arange(0, BN)
    toks = tl.arange(0, MP)
    chunk: tl.constexpr = K // SPLIT
    acc = tl.zeros((BN, MP), dtype=tl.float32)
    for k0 in tl.range(0, chunk, BK, num_stages=STAGES):
        cols = pid_k * chunk + k0 + tl.arange(0, BK)
        w = tl.load(W + rows[:, None] * K + cols[None, :])
        x = tl.load(X + toks[None, :] * K + cols[:, None], mask=toks[None, :] < M, other=0).to(tl.float32)
        acc += tl.sum(w[:, :, None] * x[None, :, :], axis=1)
    mask = toks[None, :] < M
    if SPLIT == 1:
        tl.store(Y + toks[None, :] * N + rows[:, None], acc, mask=mask)
    else:
        tl.store(P + pid_k * M * N + toks[None, :] * N + rows[:, None], acc, mask=mask)
        if tl.atomic_add(C + pid_n, 1, sem="acq_rel") == SPLIT - 1:
            total = tl.zeros((BN, MP), dtype=tl.float32)
            for s in range(SPLIT):
                total += tl.load(P + s * M * N + toks[None, :] * N + rows[:, None], mask=mask, other=0,
                                 cache_modifier=".cg")
            tl.store(Y + toks[None, :] * N + rows[:, None], total, mask=mask)
            tl.atomic_xchg(C + pid_n, 0)


@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 _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))


@triton.jit
def _paired_router_fused(X, W, P, Y, C, M, N: tl.constexpr, K: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr,
                         BK: tl.constexpr, SPLIT: tl.constexpr, STAGES: tl.constexpr):
    # _paired_router over contiguous K-chunks, with the split partials summed in a fixed
    # order by the last CTA to finish each tile, in the same launch.
    pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    rows = pid_m * BM + tl.arange(0, BM)
    pairs = pid_n * (BN * 2) + tl.arange(0, BN * 2)
    chunk: tl.constexpr = K // SPLIT
    acc = tl.zeros((BM, BN * 2), dtype=tl.float32)
    for start in tl.range(0, chunk, BK, num_stages=STAGES):
        k = pid_k * chunk + start + tl.arange(0, BK)
        x = tl.load(X + rows[:, None] * K + k[None, :], mask=rows[:, None] < M, other=0)
        w = tl.load(W + k[:, None] * (2 * N) + pairs[None, :])
        acc += tl.dot(x, w)
    result = tl.sum(tl.reshape(acc, (BM, BN, 2)), axis=2)
    cols = pid_n * BN + tl.arange(0, BN)
    mask = (rows[:, None] < M) & (cols[None, :] < N)
    if SPLIT == 1:
        tl.store(Y + rows[:, None] * N + cols[None, :], result, mask=mask)
    else:
        tl.store(P + pid_k * M * N + rows[:, None] * N + cols[None, :], result, mask=mask)
        tile = pid_m * tl.num_programs(1) + pid_n
        if tl.atomic_add(C + tile, 1, sem="acq_rel") == SPLIT - 1:
            total = tl.zeros((BM, BN), dtype=tl.float32)
            for s in range(SPLIT):
                total += tl.load(P + s * M * N + rows[:, None] * N + cols[None, :], mask=mask, other=0,
                                 cache_modifier=".cg")
            tl.store(Y + rows[:, None] * N + cols[None, :], total, mask=mask)
            tl.atomic_xchg(C + tile, 0)


@triton.jit
def _router_tiled(X, W, P, Y, C, M, N: tl.constexpr, K: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr,
                  BK: tl.constexpr, SPLIT: tl.constexpr, STAGES: tl.constexpr):
    # As _paired_router_fused, over W repacked [N / BN, K, 2 BN] so each CTA streams one
    # contiguous region instead of a 2 BN-column strip of every row.
    pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    rows = pid_m * BM + tl.arange(0, BM)
    pairs = tl.arange(0, BN * 2)
    chunk: tl.constexpr = K // SPLIT
    base = W + pid_n * K * (2 * BN)
    acc = tl.zeros((BM, BN * 2), dtype=tl.float32)
    for start in tl.range(0, chunk, BK, num_stages=STAGES):
        k = pid_k * chunk + start + tl.arange(0, BK)
        x = tl.load(X + rows[:, None] * K + k[None, :], mask=rows[:, None] < M, other=0)
        w = tl.load(base + k[:, None] * (2 * BN) + pairs[None, :])
        acc += tl.dot(x, w)
    result = tl.sum(tl.reshape(acc, (BM, BN, 2)), axis=2)
    cols = pid_n * BN + tl.arange(0, BN)
    mask = (rows[:, None] < M) & (cols[None, :] < N)
    if SPLIT == 1:
        tl.store(Y + rows[:, None] * N + cols[None, :], result, mask=mask)
    else:
        tl.store(P + pid_k * M * N + rows[:, None] * N + cols[None, :], result, mask=mask)
        tile = pid_m * tl.num_programs(1) + pid_n
        if tl.atomic_add(C + tile, 1, sem="acq_rel") == SPLIT - 1:
            total = tl.zeros((BM, BN), dtype=tl.float32)
            for s in range(SPLIT):
                total += tl.load(P + s * M * N + rows[:, None] * N + cols[None, :], mask=mask, other=0,
                                 cache_modifier=".cg")
            tl.store(Y + rows[:, None] * N + cols[None, :], total, mask=mask)
            tl.atomic_xchg(C + tile, 0)


@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 _empty(x, m, n, layer):
    return torch.empty((m, n), device=x.device, dtype=torch.float32 if layer == "router" else torch.bfloat16)


def _split_sum(partial, out, split):
    size = out.numel()
    _sum_splits[(triton.cdiv(size, 256),)](partial, out, size, split, 256)
    return out


def m_cublas(x, state, layer, n, k, p):
    return _bf16_gemm_dispatch_impl(x, state[layer][:n], None)


def m_hopper(x, state, layer, n, k, p):
    return hopper_bf16_gemv(x, state[layer][:n])


def m_gemv(x, state, layer, n, k, p):
    out = _empty(x, 1, n, layer)
    _gemv[(triton.cdiv(n, 4),)](x, state[layer], out, n, k, 4, triton.next_power_of_2(k), num_warps=8)
    return out


def m_vector(x, state, layer, n, k, p):
    bn, warps = p
    m = x.shape[0]
    out = _empty(x, m, n, layer)
    _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


def m_trans(x, state, layer, n, k, p):
    return torch.mm(x, state["t_" + layer][:, :n])


def m_widek(x, state, layer, n, k, p):
    split, bm, bn, bk, warps, stages = p
    m = x.shape[0]
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32)
    _dense_split[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, state[layer], partial, m, n, k, bm, bn, bk, split, num_warps=warps, num_stages=stages)
    return _split_sum(partial, _empty(x, m, n, layer), split)


def m_splitf(x, state, layer, n, k, p):
    bm, bn, bk, split, warps, stages = p
    m = x.shape[0]
    out = _empty(x, m, n, layer)
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32) if split > 1 else out
    _split_fused[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, state[layer], partial, out, state["tile_counters"], m, n, k, bm, bn, bk, split, stages,
        num_warps=warps)
    return out


def m_splitft(x, state, layer, n, k, p):
    bm, bn, bk, split, warps, stages = p
    m = x.shape[0]
    wt = state["t_" + layer]
    out = _empty(x, m, n, layer)
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32) if split > 1 else out
    _split_fused_t[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, wt, partial, out, state["tile_counters"], m, n, wt.stride(0), k, bm, bn, bk, split, stages,
        num_warps=warps)
    return out


def m_streamk(x, state, layer, n, k, p):
    bm, bn, bk, per_sm, warps, stages = p
    m = x.shape[0]
    tiles_n = triton.cdiv(n, bn)
    units = triton.cdiv(m, bm) * tiles_n * (k // bk)
    progs = min(units, SMS * per_sm)
    out = _empty(x, m, n, layer)
    partial = torch.empty((progs, 2, bm, bn), device=x.device, dtype=torch.float32)
    _streamk[(progs,)](x, state[layer], out, partial, state["tile_counters"], m, n, k, bm, bn, bk,
                       tiles_n, units, stages, num_warps=warps)
    return out


def m_rsmall(x, state, layer, n, k, p):
    bn, bk, warps, stages = p
    m = x.shape[0]
    out = _empty(x, m, n, layer)
    _router_small[(n // bn,)](x, state["router"], out, m, n, k, bn, bk, max(2, triton.next_power_of_2(m)),
                              stages, num_warps=warps)
    return out


def m_rfused(x, state, layer, n, k, p):
    bm, bn, bk, split, warps, stages = p
    m = x.shape[0]
    out = _empty(x, m, n, layer)
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32) if split > 1 else out
    _paired_router_fused[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, state["router_interleaved"], partial, out, state["tile_counters"], m, n, k, bm, bn, bk, split, stages,
        num_warps=warps)
    return out


def m_rtiled(x, state, layer, n, k, p):
    bm, bn, bk, split, warps, stages = p
    m = x.shape[0]
    out = _empty(x, m, n, layer)
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32) if split > 1 else out
    _router_tiled[(triton.cdiv(m, bm), n // bn, split)](
        x, state[f"router_tiled_{bn}"], partial, out, state["tile_counters"], m, n, k, bm, bn, bk, split, stages,
        num_warps=warps)
    return out


def m_rsplit(x, state, layer, n, k, p):
    bn, bk, split, warps = p
    m = x.shape[0]
    out = _empty(x, m, n, layer)
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32) if split > 1 else out
    stages = max(1, min(8, k // split // bk))
    _router_split[(n // bn, split)](x, state["router"], partial, out, state["tile_counters"], m, n, k, bn, bk,
                                    max(2, triton.next_power_of_2(m)), split, stages, num_warps=warps)
    return out


def m_r009(x, state, layer, n, k, p):
    m = x.shape[0]
    if 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:]
    split, bm = 16, 16
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32)
    _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)
    return _split_sum(partial, _empty(x, m, n, layer), split)


def m_rpair(x, state, layer, n, k, p):
    bm, bn, bk, split = p[:4]
    tuning = {"num_warps": p[4], "num_stages": p[5]} if len(p) > 4 else {"num_warps": 4}
    m = x.shape[0]
    out = _empty(x, m, n, layer)
    target = out if split == 1 else torch.empty((split, m, n), device=x.device, dtype=torch.float32)
    _paired_router[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, state["router_interleaved"], target, m, n, k, bm, bn, bk, split, **tuning)
    return out if split == 1 else _split_sum(target, out, split)


METHODS = {"cublas": m_cublas, "hopper": m_hopper, "gemv": m_gemv, "vector": m_vector, "trans": m_trans,
           "widek": m_widek, "splitf": m_splitf, "splitft": m_splitft, "streamk": m_streamk, "r009": m_r009, "rsmall": m_rsmall, "rvector": m_vector, "rpair": m_rpair, "rfused": m_rfused, "rtiled": m_rtiled, "rsplit": m_rsplit}


TILE_COUNTERS = 4096
SMS = 132  # H200 SXM; the judge's GPU


def applicable(method, layer, m, n, k, params=()):
    """Whether a method computes this shape exactly as specified."""
    if layer == "router":
        if method == "rsplit":
            return m <= 32 and k % (params[1] * params[2]) == 0
        if method == "rtiled":
            return k % (params[2] * params[3]) == 0 and params[1] in (16, 32, 64)
        if method == "rfused":
            return k % (params[2] * params[3]) == 0
        return method in ("r009", "rpair") or (method == "rvector" and m <= 8) or (method == "rsmall" and m <= 32)
    if method in ("r009", "rpair", "rvector", "rsmall", "rfused", "rtiled", "rsplit"):
        return False
    if method == "hopper":
        return use_hopper_bf16_gemv(m, n, k)
    if method == "gemv":
        return m == 1
    if method == "vector":
        return m <= 8
    if method == "streamk":
        bm, bn, bk = params[:3]
        return k % bk == 0 and triton.cdiv(m, bm) * triton.cdiv(n, bn) <= TILE_COUNTERS and m <= 4 * bm
    if method in ("splitf", "splitft"):
        bm, bn, bk, split = params[:4]
        return (k % (bk * split) == 0 and triton.cdiv(m, bm) * triton.cdiv(n, bn) <= TILE_COUNTERS
                and m <= 4 * bm)
    return True


def choose(layer, m, n, k):
    for limit, method, params in PLAN[layer]:
        if m <= limit:
            break
    if not applicable(method, layer, m, n, k, params):
        method, params = ("r009", ()) if layer == "router" else ("cublas", ())
    return method, params


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)
    n, k = state["router"].shape
    state["router_interleaved"] = torch.stack(
        (state["router_hi_t"], state["router_lo_t"]), dim=-1).reshape(k, 2 * n).contiguous()
    for bn in (16, 32, 64):  # lossless repack: [N / bn, K, 2 bn] of the interleaved hi/lo pair
        state[f"router_tiled_{bn}"] = state["router_interleaved"].view(k, n // bn, 2 * bn).permute(1, 0, 2).contiguous()
    state["tile_counters"] = torch.zeros(TILE_COUNTERS, dtype=torch.int32, device=state["router"].device)
    if any(method in ("trans", "splitft") for plan in PLAN.values() for _, method, _ in plan):
        for name in DENSE:
            state["t_" + name] = state[name].t().contiguous()
    return state


def kernel_fn(data: Inputs, state) -> torch.Tensor:
    x, cfg = data["x"], data["config"]
    layer, n, k = cfg["layer"], cfg["n"], cfg["k"]
    method, params = choose(layer, x.shape[0], n, k)
    return METHODS[method](x, state, layer, n, k, params)