# 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)), (8, "cluster", (8, 4, 2)), (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, "flashinfer", ("tinygemm",)), (1 << 30, "trans", ())],
    "router": [(1, "rvector", (1, 8)), (16, "ratomic", (16, 64, 64, 32, 4, 4)), (32, "ratomic", (32, 32, 64, 32, 4, 4)),
               (128, "ratomic", (64, 64, 128, 16, 4, 4)), (1 << 30, "ratomic", (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 _dense_split_fused(X, W, P, Y, C, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                       BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, SPLIT: tl.constexpr):
    # _dense_split's main loop (splits interleaved along K), with its partials summed in a
    # fixed order by the last split to finish the tile, in the same launch.
    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)
    out_mask = (rm[:, None] < M) & (rn[None, :] < N)
    tl.store(P + tl.program_id(2) * M * N + rm[:, None] * N + rn[None, :], acc, out_mask)
    tile = tl.program_id(0) * tl.num_programs(1) + tl.program_id(1)
    if tl.atomic_add(C + tile, 1, sem="acq_rel") == SPLIT - 1:
        sp = tl.arange(0, SPLIT)
        total = tl.sum(tl.load(P + sp[:, None, None] * M * N + rm[None, :, None] * N + rn[None, None, :],
                               mask=out_mask[None, :, :], other=0, cache_modifier=".cg"), axis=0)
        tl.store(Y + rm[:, None] * N + rn[None, :], total.to(tl.bfloat16), mask=out_mask)
        tl.atomic_xchg(C + tile, 0)


@triton.jit
def _paired_router_split_fused(X, W, P, Y, C, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                               BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, SPLIT: tl.constexpr):
    # _paired_router's main loop with the same in-launch, fixed-order split reduction.
    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)
    result = tl.sum(tl.reshape(acc, (BM, BN, 2)), axis=2)
    cols = tl.program_id(1) * BN + tl.arange(0, BN)
    mask = (rows[:, None] < M) & (cols[None, :] < N)
    tl.store(P + tl.program_id(2) * M * N + rows[:, None] * N + cols[None, :], result, mask)
    tile = tl.program_id(0) * tl.num_programs(1) + tl.program_id(1)
    if tl.atomic_add(C + tile, 1, sem="acq_rel") == SPLIT - 1:
        sp = tl.arange(0, SPLIT)
        total = tl.sum(tl.load(P + sp[:, None, None] * M * N + rows[None, :, None] * N + cols[None, None, :],
                               mask=mask[None, :, :], other=0, cache_modifier=".cg"), axis=0)
        tl.store(Y + rows[:, None] * N + cols[None, :], total, mask=mask)
        tl.atomic_xchg(C + tile, 0)


@triton.jit
def _paired_router_atomic(X, W, Y, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                          BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, SPLIT: tl.constexpr):
    # _paired_router's main loop; each split adds its FP32 partial into the zeroed FP32 output
    # with FP32 atomics, so no reduction pass waits on the slowest split.
    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)
    result = tl.sum(tl.reshape(acc, (BM, BN, 2)), axis=2)
    cols = tl.program_id(1) * BN + tl.arange(0, BN)
    tl.atomic_add(Y + rows[:, None] * N + cols[None, :], result,
                  mask=(rows[:, None] < M) & (cols[None, :] < N), sem="relaxed")


@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:
            sp = tl.arange(0, SPLIT)
            total = tl.sum(tl.load(P + sp[:, None, None] * M * N + rm[None, :, None] * N + rn[None, None, :],
                                   mask=out_mask[None, :, :], other=0, cache_modifier=".cg"), axis=0)
            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:
            sp = tl.arange(0, SPLIT)
            total = tl.sum(tl.load(P + sp[:, None, None] * M * N + rm[None, :, None] * N + rn[None, None, :],
                                   mask=out_mask[None, :, :], other=0, cache_modifier=".cg"), axis=0)
            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:
            sp = tl.arange(0, SPLIT)
            total = tl.sum(tl.load(P + sp[:, None, None] * M * N + rows[None, :, None] * N + cols[None, None, :],
                                   mask=mask[None, :, :], other=0, cache_modifier=".cg"), axis=0)
            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:
            sp = tl.arange(0, SPLIT)
            total = tl.sum(tl.load(P + sp[:, None, None] * M * N + rows[None, :, None] * N + cols[None, None, :],
                                   mask=mask[None, :, :], other=0, cache_modifier=".cg"), axis=0)
            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_widekf(x, state, layer, n, k, p):
    split, bm, bn, bk, 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)
    _dense_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,
        num_warps=warps, num_stages=stages)
    return out


