# 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, "cluster", (8, 3, 2, 128)), (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); }
// The same for BK-wide rows (BK / 8 chunks): only the low 3 chunk bits are permuted.
template <int BK>
__device__ __forceinline__ int swzk(int row, int chunk) { return row * BK + ((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, int BK>
__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 * BK, X_ELEMS = TOKS * BK, STAGE_ELEMS = W_ELEMS + X_ELEMS, CH = BK / 8;
  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 / BK) / 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 * CH) / THREADS; ++j) {
      const int idx = tid + j * THREADS, r = idx / CH, c = idx % CH;
      cp_async16_zfill(ws + swzk<BK>(r, c), w + (size_t)(n0 + r) * K + kt * BK + c * 8, true);
    }
    for (int idx = tid; idx < TOKS * CH; idx += THREADS) {
      const int r = idx / CH, c = idx % CH;
      const bool valid = r < M;
      cp_async16_zfill(xs + swzk<BK>(r, c), x + (size_t)(valid ? r : 0) * K + kt * BK + 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 < BK / 16; ++ks) {
      uint32_t a0, a1, a2, a3;
      ldsm_x4(a0, a1, a2, a3, ws + swzk<BK>(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 + swzk<BK>(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, int BK>
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) * BK * 2 + (size_t)TOKS * ROWS * 4;
  auto kernel = cluster_gemm<MT, STAGES, NW, BK>;
  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));
}

// Large-batch variant: CTA tile TM tokens x TN output features, warps arranged
// (TM / 64) x (TN / 32), each warp 64 tokens x 32 features (4 x 4 MMA tiles). Tokens are the
// MMA's A operand (x rows), features its B operand (W rows, i.e. a column-major K x N).
// grid (N / TN, S, ceil(M / TM)); the S CTAs along y form a cluster and reduce as above.
template <int TM, int TN, int STAGES>
__global__ void __launch_bounds__((TM / 64) * (TN / 32) * 32) cluster_gemm_big(
    const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ w,
    __nv_bfloat16* __restrict__ out, int M, int N, int K) {
  constexpr int WN = TN / 32, THREADS = (TM / 64) * WN * 32;
  constexpr int W_ELEMS = TN * 64, X_ELEMS = TM * 64, STAGE_ELEMS = W_ELEMS + X_ELEMS;
  extern __shared__ __align__(128) __nv_bfloat16 smem[];
  cg::cluster_group cluster = cg::this_cluster();
  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
  const int wm = warp / WN, wn = warp % WN;
  const int S = gridDim.y, rank = blockIdx.y;
  const int n0 = blockIdx.x * TN, m0 = blockIdx.z * TM;
  const int count = (K / 64) / S;

  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 < (TN * 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);
    }
#pragma unroll
    for (int j = 0; j < (TM * 8) / THREADS; ++j) {
      const int idx = tid + j * THREADS, r = idx >> 3, c = idx & 7;
      const bool valid = m0 + r < M;
      cp_async16_zfill(xs + swz(r, c), x + (size_t)(valid ? m0 + 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[4][4][4];
#pragma unroll
  for (int i = 0; i < 4; ++i)
#pragma unroll
    for (int j = 0; j < 4; ++j)
#pragma unroll
      for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.f;

  for (int it = 0; it < count; ++it) {
    cp_wait<STAGES - 2>();
    __syncthreads();
    if (it + STAGES - 1 < count) load((it + STAGES - 1) % STAGES, it + STAGES - 1);
    cp_commit();
    const __nv_bfloat16* ws = smem + (it % STAGES) * STAGE_ELEMS;
    const __nv_bfloat16* xs = ws + W_ELEMS;
#pragma unroll
    for (int ks = 0; ks < 4; ++ks) {
      uint32_t a[4][4], b[4][2];
#pragma unroll
      for (int i = 0; i < 4; ++i)
        ldsm_x4(a[i][0], a[i][1], a[i][2], a[i][3], xs + swz(wm * 64 + i * 16 + (lane & 15), ks * 2 + (lane >> 4)));
#pragma unroll
      for (int j = 0; j < 2; ++j)
        ldsm_x4(b[2 * j][0], b[2 * j][1], b[2 * j + 1][0], b[2 * j + 1][1],
                ws + swz(wn * 32 + j * 16 + (lane & 7) + ((lane >> 4) << 3), ks * 2 + ((lane >> 3) & 1)));
#pragma unroll
      for (int i = 0; i < 4; ++i)
#pragma unroll
        for (int j = 0; j < 4; ++j) mma16816(acc[i][j], a[i][0], a[i][1], a[i][2], a[i][3], b[j][0], b[j][1]);
    }
  }

  // Overlay the FP32 partial tile [TM][TN] on the drained pipeline buffers.
  cp_wait<0>();
  __syncthreads();
  float* red = reinterpret_cast<float*>(smem);
  const int g = lane >> 2, t = lane & 3;
#pragma unroll
  for (int i = 0; i < 4; ++i)
#pragma unroll
    for (int j = 0; j < 4; ++j) {
      const int tok = wm * 64 + i * 16 + g, col = wn * 32 + j * 8 + 2 * t;
      red[tok * TN + col] = acc[i][j][0];
      red[tok * TN + col + 1] = acc[i][j][1];
      red[(tok + 8) * TN + col] = acc[i][j][2];
      red[(tok + 8) * TN + col + 1] = acc[i][j][3];
    }
  cluster.sync();
  const int rows_valid = (M - m0 < TM ? M - m0 : TM);
  const int valid_e = rows_valid * TN;
  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)(m0 + e / TN) * N + n0 + e % TN] = __float2bfloat16(sum);
  }
  cluster.sync();
}

