vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.gemm.warp_splitk

FlashInfer's cp.async-pipelined BF16 GEMM for M <= 32 whose four compute warps hold split-K partials (warp split-K), kernel unchanged; the autotune heuristics are dropped for per-shape tactics chosen by decode.py and compiled before capture.

The weight streams through two cp.async loader warps with an L2 evict_first hint (one tile is issued before the PDL wait), the activations through one or two loader warps after it; four mma.sync warps each own a K quarter of every stage and reduce through shared memory. Small output tiles (16 or 32 weight rows) give many CTAs that co-reside on an SM, which is what keeps a decode projection's weight stream at HBM speed.

Classes:

• PreparedWarpSplitK –

  out[M, N] = x[M, K] @ weight[N, K]^T for M <= 32 with a tactic per M;

Functions:

• compile_warp_splitk –

  Code for an [n, k] weight with dynamic M (at most 32 rows), PDL on.

PreparedWarpSplitK

out[M, N] = x[M, K] @ weight[N, K]^T for M <= 32 with a tactic per M; x and out packed row-major, 32 B aligned.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

class PreparedWarpSplitK:
    """``out[M, N] = x[M, K] @ weight[N, K]^T`` for M <= 32 with a tactic per M;
    ``x`` and ``out`` packed row-major, 32 B aligned."""

    def __init__(self, weight, tactics):
        n, k = weight.shape
        if weight.dtype != _torch.bfloat16 or weight.stride() != (k, 1):
            raise ValueError("warp split-K weights are packed BF16 [N, K]")
        if weight.data_ptr() % 32 or any(m > _MAX_M for m in tactics):
            raise ValueError("warp split-K: 32 B aligned weight, M <= 32")
        self.weight, self.tactics = weight, dict(tactics)
        index = weight.device.index
        self.codes = {
            m: compile_warp_splitk(index, n, k, t) for m, t in self.tactics.items()
        }
        self.weight_view = _view(weight)

    def run(self, x, out):
        code = self.codes[x.shape[0]]
        stream = _cuda.CUstream(_torch.cuda.current_stream(x.device).cuda_stream)
        # The kernel never reads its bias argument; a row of ``out`` fills it.
        bias = _view(out[0])
        code(self.weight_view, _view(x, dynamic_m=True), bias,
             _view(out, dynamic_m=True), stream)  # fmt: skip
        return out

_a_fragment_offset(lane, compute_warp, m_iter, phase, k_tile)

Element offset of one lane's A ldmatrix.x4 address within a stage.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

def _a_fragment_offset(lane, compute_warp, m_iter, phase, k_tile):
    """Element offset of one lane's A ``ldmatrix.x4`` address within a stage."""
    linear = lane % 16 + (lane // 16) * 128 + phase * 256
    row = linear % _MMA_SHAPE[0] + m_iter * _MMA_SHAPE[0]
    col = linear // _MMA_SHAPE[0] + compute_warp * (k_tile // _COMPUTE_WARPS)
    return row * k_tile + (col ^ ((row % 8) * 8))

_b_pair_offset(lane, compute_warp, n_iter, pair, token_tile, k_tile)

Element offset of one lane's B ldmatrix.x4 address within a stage.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

def _b_pair_offset(lane, compute_warp, n_iter, pair, token_tile, k_tile):
    """Element offset of one lane's B ``ldmatrix.x4`` address within a stage."""
    b_atom = lane % 8 + ((lane % 32) // 8) * (token_tile * 8)
    linear = b_atom + pair * 2 * token_tile * 16 + n_iter * 8
    row = linear % token_tile
    col = linear // token_tile + compute_warp * (k_tile // _COMPUTE_WARPS)
    return row * k_tile + (col ^ ((row % 8) * 8))

_cp_async_16_with_policy(dst, src, policy, *, loc=None, ip=None)

16-byte cp.async.cg global->shared carrying an L2 cache-policy hint.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

@dsl_user_op
def _cp_async_16_with_policy(dst, src, policy, *, loc=None, ip=None):
    """16-byte cp.async.cg global->shared carrying an L2 cache-policy hint."""
    llvm.inline_asm(
        None,
        [
            dst.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip),
            src.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip),
            policy.ir_value(loc=loc, ip=ip),
        ],
        "cp.async.cg.shared.global.L2::cache_hint [$0], [$1], 16, $2;",
        "r,l,l",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
    )

_issue_weight_tile(smem, weight, tidx, output_tile_idx, k_tile, stage, policy, output_tile, k_tile_size, with_policy, *, loc=None, ip=None)

Issue one A-loader thread's cp.async copies for weight K tile k_tile.