def m_rpairf(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)
    partial = torch.empty((split, m, n), device=x.device, dtype=torch.float32)
    _paired_router_split_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, **tuning)
    return out


def m_ratomic(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 = torch.zeros((m, n), device=x.device, dtype=torch.float32)
    _paired_router_atomic[(triton.cdiv(m, bm), triton.cdiv(n, bn), split)](
        x, state["router_interleaved"], out, m, n, k, bm, bn, bk, split, **tuning)
    return out


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_deepgemm(x, state, layer, n, k, p):
    import deep_gemm
    out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
    deep_gemm.bf16_gemm_nt(x, state[layer][:n], out)
    return out


def m_flashinfer(x, state, layer, n, k, p):
    import flashinfer
    # mm_bf16 takes the weight as a column-major (k, n) matrix: the transpose view of W[:n].
    return flashinfer.mm_bf16(x, state[layer][:n].t(), backend=p[0])


def m_quack(x, state, layer, n, k, p):
    from quack.gemm import gemm as quack_gemm
    bm, bn, cm, cn, ck, pingpong = p
    out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
    quack_gemm(x.unsqueeze(0), state[layer][:n].unsqueeze(0), out.unsqueeze(0), None, None, bm, bn, cm, cn,
               cluster_K=ck, pingpong=pingpong, persistent=True)
    return out


_CLUSTER_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cooperative_groups.h>
#include <cuda_bf16.h>
namespace cg = cooperative_groups;

__device__ __forceinline__ void cp_async16_zfill(void* dst, const void* src, bool valid) {
  const unsigned s = static_cast<unsigned>(__cvta_generic_to_shared(dst));
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" :: "r"(s), "l"(src), "r"(valid ? 16 : 0) : "memory");
}
__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;" ::: "memory"); }
template <int N> __device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;" :: "n"(N) : "memory"); }

__device__ __forceinline__ void ldsm_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, const void* p) {
  const unsigned s = static_cast<unsigned>(__cvta_generic_to_shared(p));
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
               : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(s));
}
__device__ __forceinline__ void ldsm_x2(uint32_t& r0, uint32_t& r1, const void* p) {
  const unsigned s = static_cast<unsigned>(__cvta_generic_to_shared(p));
  asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];" : "=r"(r0), "=r"(r1) : "r"(s));
}
__device__ __forceinline__ void mma16816(float* c, uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
                                         uint32_t b0, uint32_t b1) {
  asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
               "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
               : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
               : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// Tiles of 64 elements of K; a row of a tile is 128 bytes = 8 chunks of 16 bytes, and chunk c
// of row r lives at chunk c ^ (r & 7), so ldmatrix's 8-row reads hit distinct banks.
__device__ __forceinline__ int swz(int row, int chunk) { return row * 64 + ((chunk ^ (row & 7)) << 3); }

// out[M, N] = x[M, K] W[N, K]^T. CTA (bx, by) owns 16 NW output features (rows of W)
// and the K-tiles by, by + S, by + 2S, ...; the S CTAs of a cluster (grid.y = S) sum their FP32
// partials through distributed shared memory in rank order and write BF16.
template <int MT, int STAGES, int NW>
__global__ void __launch_bounds__(NW * 32) cluster_gemm(const __nv_bfloat16* __restrict__ x,
                                                        const __nv_bfloat16* __restrict__ w,
                                                        __nv_bfloat16* __restrict__ out, int M, int N, int K) {
  constexpr int ROWS = 16 * NW, TOKS = MT * 8, THREADS = NW * 32;
  constexpr int W_ELEMS = ROWS * 64, X_ELEMS = TOKS * 64, STAGE_ELEMS = W_ELEMS + X_ELEMS;
  extern __shared__ __align__(128) __nv_bfloat16 smem[];
  float* red = reinterpret_cast<float*>(smem + STAGES * STAGE_ELEMS);  // [TOKS][ROWS]
  cg::cluster_group cluster = cg::this_cluster();
  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
  const int S = gridDim.y, rank = blockIdx.y;
  const int n0 = blockIdx.x * ROWS;
  const int count = (K / 64) / S;  // this CTA's K-tiles

  auto load = [&](int slot, int i) {
    const int kt = rank + i * S;
    __nv_bfloat16* ws = smem + slot * STAGE_ELEMS;
    __nv_bfloat16* xs = ws + W_ELEMS;
#pragma unroll
    for (int j = 0; j < (ROWS * 8) / THREADS; ++j) {
      const int idx = tid + j * THREADS, r = idx >> 3, c = idx & 7;
      cp_async16_zfill(ws + swz(r, c), w + (size_t)(n0 + r) * K + kt * 64 + c * 8, true);
    }
    for (int idx = tid; idx < TOKS * 8; idx += THREADS) {
      const int r = idx >> 3, c = idx & 7;
      const bool valid = r < M;
      cp_async16_zfill(xs + swz(r, c), x + (size_t)(valid ? r : 0) * K + kt * 64 + c * 8, valid);
    }
  };

#pragma unroll
  for (int s = 0; s < STAGES - 1; ++s) {
    if (s < count) load(s, s);
    cp_commit();
  }

  float acc[MT][4];
#pragma unroll
  for (int mt = 0; mt < MT; ++mt)
#pragma unroll
    for (int i = 0; i < 4; ++i) acc[mt][i] = 0.f;

  for (int i = 0; i < count; ++i) {
    cp_wait<STAGES - 2>();
    __syncthreads();
    if (i + STAGES - 1 < count) load((i + STAGES - 1) % STAGES, i + STAGES - 1);
    cp_commit();
    const __nv_bfloat16* ws = smem + (i % STAGES) * STAGE_ELEMS;
    const __nv_bfloat16* xs = ws + W_ELEMS;
#pragma unroll
    for (int ks = 0; ks < 4; ++ks) {
      uint32_t a0, a1, a2, a3;
      ldsm_x4(a0, a1, a2, a3, ws + swz(warp * 16 + (lane & 15), ks * 2 + (lane >> 4)));
#pragma unroll
      for (int mt = 0; mt < MT; ++mt) {
        uint32_t b0, b1;
        ldsm_x2(b0, b1, xs + swz(mt * 8 + (lane & 7), ks * 2 + ((lane >> 3) & 1)));
        mma16816(acc[mt], a0, a1, a2, a3, b0, b1);
      }
    }
  }

  // C fragment: rows warp*16 + g (+8), tokens mt*8 + 2t (+1).
  const int g = lane >> 2, t = lane & 3;
#pragma unroll
  for (int mt = 0; mt < MT; ++mt) {
    const int row = warp * 16 + g, tok = mt * 8 + 2 * t;
    red[tok * ROWS + row] = acc[mt][0];
    red[(tok + 1) * ROWS + row] = acc[mt][1];
    red[tok * ROWS + row + 8] = acc[mt][2];
    red[(tok + 1) * ROWS + row + 8] = acc[mt][3];
  }
  cluster.sync();
  const int valid_e = (M < TOKS ? M : TOKS) * ROWS;
  const int lo = (int)((long)valid_e * rank / S), hi = (int)((long)valid_e * (rank + 1) / S);
  for (int e = lo + tid; e < hi; e += THREADS) {
    float sum = 0.f;
    for (int q = 0; q < S; ++q) sum += cluster.map_shared_rank(red, q)[e];
    out[(size_t)(e / ROWS) * N + n0 + e % ROWS] = __float2bfloat16(sum);
  }
  cluster.sync();
}

template <int MT, int STAGES, int NW>
void launch_cluster(const torch::Tensor& x, const torch::Tensor& w, torch::Tensor& out, int S) {
  const int M = x.size(0), K = x.size(1), N = out.size(1);
  constexpr int ROWS = 16 * NW, TOKS = MT * 8;
  const size_t smem = (size_t)STAGES * (ROWS + TOKS) * 64 * 2 + (size_t)TOKS * ROWS * 4;
  auto kernel = cluster_gemm<MT, STAGES, NW>;
  static bool configured = false;
  if (!configured) {
    C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 226 * 1024));
    C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1));
    configured = true;
  }
  cudaLaunchConfig_t cfg = {};
  cfg.gridDim = dim3(N / ROWS, S);
  cfg.blockDim = dim3(NW * 32);
  cfg.dynamicSmemBytes = smem;
  cfg.stream = at::cuda::getCurrentCUDAStream();
  cudaLaunchAttribute attr[1];
  attr[0].id = cudaLaunchAttributeClusterDimension;
  attr[0].val.clusterDim.x = 1;
  attr[0].val.clusterDim.y = S;
  attr[0].val.clusterDim.z = 1;
  cfg.attrs = attr;
  cfg.numAttrs = 1;
  C10_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kernel, reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
                                    reinterpret_cast<const __nv_bfloat16*>(w.data_ptr()),
                                    reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), M, N, K));
}