template <int TM, int TN, int STAGES>
cudaLaunchConfig_t big_config(int M, int N, int S, cudaLaunchAttribute* attr) {
  cudaLaunchConfig_t cfg = {};
  cfg.gridDim = dim3(N / TN, S, (M + TM - 1) / TM);
  cfg.blockDim = dim3((TM / 64) * (TN / 32) * 32);
  const size_t pipe = (size_t)STAGES * (TN + TM) * 64 * 2, red = (size_t)TM * TN * 4;
  cfg.dynamicSmemBytes = pipe > red ? pipe : red;
  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;
  auto kernel = cluster_gemm_big<TM, TN, STAGES>;
  C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 226 * 1024));
  C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1));
  return cfg;
}

template <int TM, int TN, int STAGES>
void launch_big(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);
  cudaLaunchAttribute attr[1];
  cudaLaunchConfig_t cfg = big_config<TM, TN, STAGES>(M, N, S, attr);
  cfg.stream = at::cuda::getCurrentCUDAStream();
  C10_CUDA_CHECK(cudaLaunchKernelEx(&cfg, cluster_gemm_big<TM, TN, STAGES>,
                                    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));
}

template <int TM, int TN, int STAGES>
int64_t big_clusters(int M, int N, int S) {
  cudaLaunchAttribute attr[1];
  cudaLaunchConfig_t cfg = big_config<TM, TN, STAGES>(M, N, S, attr);
  int clusters = 0;
  C10_CUDA_CHECK(cudaOccupancyMaxActiveClusters(&clusters, cluster_gemm_big<TM, TN, STAGES>, &cfg));
  return clusters;
}

void big_mm(torch::Tensor x, torch::Tensor w, torch::Tensor out, int64_t s, int64_t tm, int64_t tn, int64_t stages) {
@BIGDISPATCH@
  TORCH_CHECK(false, "big_mm: no instance for TM=", tm, " TN=", tn, " STAGES=", stages);
}

int64_t big_occupancy(int64_t m, int64_t n, int64_t s, int64_t tm, int64_t tn, int64_t stages) {
@BIGOCCUPANCY@
  return -1;
}

template <int MT, int STAGES, int NW, int BK>
int64_t active_clusters(int N, int S) {
  constexpr int ROWS = 16 * NW, TOKS = MT * 8;
  const size_t smem = (size_t)STAGES * (ROWS + TOKS) * BK * 2 + (size_t)TOKS * ROWS * 4;
  auto kernel = cluster_gemm<MT, STAGES, NW, BK>;
  C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 226 * 1024));
  C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1));
  cudaLaunchConfig_t cfg = {};
  cfg.gridDim = dim3(N / ROWS, S);
  cfg.blockDim = dim3(NW * 32);
  cfg.dynamicSmemBytes = smem;
  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;
  int clusters = 0;
  C10_CUDA_CHECK(cudaOccupancyMaxActiveClusters(&clusters, kernel, &cfg));
  return clusters;
}