Plain Python control flow only: this helper is not preprocessed by the DSL. with_policy selects the .L2::cache_hint form; the pre-barrier prologue tile must use the plain form because ptxas 13.0 misencodes the hinted copy in that position (LDGSTS ... [R+UR], desc[UR<odd>]).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

@dsl_user_op
def _issue_weight_tile(
    smem,
    weight,
    tidx,
    output_tile_idx,
    k_tile,
    stage,
    policy,
    output_tile,
    k_tile_size,
    with_policy,
    *,
    loc=None,
    ip=None,
):
    """Issue one A-loader thread's cp.async copies for weight K tile ``k_tile``.

    Plain Python control flow only: this helper is not preprocessed by the DSL.
    ``with_policy`` selects the ``.L2::cache_hint`` form; the pre-barrier
    prologue tile must use the plain form because ptxas 13.0 misencodes the
    hinted copy in that position (``LDGSTS ... [R+UR], desc[UR<odd>]``).
    """
    copy_count = output_tile * k_tile_size // (8 * _A_LOADER_THREADS)
    for copy_idx in range(copy_count):
        linear = tidx * 8 + copy_idx * _A_LOADER_THREADS * 8
        row = linear // k_tile_size
        col = linear % k_tile_size
        dst_offset = (
            stage * output_tile * k_tile_size
            + row * k_tile_size
            + (col ^ ((row % 8) * 8))
        )
        src_offset = weight.layout(
            (
                output_tile_idx * output_tile + row,
                k_tile * k_tile_size + col,
            )
        )
        if with_policy:
            # CuTe pointers: ``toint()`` yields the 32-bit shared address for
            # SMEM and the 64-bit global address for GMEM, as the PTX expects.
            _cp_async_16_with_policy(
                smem.iterator + dst_offset,
                weight.iterator + src_offset,
                policy,
                loc=loc,
                ip=ip,
            )
        else:
            prims.cp_async_shared_global(
                smem.iterator.raw_ptr() + dst_offset,
                weight.iterator.raw_ptr() + src_offset,
                16,
                "cg",
                loc=loc,
                ip=ip,
            )

_l2_policy_evict_first(*, loc=None, ip=None)

64-bit L2 evict_first policy (fraction 1.0).

Must stay a runtime createpolicy: with the equivalent constant, ptxas 13.0 emits the invalid LDGSTS ... [R+UR], desc[UR1] encoding.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

@dsl_user_op
def _l2_policy_evict_first(*, loc=None, ip=None):
    """64-bit L2 evict_first policy (fraction 1.0).

    Must stay a runtime ``createpolicy``: with the equivalent constant, ptxas
    13.0 emits the invalid ``LDGSTS ... [R+UR], desc[UR1]`` encoding.
    """
    return cutlass.Int64(
        llvm.inline_asm(
            cutlass.Int64.mlir_type,
            [],
            "createpolicy.fractional.L2::evict_first.b64 $0, 1.0;",
            "=l",
            has_side_effects=False,
            is_align_stack=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
        )
    )

_ldmatrix_x4(smem, offset, *, loc=None, ip=None)

ldmatrix.x4 of the 16-bit fragment offset elements into smem.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

@dsl_user_op
def _ldmatrix_x4(smem, offset, *, loc=None, ip=None):
    """``ldmatrix.x4`` of the 16-bit fragment ``offset`` elements into ``smem``."""
    return prims.ldmatrix(
        smem.iterator.raw_ptr() + offset,
        num=4,
        layout=prims.MMALayout.ROW,
        loc=loc,
        ip=ip,
    )

compile_warp_splitk(device_index, n, k, tactic) cached

Code for an [n, k] weight with dynamic M (at most 32 rows), PDL on.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/warp_splitk.py

@functools.cache
def compile_warp_splitk(device_index: int, n: int, k: int, tactic: WarpSplitKTactic):
    """Code for an [n, k] weight with dynamic M (at most 32 rows), PDL on."""
    check_tactic(tactic, n, k)
    dtype = _torch.bfloat16
    with _torch.cuda.device(device_index):
        shapes: tuple[tuple[tuple[int, ...], bool], ...] = (
            ((n, k), False),
            ((_MAX_M, k), True),
            ((n,), False),
        )
        shapes += (((_MAX_M, n), True),)
        args = tuple(
            _view(_torch.empty(s, device="cuda", dtype=dtype), dynamic_m=d)
            for s, d in shapes
        )
        stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream)
        return cute_ext.compile(
            CpAsyncWarpSplitKKernel(tactic, True, False), *args, stream
        )