void cluster_mm(torch::Tensor x, torch::Tensor w, torch::Tensor out, int64_t s, int64_t mt, int64_t stages, int64_t nw) {
@DISPATCH@
  TORCH_CHECK(false, "cluster_mm: no instance for MT=", mt, " STAGES=", stages, " NW=", nw);
}
"""
_CLUSTER_DECL = "void cluster_mm(torch::Tensor x, torch::Tensor w, torch::Tensor out, int64_t s, int64_t mt, int64_t stages, int64_t nw);"
CLUSTER_INSTANCES = None  # (MT, STAGES, NW) to compile; None means those PLAN uses
_CLUSTER_EXT = {}


def _cluster_instances():
    if CLUSTER_INSTANCES is not None:
        return sorted(CLUSTER_INSTANCES)
    return sorted({(mt, p[1], p[2]) for layer, plan in PLAN.items() for limit, method, p in plan if method == "cluster"
                   for mt in (1, 2, 4)})


def _cluster_ext():
    inst = tuple(_cluster_instances())
    if inst not in _CLUSTER_EXT:
        import hashlib
        from torch.utils.cpp_extension import load_inline
        dispatch = "\n".join(f"  if (mt == {a} && stages == {b} && nw == {c}) {{ launch_cluster<{a}, {b}, {c}>(x, w, out, s); return; }}"
                             for a, b, c in inst)
        src = _CLUSTER_CUDA.replace("@DISPATCH@", dispatch)
        name = "okc_cluster_" + hashlib.sha256(src.encode()).hexdigest()[:12]
        _CLUSTER_EXT[inst] = load_inline(name=name, cpp_sources=_CLUSTER_DECL, cuda_sources=src,
                                         functions=["cluster_mm"], extra_cuda_cflags=["-O3", "-std=c++20"],
                                         verbose=False)
    return _CLUSTER_EXT[inst]


def m_cluster(x, state, layer, n, k, p):
    split, stages, nw = p
    m = x.shape[0]
    out = torch.empty((m, n), device=x.device, dtype=torch.bfloat16)
    _cluster_ext().cluster_mm(x, state[layer], out, split, (m + 7) // 8 if m <= 16 else 4, stages, nw)
    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, "deepgemm": m_deepgemm, "cluster": m_cluster, "quack": m_quack, "flashinfer": m_flashinfer, "widekf": m_widekf, "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, "rpairf": m_rpairf, "ratomic": m_ratomic, "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 == "ratomic":
            return True
        if method == "rpairf":
            return triton.cdiv(m, params[0]) * triton.cdiv(n, params[1]) <= TILE_COUNTERS
        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", "rpairf", "ratomic"):
        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 == "cluster":
        split, stages, nw = params
        rows = 16 * nw
        smem = stages * (rows + 32) * 64 * 2 + 32 * rows * 4
        return m <= 32 and n % rows == 0 and k % (64 * split) == 0 and split <= 16 and smem <= 220 * 1024
    if method == "widekf":
        return triton.cdiv(m, params[1]) * triton.cdiv(n, params[2]) <= TILE_COUNTERS
    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()
    if any(method == "cluster" for plan in PLAN.values() for _, method, _ in plan) or CLUSTER_INSTANCES:
        _cluster_ext()
    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)