int64_t cluster_occupancy(int64_t n, int64_t s, int64_t mt, int64_t stages, int64_t nw, int64_t bk) {
@OCCUPANCY@
  return -1;
}

void cluster_mm(torch::Tensor x, torch::Tensor w, torch::Tensor out, int64_t s, int64_t mt, int64_t stages, int64_t nw, int64_t bk) {
@DISPATCH@
  TORCH_CHECK(false, "cluster_mm: no instance for MT=", mt, " STAGES=", stages, " NW=", nw, " BK=", bk);
}
"""
_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, int64_t bk);\n"
                 "int64_t cluster_occupancy(int64_t n, int64_t s, int64_t mt, int64_t stages, int64_t nw, int64_t bk);\n"
                 "void big_mm(torch::Tensor x, torch::Tensor w, torch::Tensor out, int64_t s, int64_t tm, int64_t tn, int64_t stages);\n"
                 "int64_t big_occupancy(int64_t m, int64_t n, int64_t s, int64_t tm, int64_t tn, int64_t stages);")
BIG_INSTANCES = ((64, 128, 3), (64, 128, 4), (128, 128, 2), (128, 128, 3), (128, 64, 3), (128, 64, 4), (64, 64, 4))
CLUSTER_INSTANCES = None  # (MT, STAGES, NW, BK) 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], p[3] if len(p) > 3 else 64) 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} && bk == {d}) {{ launch_cluster<{a}, {b}, {c}, {d}>(x, w, out, s); return; }}"
                             for a, b, c, d in inst)
        occupancy = "\n".join(f"  if (mt == {a} && stages == {b} && nw == {c} && bk == {d}) return active_clusters<{a}, {b}, {c}, {d}>(n, s);"
                              for a, b, c, d in inst)
        big = "\n".join(f"  if (tm == {a} && tn == {b} && stages == {c}) {{ launch_big<{a}, {b}, {c}>(x, w, out, s); return; }}"
                        for a, b, c in BIG_INSTANCES)
        big_occ = "\n".join(f"  if (tm == {a} && tn == {b} && stages == {c}) return big_clusters<{a}, {b}, {c}>(m, n, s);"
                            for a, b, c in BIG_INSTANCES)
        src = (_CLUSTER_CUDA.replace("@DISPATCH@", dispatch).replace("@OCCUPANCY@", occupancy)
               .replace("@BIGDISPATCH@", big).replace("@BIGOCCUPANCY@", big_occ))
        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", "cluster_occupancy", "big_mm", "big_occupancy"], 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[:3]
    bk = p[3] if len(p) > 3 else 64
    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, bk)
    return out


def m_big(x, state, layer, n, k, p):
    split, tm, tn, stages = p
    out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
    _cluster_ext().big_mm(x, state[layer], out, split, tm, tn, stages)
    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, "big": m_big, "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 == "big":
        split, tm, tn, stages = params
        return (m > 32 and (tm, tn, stages) in BIG_INSTANCES and n % tn == 0 and k % (64 * split) == 0
                and split <= 16)
    if method == "cluster":
        split, stages, nw = params[:3]
        bk = params[3] if len(params) > 3 else 64
        rows, toks = 16 * nw, (8 if m <= 8 else 16 if m <= 16 else 32)
        smem = stages * (rows + toks) * bk * 2 + toks * rows * 4
        return m <= 32 and n % rows == 0 and k % (bk * 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 in ("cluster", "big") 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)

"""Extra exact-BF16 methods for the 1.32x projection round.

This fragment is appended to the immutable, officially scored 033 source.
The default dispatch stays identical until a measured plan is selected.
"""


@triton.jit
def _c2_dense_packed(X, W, Y, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                     BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
                     STAGES: tl.constexpr):
    # Each output-feature tile owns a contiguous [K, BN] block of exact weights.
    # With BN=16 the pooler launches 256 independent output tiles, without a
    # split-K partial buffer or a separate reduction kernel.
    pm, pn = tl.program_id(0), tl.program_id(1)
    rows = pm * BM + tl.arange(0, BM)
    columns = pn * BN + tl.arange(0, BN)
    kk = tl.arange(0, BK)
    nn = tl.arange(0, BN)
    accum = tl.zeros((BM, BN), tl.float32)
    for start in tl.range(0, K, BK, num_stages=STAGES):
        k = start + kk
        a = tl.load(X + rows[:, None] * K + k[None, :],
                    (rows[:, None] < M) & (k[None, :] < K), other=0)
        b = tl.load(W + pn * K * BN + k[:, None] * BN + nn[None, :],
                    (k[:, None] < K) & (columns[None, :] < N), other=0)
        accum += tl.dot(a, b)
    tl.store(Y + rows[:, None] * N + columns[None, :], accum.to(tl.bfloat16),
             (rows[:, None] < M) & (columns[None, :] < N))


@triton.jit
def _c2_router_wide(X, W, Y, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
                    SPLIT: tl.constexpr, STAGES: tl.constexpr):
    # Wider output tiles reduce the number of CTAs re-reading the activations.
    # High/low products retain separate FP32 accumulators until the epilogue.
    pm, pn, ps = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    rows = pm * BM + tl.arange(0, BM)
    pair_columns = tl.arange(0, 2 * BN)
    kk = tl.arange(0, BK)
    acc = tl.zeros((BM, 2 * BN), tl.float32)
    chunk: tl.constexpr = K // SPLIT
    for start in tl.range(0, chunk, BK, num_stages=STAGES):
        k = ps * chunk + start + kk
        a = tl.load(X + rows[:, None] * K + k[None, :],
                    (rows[:, None] < M) & (k[None, :] < K), other=0)
        b = tl.load(W + pn * K * (2 * BN) + k[:, None] * (2 * BN) + pair_columns[None, :],
                    (k[:, None] < K) & (pn * BN * 2 + pair_columns[None, :] < 2 * N), other=0)
        acc += tl.dot(a, b)
    total = tl.sum(tl.reshape(acc, (BM, BN, 2)), axis=2)
    cols = pn * BN + tl.arange(0, BN)
    mask = (rows[:, None] < M) & (cols[None, :] < N)
    if SPLIT == 1:
        tl.store(Y + rows[:, None] * N + cols[None, :], total, mask)
    else:
        tl.atomic_add(Y + rows[:, None] * N + cols[None, :], total, mask=mask, sem="relaxed")


def m_c2_packed(x, state, layer, n, k, params):
    bm, bn, bk, warps, stages = params
    out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
    _c2_dense_packed[(triton.cdiv(x.shape[0], bm), triton.cdiv(n, bn))](
        x, state[f"c2_dense_{layer}_{bn}"], out, x.shape[0], n, k,
        bm, bn, bk, stages, num_warps=warps)
    return out


def m_c2_router(x, state, layer, n, k, params):
    bm, bn, bk, split, warps, stages = params
    allocate = torch.zeros if split > 1 else torch.empty
    out = allocate((x.shape[0], n), device=x.device, dtype=torch.float32)
    _c2_router_wide[(triton.cdiv(x.shape[0], bm), triton.cdiv(n, bn), split)](
        x, state[f"c2_router_{bn}"], out, x.shape[0], n, k,
        bm, bn, bk, split, stages, num_warps=warps)
    return out


METHODS["c2packed"] = m_c2_packed
METHODS["c2router"] = m_c2_router
_c2_base_applicable = applicable
_c2_base_prepare = prepare


def applicable(method, layer, m, n, k, params=()):
    if method == "c2packed":
        bm, bn, bk, warps, stages = params
        return layer != "router" and bn in (16, 32, 64) and n % bn == 0 and k % bk == 0
    if method == "c2router":
        bm, bn, bk, split, warps, stages = params
        return layer == "router" and bn in (64, 128, 256) and n % bn == 0 and k % (bk * split) == 0
    return _c2_base_applicable(method, layer, m, n, k, params)


def prepare(weights: Weights):
    state = _c2_base_prepare(weights)
    # Lossless layout transformations only. No evaluated activations are present.
    for layer in DENSE:
        n, k = weights[layer].shape
        for bn in (16, 32, 64):
            state[f"c2_dense_{layer}_{bn}"] = state[layer].view(n // bn, bn, k).permute(0, 2, 1).contiguous()
    n, k = weights["router"].shape
    for bn in (64, 128, 256):
        state[f"c2_router_{bn}"] = state["router_interleaved"].view(k, n // bn, 2 * bn).permute(1, 0, 2).contiguous()
    return state

"""Explicit cuBLASLt heuristic selection, FP32 compute/reductions, BF16 output."""
_LT_CPP = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublasLt.h>
#include <map>
#include <memory>
#include <tuple>
#include <vector>
#define LT_OK(x) TORCH_CHECK((x)==CUBLAS_STATUS_SUCCESS, "cuBLASLt failed: ", #x)
struct LTPlan {
 cublasLtMatmulDesc_t op;
 cublasLtMatrixLayout_t a,b,c;
 std::vector<cublasLtMatmulHeuristicResult_t> algos;
 ~LTPlan(){cublasLtMatmulDescDestroy(op);cublasLtMatrixLayoutDestroy(a);cublasLtMatrixLayoutDestroy(b);cublasLtMatrixLayoutDestroy(c);}
};
static cublasLtHandle_t lt_handle(){static cublasLtHandle_t h=nullptr; if(!h) LT_OK(cublasLtCreate(&h));return h;}
static LTPlan& lt_plan(int m,int n,int k,int trans,int lda) {
 static std::map<std::tuple<int,int,int,int,int,int>,std::unique_ptr<LTPlan>> cache;
 auto key=std::make_tuple(at::cuda::current_device(),m,n,k,trans,lda);
 auto it=cache.find(key);if(it!=cache.end()) return *it->second;
 auto p=std::make_unique<LTPlan>();
 LT_OK(cublasLtMatmulDescCreate(&p->op,CUBLAS_COMPUTE_32F,CUDA_R_32F));
 cublasOperation_t op_a=(trans==1)?CUBLAS_OP_N:CUBLAS_OP_T;
 LT_OK(cublasLtMatmulDescSetAttribute(p->op,CUBLASLT_MATMUL_DESC_TRANSA,&op_a,sizeof(op_a)));
 cublasOperation_t op_b=(trans==3)?CUBLAS_OP_T:CUBLAS_OP_N;
 LT_OK(cublasLtMatmulDescSetAttribute(p->op,CUBLASLT_MATMUL_DESC_TRANSB,&op_b,sizeof(op_b)));
 if(trans<2) {
  LT_OK(cublasLtMatrixLayoutCreate(&p->a,CUDA_R_16BF,trans?n:k,trans?k:n,lda));
  LT_OK(cublasLtMatrixLayoutCreate(&p->b,CUDA_R_16BF,k,m,k));
  LT_OK(cublasLtMatrixLayoutCreate(&p->c,CUDA_R_16BF,n,m,n));
 } else {
  LT_OK(cublasLtMatrixLayoutCreate(&p->a,CUDA_R_16BF,k,m,k));
  LT_OK(cublasLtMatrixLayoutCreate(&p->b,CUDA_R_16BF,(trans==3)?n:k,(trans==3)?k:n,lda));
  LT_OK(cublasLtMatrixLayoutCreate(&p->c,CUDA_R_16BF,m,n,m));
 }
 cublasLtMatmulPreference_t pref;LT_OK(cublasLtMatmulPreferenceCreate(&pref));
 uint64_t ws=32ull<<20;uint32_t reduction=CUBLASLT_REDUCTION_SCHEME_COMPUTE_TYPE;
 LT_OK(cublasLtMatmulPreferenceSetAttribute(pref,CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,&ws,sizeof(ws)));
 LT_OK(cublasLtMatmulPreferenceSetAttribute(pref,CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK,&reduction,sizeof(reduction)));
 p->algos.resize(64);int count=0;
 auto status=cublasLtMatmulAlgoGetHeuristic(lt_handle(),p->op,p->a,p->b,p->c,p->c,pref,64,p->algos.data(),&count);
 cublasLtMatmulPreferenceDestroy(pref);LT_OK(status);
 p->algos.resize(count);TORCH_CHECK(count>0,"no cuBLASLt algorithm");
 for(int i=0;i<std::min(count,10);++i){
  for(int split: {2,4,8,16}){
   auto algo=p->algos[i].algo;
   int reduction=CUBLASLT_REDUCTION_SCHEME_COMPUTE_TYPE;
   if(cublasLtMatmulAlgoConfigSetAttribute(&algo,CUBLASLT_ALGO_CONFIG_SPLITK_NUM,&split,sizeof(split))!=CUBLAS_STATUS_SUCCESS) continue;
   if(cublasLtMatmulAlgoConfigSetAttribute(&algo,CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME,&reduction,sizeof(reduction))!=CUBLAS_STATUS_SUCCESS) continue;
   cublasLtMatmulHeuristicResult_t checked{};
   auto result=cublasLtMatmulAlgoCheck(lt_handle(),p->op,p->a,p->b,p->c,p->c,&algo,&checked);
   if(result==CUBLAS_STATUS_SUCCESS && checked.state==CUBLAS_STATUS_SUCCESS && checked.workspaceSize<=ws){checked.algo=algo;p->algos.push_back(checked);}
  }
 }

 auto &ref=*p;cache.emplace(key,std::move(p));return ref;
}
int64_t lt_count(int64_t m,int64_t n,int64_t k,int64_t trans,int64_t lda){return lt_plan(m,n,k,trans,lda).algos.size();}
void lt_mm(torch::Tensor x,torch::Tensor w,torch::Tensor y,torch::Tensor ws,int64_t trans,int64_t index){
 auto &p=lt_plan(x.size(0),trans>=2?y.size(0):y.size(1),x.size(1),trans,(trans==1||trans==3)?w.size(1):x.size(1));
 if(index>=p.algos.size()) index=0;
 TORCH_CHECK(index>=0,"invalid algorithm index");
 float alpha=1.f,beta=0.f;
 LT_OK(cublasLtMatmul(lt_handle(),p.op,&alpha,(trans>=2?x.data_ptr():w.data_ptr()),p.a,(trans>=2?w.data_ptr():x.data_ptr()),p.b,&beta,y.data_ptr(),p.c,
                    y.data_ptr(),p.c,&p.algos[index].algo,ws.data_ptr(),ws.numel(),at::cuda::getCurrentCUDAStream()));
}
'''
_LT_EXT = None


def _lt_ext():
    global _LT_EXT
    if _LT_EXT is None:
        import hashlib
        from torch.utils.cpp_extension import load_inline
        _LT_EXT = load_inline(name='okc_lt_' + hashlib.sha256(_LT_CPP.encode()).hexdigest()[:12],
                              cpp_sources=_LT_CPP, functions=['lt_mm', 'lt_count'],
                              with_cuda=True, extra_cflags=['-O3'], extra_ldflags=['-lcublasLt'], verbose=False)
    return _LT_EXT


_c7_fallback_plan = {layer: list(entries) for layer, entries in PLAN.items()}


def m_c7_lt(x, state, layer, n, k, params):
    trans, index = params
    # Swapped output has leading dimension M, so retain the original path when
    # its row stride does not meet the 16-byte BF16 alignment used by these tiles.
    if trans >= 2 and x.shape[0] % 8:
        for limit, method, p in _c7_fallback_plan[layer]:
            if x.shape[0] <= limit:
                break
        return METHODS[method](x, state, layer, n, k, p)
    w = state['t_' + layer] if trans in (1, 3) else state[layer]
    out = torch.empty((n, x.shape[0]) if trans >= 2 else (x.shape[0], n), device=x.device, dtype=torch.bfloat16)
    _lt_ext().lt_mm(x, w, out, state['c7_workspace'], trans, index)
    return out.t() if trans >= 2 else out


METHODS['c7lt'] = m_c7_lt
_c7_prepare = prepare
_c7_applicable = applicable


def applicable(method, layer, m, n, k, params=()):
    if method == 'c7lt':
        return layer in DENSE
    return _c7_applicable(method, layer, m, n, k, params)


def prepare(weights: Weights):
    state = _c7_prepare(weights)
    _lt_ext()
    state['c7_workspace'] = torch.empty(32 << 20, dtype=torch.uint8, device=state['router'].device)
    for layer in DENSE:
        if 't_' + layer not in state:
            state['t_' + layer] = state[layer].t().contiguous()
    return state

PLAN['qkv_gate'] = [(1, 'vector', (1,4)), (16, 'cluster', (8,3,2,128)), (32, 'c7lt', (3,8)), (1<<30, 'cublas', ())]
PLAN['pooler'] = [(1, 'vector', (1,8)), (16, 'c2packed', (16,32,512,2,4)), (32, 'c2packed', (32,32,512,2,3)), (64, 'trans', ()), (128, 'c7lt', (3,1)), (1<<30, 'c7lt', (0,3))]
PLAN['unpooler'] = [(1, 'vector', (1,8)), (8, 'flashinfer', ('tinygemm',)), (128, 'trans', ()), (1<<30, 'c7lt', (3,0))]