[Perf] Kimi K3 Specialized Decode: CuTeDSL Optimized GB300 (1.5-2x) (#61867)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

vllm/models/kimi_k3/nvidia/model.py

@@ -4,7 +4,7 @@
 
 import math
 from collections.abc import Hashable, Iterable
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from torch import nn
@@ -121,6 +121,12 @@ from ..common.mm_preprocess import (
     KimiK3MultiModalProcessor,
     KimiK3ProcessingInfo,
 )
+from .specialized_decode.runtime import specialized_decode_config
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.model_states.interface import ModelState
+
+    from .specialized_decode.runtime.decode import SpecializedDecode
 
 logger = init_logger(__name__)
 
@@ -1421,8 +1427,34 @@ class KimiLinearForCausalLM(
         self.logits_processor = LogitsProcessor(
             self.config.vocab_size, scale=logit_scale
         )
+        # decode_backend=latency_specialized: a specialized decode executor over
+        # self.model, built by the model state once the weights are loaded.
+        self.specialized_decode_config = specialized_decode_config(vllm_config)
+        self.specialized_decode: SpecializedDecode | None = None
+        # The final norm is applied in compute_logits unless the specialized decode
+        # applies it in forward.
+        self.needs_norm_in_logits = True
+
+    def build_specialized_decode(self) -> None:
+        from .specialized_decode.runtime.decode import SpecializedDecode
+
+        self.specialized_decode = SpecializedDecode(
+            self.model, self.specialized_decode_config
+        )
+        # From here on forward() returns normalized hidden states on every step: the
+        # specialized decode's last AttnRes applies the norm, and forward() applies it
+        # to each step that falls back to the stock forward (prefills, larger batches,
+        # PIECEWISE steps).
+        self.needs_norm_in_logits = False
+
+    def get_model_state_cls(self) -> "type[ModelState]":
+        from .model_state import KimiK3ModelState
+
+        return KimiK3ModelState
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        if self.specialized_decode is not None:
+            return self.specialized_decode.embed_input_ids(input_ids)
         return self.model.embed_input_ids(input_ids)
 
     def make_empty_intermediate_tensors(
@@ -1441,9 +1473,18 @@ class KimiLinearForCausalLM(
         inputs_embeds: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
-        return self.model(
+        decode = self.specialized_decode
+        if decode is not None:
+            step = decode.plan(positions, intermediate_tensors)
+            if step is not None:
+                return decode.run(step, input_ids, inputs_embeds)
+        hidden_states = self.model(
             input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
         )
+        if not self.needs_norm_in_logits:
+            assert decode is not None
+            hidden_states = decode.normalize(hidden_states)
+        return hidden_states
 
     @classmethod
     def get_mamba_state_dtype_from_config(
@@ -1502,7 +1543,8 @@ class KimiLinearForCausalLM(
     ) -> torch.Tensor | None:
         # The model's final norm is applied here (not at the end of forward) so
         # that the pre-norm hidden states can be fed to the MTP draft model.
-        hidden_states = self.model.norm(hidden_states, None)
+        if self.needs_norm_in_logits:
+            hidden_states = self.model.norm(hidden_states, None)
         return self.logits_processor(self.lm_head, hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1926,6 +1968,9 @@ class KimiK3ForConditionalGeneration(
     def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
         return self.language_model.compute_logits(hidden_states)
 
+    def get_model_state_cls(self) -> "type[ModelState]":
+        return self.language_model.get_model_state_cls()
+
     def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
         return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(
             input_buffers, **kwargs

vllm/models/kimi_k3/nvidia/model_state.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi K3's model runner state: the hybrid (KDA + MLA) state, plus the lifecycle of
+the specialized decode when ``decode_backend=latency_specialized`` selects it."""
+
+from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
+
+
+class KimiK3ModelState(MambaHybridModelState):
+    """Builds the specialized decode once the weights are loaded and the aux layers
+    set, and prepares it for every bound KV cache (the profiling cache, then the
+    real one)."""
+
+    def __init__(self, vllm_config, model, encoder_cache, device):
+        super().__init__(vllm_config, model, encoder_cache, device)
+        if hasattr(model, "get_language_model"):
+            model = model.get_language_model()
+        self.language_model = model
+        if model.specialized_decode_config is not None:
+            model.build_specialized_decode()
+
+    def initialize_kv_cache(self, kv_cache_config, block_tables):
+        super().initialize_kv_cache(kv_cache_config, block_tables)
+        decode = self.language_model.specialized_decode
+        if decode is not None:
+            decode.prepare_kv()

vllm/models/kimi_k3/nvidia/specialized_decode/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/kernel.py

@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Seven-CTA cluster AttnRes: one cluster per token, 1024 hidden columns per CTA.
+
+The delta is a tensor (read after the PDL wait) or ``slots`` Lamport mailbox slots
+summed in slot order, rounded to BF16 and re-armed. The final boundary rounds the
+mixture to BF16 before its RMSNorm. Prefix and saved blocks must be complete at launch.
+"""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Uint32, Uint64
+
+from ..common.mailbox import poll, rearm
+from ..common.primitives import (
+    alloc,
+    at,
+    bf16,
+    bf16x2,
+    dot,
+    fma,
+    group_barrier,
+    launch,
+    load_words,
+    push_float,
+    scatter_reduce,
+    store_words,
+    wait,
+    zeros,
+)
+
+
+@cute.jit
+def round_pair(t, row: cutlass.Constexpr, j: cutlass.Constexpr):
+    word = bf16x2(t[row, 2 * j], t[row, 2 * j + 1])
+    t[row, 2 * j], t[row, 2 * j + 1] = bf16(word, 0), bf16(word, 1)
+
+
+@cute.jit
+def unpack(words, out, row: cutlass.Constexpr):
+    for i in cutlass.range_constexpr(8):
+        out[row, i] = bf16(words[i // 2], i % 2)
+
+
+def block_ptr(blocks, token, block, col):
+    """``&blocks[token, block, col]`` from the tensor's strides."""
+    return blocks.iterator + token * blocks.stride[0] + block * blocks.stride[1] + col
+
+
+@cute.jit
+def rstd(sum_sq, eps):
+    """Reciprocal RMS of a 7168-wide row from its sum of squares."""
+    return cute.math.rsqrt(fma(sum_sq, Float32(1 / 7168), eps), fastmath=True)
+
+
+@cute.jit
+def dot8(a, ar: cutlass.Constexpr, b, br: cutlass.Constexpr):
+    return dot((a[ar, i] for i in range(8)), (b[br, i] for i in range(8)), Float32(0))
+
+
+@cute.jit
+def arm_exchanges(bars, counts):
+    """Thread 0 arms barrier ``r`` for ``counts[r]`` FP32 pushes from each of 7 CTAs."""
+    if cute.arch.thread_idx()[0] == 0:
+        for r, count in enumerate(counts):
+            if cutlass.const_expr(count > 0):
+                cute.arch.mbarrier_init(bars.iterator + r, 1)
+                cute.arch.mbarrier_arrive_and_expect_tx(bars.iterator + r, 28 * count)
+        cute.arch.mbarrier_init_fence()
+
+
+class ClusterAttnRes:
+    def __init__(self, nb, *, slots=1, has_delta=True, write=-1, final=False):
+        self.nb, self.slots, self.has_delta = nb, slots, has_delta
+        self.write, self.final = write, final
+        self.npre = 2 * nb
+        self.npost = (2 if final else 4) if nb else 1
+        self.kpre = 1 << max(self.npre - 1, 0).bit_length()
+        self.kpost = 1 << max(self.npost - 1, 0).bit_length()
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, (7, args[1].shape[0], 1), 128, 7)
+
+    @cute.jit
+    def exchange(
+        self, values, count: cutlass.Constexpr, kpad: cutlass.Constexpr, bufs, r, ids
+    ):
+        """Cluster-sum ``count`` per-thread partials into round ``r``'s total row."""
+        wb, peer_buf, total, bars = bufs
+        peer_buf, total, bar = peer_buf[r, None], total[r, None], bars.iterator + r
+        rank, _, tid = ids
+        lane, warp = tid % 32, tid // 32
+        reduced, index = scatter_reduce(values, lane, kpad)
+        if (lane & (32 // kpad - 1)) == 0 and index < count:
+            wb[warp * 16 + index] = reduced
+        group_barrier(0)
+        if tid < count:
+            value = (wb[tid] + wb[16 + tid]) + (wb[32 + tid] + wb[48 + tid])
+            for peer in cutlass.range_constexpr(7):
+                push_float(peer_buf.iterator + rank * 16 + tid, value, bar, Int32(peer))
+        wait(bar)
+        if tid < count:
+            sums = [Float32(0), Float32(0)]
+            for peer in cutlass.range_constexpr(7):
+                sums[peer % 2] += peer_buf[peer * 16 + tid]
+            total[tid] = sums[0] + sums[1]
+        group_barrier(0)
+
+    @cute.kernel
+    def kernel(self, args):
+        delta, prefix, blocks, norm, query, out_norm, output, eps, out_eps = args[0:9]
+        rank, token, _ = cute.arch.block_idx()
+        tid, _, _ = cute.arch.thread_idx()
+        bufs = self.setup(cutlass.utils.SmemAllocator())
+        cute.arch.cluster_arrive_relaxed()
+        cute.arch.griddepcontrol_launch_dependents()
+        state = self.state()
+        rows, ids = (prefix, blocks, norm, query, out_norm), (rank, token, tid)
+        self.pre(state, bufs, rows, eps, ids)
+        self.post(state, bufs, delta, rows, output, eps, out_eps, ids, args[9:])
+        cute.arch.griddepcontrol_wait()
+
+    def state(self):
+        """Carried from pre() to post(); the rest reloads so a GEMM can run between."""
+        return cute.make_rmem_tensor((1, 8), Float32), cute.make_rmem_tensor(2, Float32)
+
+    @cute.jit
+    def setup(self, smem):
+        warp_buf = alloc(smem, Float32, 64)
+        peer_buf = alloc(smem, Float32, (3, 112))
+        total = alloc(smem, Float32, (3, 16))
+        bars = alloc(smem, Uint64, 3, 8)
+        arm_exchanges(bars, (self.npre, self.npost, int(self.final)))
+        return warp_buf, peer_buf, total, bars
+
+    @cute.jit
+    def pre(self, state, bufs, rows, eps: Float32, ids):
+        """Inputs complete at launch, then the saved-block statistics exchange."""
+        prefix, blocks, norm, query, out_norm = rows
+        rank, token, tid = ids
+        warp_buf, peer_buf, total, bars = bufs
+        mixture, scal = state
+        # Rows: saved blocks, combined weight, gamma, old prefix, new prefix.
+        v = cute.make_rmem_tensor((self.nb + 4, 8), Float32)
+        col = (rank * 128 + tid) * 8
+        W = self.nb
+        for source in cutlass.range_constexpr(self.nb):
+            words = load_words(
+                block_ptr(blocks, token, source, col), count=4, mutable=True
+            )
+            unpack(words, v, source)
+        self.weights(v, norm, query, col)
+        cute.arch.cluster_wait()
+        maximum, exp_sum = Float32(0), Float32(0)
+        mixture.fill(0)
+        if cutlass.const_expr(self.nb > 0):
+            values = zeros(self.kpre, Float32)
+            for source in cutlass.range_constexpr(self.nb):
+                values[source] = dot8(v, source, v, source)
+                values[self.nb + source] = dot8(v, source, v, W)
+            self.exchange(values, self.npre, self.kpre, bufs, 0, ids)
+            logits = cute.make_rmem_tensor(self.nb, Float32)
+            for b in cutlass.range_constexpr(self.nb):
+                logits[b] = total[0, self.nb + b] * rstd(total[0, b], eps)
+            maximum = logits[0]
+            for source in cutlass.range_constexpr(1, self.nb):
+                maximum = cute.arch.fmax(maximum, logits[source])
+            for source in cutlass.range_constexpr(self.nb):
+                e = cute.math.exp(logits[source] - maximum, fastmath=True)
+                exp_sum += e
+                for i in cutlass.range_constexpr(8):
+                    mixture[0, i] = fma(e, v[source, i], mixture[0, i])
+        scal[0], scal[1] = maximum, exp_sum
+
+    @cute.jit
+    def weights(self, v, norm, query, col):
+        if cutlass.const_expr(self.nb > 0):
+            nw = load_words(norm.iterator + col, count=4)
+            qw = load_words(query.iterator + col, count=4)
+            for i in cutlass.range_constexpr(8):
+                v[self.nb, i] = bf16(nw[i // 2], i % 2) * bf16(qw[i // 2], i % 2)
+
+    @cute.jit
+    def post(
+        self, state, bufs, delta, rows, output, eps: Float32, out_eps: Float32, ids, aux
+    ):
+        """Add the delta, update prefix (and its ``aux`` copy) and saved block, write
+        the normalized output."""
+        prefix, blocks, norm, query, out_norm = rows
+        rank, token, tid = ids
+        warp_buf, peer_buf, total, bars = bufs
+        mixture, scal = state
+        maximum, exp_sum = scal[0], scal[1]
+        v = cute.make_rmem_tensor((self.nb + 4, 8), Float32)
+        col = (rank * 128 + tid) * 8
+        W, G, OLD, NEW = self.nb, self.nb + 1, self.nb + 2, self.nb + 3
+        self.weights(v, norm, query, col)
+        unpack(load_words(out_norm.iterator + col, count=4), v, G)
+        if cutlass.const_expr(self.has_delta):
+            unpack(load_words(at(prefix, token, col), count=4, mutable=True), v, OLD)
+        cur = cute.make_rmem_tensor((1, 8), Float32)
+        if cutlass.const_expr(self.slots == 0):
+            cute.arch.griddepcontrol_wait()
+            unpack(load_words(at(delta, token, col), count=4, mutable=True), cur, 0)
+        else:
+            cur.fill(0)
+            base = delta.iterator + token * 7168 + col
+            got = poll(base, delta.shape[1] * 7168, self.slots)
+            for slot in cutlass.range_constexpr(self.slots):
+                rearm(base + slot * delta.shape[1] * 7168)
+                part = cute.make_rmem_tensor((1, 8), Float32)
+                unpack(tuple(got[slot, k] for k in range(4)), part, 0)
+                for i in cutlass.range_constexpr(8):
+                    cur[0, i] = cur[0, i] + part[0, i]
+            if cutlass.const_expr(self.slots > 1):
+                for j in cutlass.range_constexpr(4):  # all-reduce output boundary
+                    round_pair(cur, 0, j)
+        words = cute.make_rmem_tensor(4, Uint32)
+        for j in cutlass.range_constexpr(4):
+            lo, hi = cur[0, 2 * j], cur[0, 2 * j + 1]
+            if cutlass.const_expr(self.has_delta):
+                lo, hi = v[OLD, 2 * j] + lo, v[OLD, 2 * j + 1] + hi
+            words[j] = bf16x2(lo, hi)
+        unpack(words, v, NEW)
+        values = zeros(self.kpost, Float32)
+        values[0] = dot8(v, NEW, v, NEW)
+        if cutlass.const_expr(self.nb > 0):
+            values[1] = dot8(v, NEW, v, W)
+            if cutlass.const_expr(not self.final):
+                values[2] = dot8(mixture, 0, v, NEW)
+                values[3] = dot8(mixture, 0, mixture, 0)
+        self.exchange(values, self.npost, self.kpost, bufs, 1, ids)
+        row = (words[0], words[1], words[2], words[3])
+        store_words(at(prefix, token, col), row)
+        for stream_out in aux:
+            store_words(at(stream_out, token, col), row)
+        if cutlass.const_expr(self.write >= 0):
+            store_words(block_ptr(blocks, token, self.write, col), row)
+        result = cute.make_rmem_tensor((1, 8), Float32)
+        if cutlass.const_expr(self.nb == 0):
+            scale = rstd(total[1, 0], out_eps)
+            for i in cutlass.range_constexpr(8):
+                result[0, i] = v[NEW, i] * scale * v[G, i]
+        else:
+            ss, dot = total[1, 0], total[1, 1]
+            logit = dot * rstd(ss, eps)
+            top = cute.arch.fmax(maximum, logit)
+            alpha = cute.math.exp(maximum - top, fastmath=True)
+            ep = cute.math.exp(logit - top, fastmath=True)
+            denominator = fma(alpha, exp_sum, ep)
+            if cutlass.const_expr(self.final):
+                inverse = Float32(1) / denominator
+                for j in cutlass.range_constexpr(4):
+                    for i in (2 * j, 2 * j + 1):
+                        result[0, i] = fma(
+                            alpha * inverse, mixture[0, i], ep * inverse * v[NEW, i]
+                        )
+                    round_pair(result, 0, j)
+                fin = cute.make_rmem_tensor(1, Float32)
+                fin[0] = dot8(result, 0, result, 0)
+                self.exchange(fin, 1, 1, bufs, 2, ids)
+                scale = rstd(total[2, 0], out_eps)
+                for i in cutlass.range_constexpr(8):
+                    result[0, i] = result[0, i] * scale * v[G, i]
+            else:
+                a2, t3 = alpha * alpha, total[1, 3]
+                ae2 = Float32(2) * alpha * ep
+                quad = fma(a2, t3, fma(ae2, total[1, 2], ep * ep * ss))
+                scale = rstd(quad, out_eps * denominator * denominator)
+                ca, cp = alpha * scale, ep * scale
+                for i in cutlass.range_constexpr(8):
+                    result[0, i] = fma(ca, mixture[0, i], cp * v[NEW, i]) * v[G, i]
+        for j in cutlass.range_constexpr(4):
+            words[j] = bf16x2(result[0, 2 * j], result[0, 2 * j + 1])
+        store_words(at(output, token, col), tuple(words[j] for j in range(4)))

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/prepared.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Prepared seven-CTA AttnRes consumers for ordinary or Lamport mailbox input."""
+
+from functools import cache
+from typing import Any
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import (
+    architecture,
+    dynamic_fake,
+    dynamic_view,
+    fake,
+    stream,
+    to_cute,
+)
+from .kernel import ClusterAttnRes
+
+
+@cache
+def compile_consumer(nb, delta, write, final, slots, cap, arch, aux=False):
+    row, weight, eps = dynamic_fake(2), fake((7168,), (1,)), cutlass.Float32(1e-6)
+    source = fake((slots, cap, 7168), (cap * 7168, 7168, 1)) if slots else row
+    kernel = ClusterAttnRes(nb, slots=slots, has_delta=delta, write=write, final=final)
+    args = (source, row, dynamic_fake(3), weight, weight, weight, row, eps, eps)
+    args += (row,) * aux  # the updated prefix row, also stored to an aux stream
+    return cute.compile(kernel, *args, make_fake_stream(), options=f"--gpu-arch={arch}")
+
+
+class PreparedAttnRes:
+    """Ordinary-input code plus the Lamport variant of its bound mailbox ring."""
+
+    def __init__(self, weights, key, eps):
+        self.device = weights[0].device
+        self.views = tuple(to_cute(w.reshape(-1)) for w in weights)
+        self.key, self.eps = key, tuple(map(cutlass.Float32, eps))
+        self.code = self._compile(0, 0)
+        self.ring: Any = None
+        self.aux: Any = None
+
+    def _compile(self, slots, cap, aux=False):
+        arch = architecture(self.device)
+        return compile_consumer(*self.key, slots, cap, arch, aux)
+
+    @classmethod
+    def for_target(cls, target, *, final=False, has_delta=True):
+        if final:
+            norm, query = target.output_attn_res_norm, target.output_attn_res_proj
+            output, nb, write = target.norm, target.num_attn_res_blocks, -1
+        else:
+            norm = target.self_attention_res_norm
+            query, output = target.self_attention_res_proj, target.input_layernorm
+            nb = target.prev_valid_blocks
+            write = target.block_write_idx if target.is_block_write_layer else -1
+        weights = (norm.weight, query.weight, output.weight)
+        eps = (norm.variance_epsilon, output.variance_epsilon)
+        return cls(weights, (nb, has_delta, write, final), eps)
+
+    def bind_ring(self, mailbox):
+        """Lamport input: sum the [slots, capacity, 7168] mailbox rows, then re-arm."""
+        shape = tuple(mailbox.shape[:2])
+        self.ring = self._compile(*shape), to_cute(mailbox), shape
+
+    def bind_aux(self):
+        """Variants that also store the updated prefix rows (a stock aux stream)."""
+        ring = None
+        if self.ring is not None:
+            slots, cap = self.ring[2]
+            ring = self._compile(slots, cap, True)
+        self.aux = self._compile(0, 0, True), ring
+
+    def run_ring(self, prefix, blocks, out, aux=None):
+        code = self.ring[0] if aux is None else self.aux[1]
+        return self._launch(code, self.ring[1], prefix, blocks, out, aux)
+
+    def run(self, delta, prefix, blocks, out, aux=None):
+        code = self.code if aux is None else self.aux[0]
+        return self._launch(code, dynamic_view(delta), prefix, blocks, out, aux)
+
+    def _launch(self, code, source, prefix, blocks, out, aux=None):
+        rows = (prefix, blocks, out) + (() if aux is None else (aux,))
+        prefix, blocks, view, *aux = map(dynamic_view, rows)
+        args = *self.views, view, *self.eps, *aux
+        code(source, prefix, blocks, *args, stream(self.device))
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/embedding.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Local CuTe vocabulary gather plus the ordinary framework TP all-reduce."""
+
+from functools import cache
+
+import cutlass as c
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_stream
+
+from vllm.distributed import get_tp_group
+
+from .host import architecture, fake
+from .primitives import launch
+
+
+class EmbeddingGather:
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, (7, args[0].shape[0], 1), 128, pdl=False)
+
+    @cute.kernel
+    def kernel(self, args):
+        ids, weight, out, start, vocab = args[0:5]
+        lane, _, _ = cute.arch.thread_idx()
+        tile, token, _ = cute.arch.block_idx()
+        token_id = c.Int64(ids[token])
+        row, column = token_id - start, tile * 1024 + lane * 8
+        values = cute.make_rmem_tensor((8,), c.BFloat16)
+        values.fill(c.BFloat16(0))
+        if (token_id >= 0) & (token_id < vocab) & (row >= 0) & (row < weight.shape[0]):
+            source = weight.iterator + row * 7168 + column
+            cute.autovec_copy(cute.make_tensor(source, cute.make_layout((8,))), values)
+        destination = out.iterator + token * 7168 + column
+        cute.autovec_copy(values, cute.make_tensor(destination, cute.make_layout((8,))))
+
+
+@cache
+def compile_gather(dtype, shard, arch):
+    batch = cute.sym_int32()
+    args: tuple = (fake((batch,), None, dtype, dtype.itemsize), fake((shard, 7168)))
+    args += (fake((batch, 7168)), c.Int64(0), c.Int64(0))
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    options = f"--gpu-arch={arch} --enable-tvm-ffi"
+    return cute.compile(EmbeddingGather(), *args, stream, options=options)
+
+
+class PreparedEmbedding:
+    def __init__(self, embedding, tp, rank):
+        weight = embedding.weight
+        # The gather moves 7 x 1024 BF16 columns of contiguous rows.
+        if weight.dim() != 2 or weight.shape[1] != 7168 or not weight.is_contiguous():
+            raise ValueError(
+                f"embedding shard must be a contiguous [V, 7168] tensor, got "
+                f"{tuple(weight.shape)}"
+            )
+        if weight.dtype != torch.bfloat16:
+            raise ValueError(f"embedding shard must be BF16, got {weight.dtype}")
+        self.group, self.weight = get_tp_group(), weight.detach()
+        self.start, self.vocab = rank * weight.shape[0], embedding.org_vocab_size
+        with torch.accelerator.device_index(weight.device.index):
+            arch, ids = architecture(weight.device), (torch.int32, torch.int64)
+            self.kernels = {t: compile_gather(t, weight.shape[0], arch) for t in ids}
+
+    def run(self, ids, out):
+        if ids.dim() != 1 or ids.dtype not in self.kernels:
+            raise ValueError(
+                f"token ids must be 1-D int32/int64, got {ids.dtype} {tuple(ids.shape)}"
+            )
+        if out.shape != (ids.shape[0], 7168) or not out.is_contiguous():
+            raise ValueError(
+                f"embedding output must be contiguous [{ids.shape[0]}, 7168]"
+            )
+        self.kernels[ids.dtype](ids, self.weight, out, self.start, self.vocab)
+        reduced = self.group.all_reduce(out)
+        if reduced is not out:
+            out.copy_(reduced)
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Host-side helpers for the prepared CuTe providers."""
+
+import ctypes
+import math
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_tensor
+
+from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail.primitives import to_cute
+
+_NAMES = ("BFloat16", "Float32", "Uint8", "Uint32", "Int32", "Int64")
+_DTYPES = {getattr(torch, name.lower()): getattr(cutlass, name) for name in _NAMES}
+
+
+def layout(tensor):
+    return tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype
+
+
+def stream(device):
+    return cuda.CUstream(torch.cuda.current_stream(device).cuda_stream)
+
+
+def architecture(device):
+    major, minor = torch.cuda.get_device_capability(device)
+    return f"sm_{major}{minor}a"
+
+
+def fake(shape, stride=None, dtype=torch.bfloat16, align=16):
+    """Row-major unless ``stride`` is given."""
+    stride = stride or tuple(math.prod(shape[i + 1 :]) for i in range(len(shape)))
+    dtype = _DTYPES.get(dtype, dtype)
+    return make_fake_tensor(dtype, shape, stride, assumed_align=align)
+
+
+def dynamic_fake(rank, dtype=torch.bfloat16):
+    shape = tuple(cute.sym_int() for _ in range(rank))
+    return fake(shape, (*(cute.sym_int(64) for _ in range(rank - 1)), 1), dtype)
+
+
+def dynamic_view(tensor, align=16):
+    return to_cute(tensor, align).mark_layout_dynamic(leading_dim=tensor.ndim - 1)
+
+
+def tensor_maps(specs, device):
+    """TMA descriptors (128 B each) in one device tensor.
+    Each spec: (pointer, dtype, dims, strides, box, swizzle, l2 promotion).
+
+    The bytes are uploaded with a copy on the current stream: the new tensor may
+    reuse a block freed by a kernel still pending on that stream (e.g. a warmup
+    launch's output), which a copy outside the stream would race with."""
+    encoded = bytearray()
+    with torch.accelerator.device_index(torch.device(device).index):
+        for pointer, dtype, dims, strides, box, swizzle, l2 in specs:
+            u64 = ([cuda.cuuint64_t(v) for v in values] for values in (dims, strides))
+            u32 = ([cuda.cuuint32_t(v) for v in box], [cuda.cuuint32_t(1)] * len(dims))
+            interleave = cuda.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE
+            fill = cuda.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
+            status, tensor_map = cuda.cuTensorMapEncodeTiled(
+                dtype, len(dims), pointer, *u64, *u32, interleave, swizzle, l2, fill
+            )
+            if status != cuda.CUresult.CUDA_SUCCESS:
+                raise RuntimeError(f"cuTensorMapEncodeTiled failed: {status}")
+            encoded += (ctypes.c_ubyte * 128).from_address(tensor_map.getPtr())
+    return torch.frombuffer(encoded, dtype=torch.uint8).to(device)
+
+
+def _checked(result, operation="CUDA driver call"):
+    """The values of a cuda-python driver result, raising on a failed status."""
+    status, *values = result
+    if status != cuda.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"{operation} failed: {status}")
+    return tuple(values)
+
+
+class PreparedLaunchAttributes:
+    def __init__(self, compiled, executor):
+        self.compiled, self.executor = compiled, executor
+
+    def __call__(self, *args, **kwargs):
+        return self.compiled(*args, **kwargs)
+
+
+def resolve_kernel(compiled, device):
+    executor = compiled.to(device)
+    library = compiled.library
+    (count,) = _checked(cuda.cuLibraryGetKernelCount(library))
+    if count != 1:  # every prepared module compiles one kernel
+        raise RuntimeError(
+            f"expected one kernel in the compiled library, found {count}"
+        )
+    (kernels,) = _checked(cuda.cuLibraryEnumerateKernels(count, library))
+    (function,) = _checked(cuda.cuKernelGetFunction(kernels[0]))
+    return executor, function
+
+
+def check_residency(compiled, device_index: int, grid: int, threads: int, what: str):
+    """Persistent grids whose CTAs wait on one another (not only within a cluster) must
+    all be resident at once: raise when the device cannot hold ``grid`` CTAs of
+    ``compiled`` (``threads`` per CTA, its dynamic shared memory) at the same time."""
+    _, function = resolve_kernel(compiled, device_index)
+    attribute = cuda.CUfunction_attribute
+    max_dynamic = attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES
+    (smem,) = _checked(cuda.cuFuncGetAttribute(max_dynamic, function))
+    (per_sm,) = _checked(
+        cuda.cuOccupancyMaxActiveBlocksPerMultiprocessor(function, threads, smem)
+    )
+    sms = torch.cuda.get_device_properties(device_index).multi_processor_count
+    if grid > per_sm * sms:
+        raise RuntimeError(
+            f"{what}: {grid} CTAs cannot be co-resident ({per_sm} per SM x {sms} SMs)"
+        )
+
+
+def prepare_shared_carveout(compiled, device_index: int, value: int):
+    """Restore ``value`` (CuTe's launcher drops preferred_smem_carveout) as a driver
+    function attribute; call once before capture."""
+    executor, function = resolve_kernel(compiled, device_index)
+    attribute = cuda.CUfunction_attribute
+    attribute = attribute.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT
+    _checked(cuda.cuFuncSetAttribute(function, attribute, value))
+    return PreparedLaunchAttributes(compiled, executor)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lamport mailboxes: BF16 rows published in 16-byte fragments; a reader polls until no
+fragment holds the -0.0 sentinel, then re-arms."""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Uint32
+
+from .primitives import load_words, published, store_relaxed
+
+SENTINEL = 0x80000000
+
+
+def rearm(pointer):
+    store_relaxed(pointer, (Uint32(SENTINEL),) * 4)
+
+
+@cute.jit
+def fetch(values, i: cutlass.Constexpr[int], pointer):
+    """Load fragment ``i`` into ``values[i]``; returns its dirty flag at bit ``i``."""
+    got = load_words(pointer, count=4, volatile=True)
+    for k in cutlass.range_constexpr(4):
+        values[i, k] = got[k]
+    return Uint32(not published(got)) << i
+
+
+@cute.jit
+def poll(base, stride, count: cutlass.Constexpr[int]):
+    """Lamport-poll ``count`` 16-byte fragments at base + i*stride (elements); all
+    loads issue before any check and only dirty fragments re-poll."""
+    values = cute.make_rmem_tensor((count, 4), Uint32)
+    pending = Uint32(0)
+    for i in cutlass.range_constexpr(count):
+        pending = pending | fetch(values, i, base + i * stride)
+    while pending != 0:
+        again = Uint32(0)
+        for i in cutlass.range_constexpr(count):
+            if (pending >> i) & 1:
+                again = again | fetch(values, i, base + i * stride)
+        pending = again
+    return values

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/policy.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-batch policy tables."""
+
+
+def batch_table(spec):
+    """Per-batch table of integers ("1-2:96,3-8:120"), tuples ("16:64.16.4.8") or names
+    ("16-64:trtllm"); later entries override earlier ones."""
+    table = {}
+    for item in filter(None, spec.split(",")):
+        span, value = item.split(":")
+        lo, _, hi = span.partition("-")
+        if "." in value:
+            value = tuple(map(int, value.split(".")))
+        elif value.lstrip("-").isdigit():
+            value = int(value)
+        table.update({b: value for b in range(int(lo), int(hi or lo) + 1)})
+    return table

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/prefetch.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""L2 weight prefetch: each warp walks a persistent device table of (address, bytes),
+spreading 16-KiB chunks over the grid in table order. The owner retains the weights
+and table through graph replay; activations and mutable caches may not be ranges
+before their producer completes."""
+
+import functools
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_stream
+
+from .host import fake
+from .primitives import launch, prefetch_bulk, thread
+
+
+class L2Prefetch:
+    def __init__(self, grid, carveout):
+        self.grid, self.carveout = grid, carveout
+
+    @cute.jit
+    def __call__(self, *args):
+        options = dict(pdl=False, preferred_smem_carveout=self.carveout)
+        launch(self.kernel, args, self.grid, 32, **options)
+
+    @cute.kernel
+    def kernel(self, args):
+        (ranges,) = args
+        cute.arch.griddepcontrol_launch_dependents()
+        tid, _, _, block = thread()
+        worker = cutlass.Int64(block * 32 + tid)
+        workers = cutlass.Int64(self.grid * 32)
+        phase, row = cutlass.Int64(0), cutlass.Int32(0)
+        while row < ranges.shape[0]:
+            base, nbytes = ranges[row, 0], ranges[row, 1]
+            if (base != 0) & (nbytes > 0):
+                chunks = (nbytes + 16384 - 1) // 16384
+                index = worker - phase % workers
+                if index < 0:
+                    index += workers
+                while index < chunks:
+                    offset = index * 16384
+                    size = nbytes - offset
+                    if size > 16384:
+                        size = cutlass.Int64(16384)
+                    size = size & cutlass.Int64(-16)
+                    if size > 0:
+                        address = (base + offset) & cutlass.Int64(-16)
+                        prefetch_bulk(address, size.to(cutlass.Int32))
+                    index += workers
+                phase += chunks
+            row += 1
+        # Early exit must not let a later grid bypass its producer.
+        cute.arch.griddepcontrol_wait()
+
+
+@functools.cache
+def compile_kernel(count, grid, carveout, device_index):
+    ranges = fake((count, 2), None, torch.int64)
+    with torch.accelerator.device_index(device_index):
+        return cute.compile(L2Prefetch(grid, carveout), ranges, make_fake_stream())

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py

@@ -0,0 +1,928 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inline-PTX, memory, synchronization and math helpers of the decode kernels."""
+
+import math
+from functools import partial, wraps
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Float64, Int32, Int64, Uint8, Uint16, Uint32, Uint64
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass.experimental import primitives as nvvm
+
+
+def ptx(assembly, constraints, operands=(), result=None, *, effects, loc=None, ip=None):
+    """One verbatim asm block; a tuple ``result`` (even a 1-tuple) returns a tuple."""
+    if isinstance(result, tuple) and len(result) == 1:
+        args = (assembly, constraints, operands, result[0])
+        return (ptx(*args, effects=effects, loc=loc, ip=ip),)
+    if result is None:
+        ir_type = None
+    elif isinstance(result, tuple):
+        ir_type = llvm.StructType.get_literal([dtype.mlir_type for dtype in result])
+    else:
+        ir_type = result.mlir_type
+    value = llvm.inline_asm(
+        ir_type,
+        [operand.ir_value(loc=loc, ip=ip) for operand in operands],
+        assembly,
+        constraints,
+        has_side_effects=effects,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+    if result is None:
+        return None
+    if isinstance(result, tuple):
+        return tuple(
+            dtype(llvm.extractvalue(dtype.mlir_type, value, [i], loc=loc, ip=ip))
+            for i, dtype in enumerate(result)
+        )
+    return result(value)
+
+
+def asm(result=None, *, pure=False):
+    """``f(*args) -> (code, constraints, operands[, result, pure])`` as one inline-PTX
+    op. Pointer operands pass their address: 32-bit for "r" (shared), 64-bit for "l"."""
+
+    def wrap(fn):
+        @dsl_user_op
+        @wraps(fn)
+        def op(*args, loc=None, ip=None, **kwargs):
+            code, constraints, operands, *rest = fn(*args, **kwargs)
+            res, fx = (*rest, pure)[:2] if rest else (result, pure)
+            kinds = [c for c in constraints.split(",") if c[:1] not in ("=", "~", "")]
+            values = []
+            for kind, v in zip(kinds, operands):
+                if hasattr(v, "toint"):
+                    v = v.toint(loc=loc, ip=ip)
+                    v = Int32(v) if kind == "r" else v
+                values.append(v)
+            return ptx(code, constraints, values, res, effects=not fx, loc=loc, ip=ip)
+
+        return op
+
+    return wrap
+
+
+def _vector(count, first=0):
+    """(".vN", "{$first,...}") of an N-word access; ("", "$first") for one word."""
+    names = ",".join(f"${first + i}" for i in range(count))
+    return (f".v{count}", "{" + names + "}") if count > 1 else ("", names)
+
+
+div_rn = partial(cute.math.div, rounding=cute.math.RoundingMode.NEAREST_EVEN)
+
+
+@asm(Float32, pure=True)
+def add_rn(a, b):
+    return "add.rn.f32 $0,$1,$2;", "=f,f,f", (Float32(a), Float32(b))
+
+
+@asm(Float32, pure=True)
+def routing_add(a: Float32, b: Float32):
+    return "add.ftz.f32 $0,$1,$2;", "=f,f,f", (a, b)
+
+
+@asm(Float32, pure=True)
+def fma(a, b, c, *, ftz=False):
+    return f"fma.rn{'.ftz' if ftz else ''}.f32 $0, $1, $2, $3;", "=f,f,f,f", (a, b, c)
+
+
+def dot(a, b, acc=None, **kwargs):
+    """``fma`` chain over ``zip(a, b)`` from ``acc``, or from ``a[0] * b[0]``."""
+    for x, y in zip(a, b):
+        acc = x * y if acc is None else fma(x, y, acc, **kwargs)
+    return acc
+
+
+@asm(Float32, pure=True)
+def fma_bf16(a, b, c, *, high=False):
+    half = "h" if high else "l"
+    code = "{ .reg .b16 al,ah,bl,bh; mov.b32 {al,ah},$1; mov.b32 {bl,bh},$2; "
+    code += f"fma.rn.f32.bf16 $0,a{half},b{half},$3; }}"
+    return code, "=f,r,r,f", (Uint32(a), Uint32(b), Float32(c))
+
+
+@asm(Uint32, pure=True)
+def round_pair_square(word):
+    return "mul.rn.bf16x2 $0,$1,$1;", "=r,r", (Uint32(word),)
+
+
+@asm(Float32, pure=True)
+def inv_rms(value, eps, *, m8=False):
+    """rsqrt(value / 3584 + eps); ``m8``: the mean as a reciprocal product refined by
+    its FMA residual."""
+    mean = "div.rn.f32 r,$1,0f45600000;"
+    if m8:
+        mean = "mul.f32 q,$1,0f39924925; neg.f32 r,q; fma.rn.f32 r,r,0f45600000,$1; "
+        mean += "fma.rn.f32 r,r,0f39924925,q;"
+    code = f"{{ .reg .f32 {'q,' * m8}r; {mean} add.f32 r,r,$2; "
+    code += "rsqrt.approx.ftz.f32 $0,r; }"
+    return code, "=f,f,f", (Float32(value), Float32(eps))
+
+
+@asm(Float32, pure=True)
+def query_norm_rsqrt(total: Float32, eps: Float32):
+    code = "{ .reg .f32 mean; div.rn.f32 mean,$1,0f44c00000; "
+    code += "add.f32 mean,mean,$2; rsqrt.approx.f32 $0,mean; }"
+    return code, "=f,f,f", (total, eps)
+
+
+@asm(Float64, pure=True)
+def reciprocal64(value: Float64):
+    step = "fma.rn.f64 e,n,r,0d3ff0000000000000; fma.rn.f64 {},r,e,r; "
+    code = "{ .reg .f64 r,e,n; neg.f64 n,$1; rcp.approx.ftz.f64 r,$1; "
+    code += step.format("r") * 2 + step.format("$0") + "}"
+    return code, "=d,d", (value,)
+
+
+@asm(Float32, pure=True)
+def sigmoid_reciprocal(denominator: Float32, *, fast_math):
+    """1 / (1 + exp(-x)) from its denominator: an approximate divide with ``fast_math``,
+    else the approximate reciprocal refined by two IEEE FMAs (out-of-range: div.rn)."""
+    if fast_math:
+        return "div.approx.ftz.f32 $0, 0f3f800000, $1;", "=f,f", (denominator,)
+    code = (
+        "{\n.reg .f32 r,error,neg_error; .reg .b32 bits; .reg .pred regular;\n"
+        "mov.b32 bits,$1; add.u32 bits,bits,0x01800000; and.b32 bits,bits,0x7f800000;\n"
+        "setp.gt.u32 regular,bits,0x01ffffff; rcp.approx.ftz.f32 r,$1;\n"
+        "fma.rn.f32 error,$1,r,0fbf800000; neg.f32 neg_error,error;\n"
+        "fma.rn.f32 $0,r,neg_error,r; @!regular div.rn.f32 $0,0f3f800000,$1;\n}"
+    )
+    return code, "=f,f", (denominator,)
+
+
+@cute.jit
+def sigmoid(value: Float32, fast_math: cutlass.Constexpr = None):
+    """1 / (1 + exp(-x)); correctly rounded when ``fast_math`` is None."""
+    if cutlass.const_expr(fast_math is None):
+        result = div_rn(Float32(1), Float32(1) + cute.math.exp(-value, fastmath=False))
+    else:
+        denominator = Float32(1) + cute.math.exp(-value, fastmath=True)
+        result = sigmoid_reciprocal(denominator, fast_math=fast_math)
+    return result
+
+
+@asm(Float32, pure=True)
+def routing_sigmoid(value: Float32):
+    code = "{ .reg .f32 h,t; mul.ftz.f32 h,$1,0f3f000000; tanh.approx.f32 t,h; "
+    code += "fma.rn.ftz.f32 $0,t,0f3f000000,0f3f000000; }"
+    return code, "=f,f", (value,)
+
+
+@asm(Float32, pure=True)
+def routing_weight(score: Float32, total: Float32, scale: Float32, *, renorm=True):
+    code = "mul.ftz.f32 $0,$3,$1;"
+    if renorm:
+        code = (
+            "{ .reg .f32 n,d; mul.ftz.f32 n,$3,$1; "
+            "add.ftz.f32 d,0f1e3ce508,$2; div.approx.ftz.f32 $0,n,d; }"
+        )
+    return code, "=f,f,f,f", (score, total, scale)
+
+
+@asm(Float32, pure=True)
+def situ(g: Float32, u: Float32, beta_lb: Float32, inv_beta: Float32, inv_lb: Float32):
+    code = "{ .reg .f32 tg,tu,ex,rc,z; mul.f32 tg,$1,$4; tanh.approx.f32 tg,tg; "
+    code += "mul.f32 tu,$2,$5; tanh.approx.f32 tu,tu; mul.f32 z,$1,0fbfb8aa3b; "
+    code += "ex2.approx.ftz.f32 ex,z; add.f32 ex,ex,0f3f800000; "
+    code += "rcp.approx.ftz.f32 rc,ex; mul.f32 z,$3,tg; mul.f32 z,z,rc; "
+    code += "mul.f32 $0,z,tu; }"
+    return code, "=f,f,f,f,f,f", (g, u, beta_lb, inv_beta, inv_lb)
+
+
+@cute.jit
+def shared_situ(gate, up, beta, inverse_beta, linear_beta, inverse_linear_beta):
+    denominator = Float32(1) + cute.math.exp(-gate, fastmath=False)
+    sigmoid = Float32(0)
+    if denominator != Float32(float("inf")):
+        sigmoid = Float32(reciprocal64(Float64(denominator)))
+    g = (beta * cute.math.tanh(gate * inverse_beta, fastmath=False)) * sigmoid
+    u = up
+    if linear_beta > 0:
+        divided = Float32(Float64(up) * inverse_linear_beta)
+        u = linear_beta * cute.math.tanh(divided, fastmath=False)
+    return g * u
+
+
+def bits_float(bits):
+    return Float32(llvm.bitcast(T.f32(), Uint32(bits).ir_value()))
+
+
+def float_bits(value: Float32) -> Int32:
+    return Int32(llvm.bitcast(T.i32(), value.ir_value()))
+
+
+def bf16(word, high=False):
+    """FP32 of the low (or high) BF16 half of ``word``."""
+    return bits_float(word & 0xFFFF0000 if high else word << 16)
+
+
+@asm((Float32, Float32), pure=True)
+def unpack_bf16(word: Uint32):
+    return "shl.b32 $0,$2,16; and.b32 $1,$2,0xffff0000;", "=f,=f,r", (word,)
+
+
+@asm((Float32, Float32), pure=True)
+def fma2(a, b, c):
+    """Lanewise a * b + c of FP32 pairs in one packed instruction (IEEE per lane)."""
+    code = "{ .reg .b64 a,b,c,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; "
+    code += "mov.b64 c,{$6,$7}; fma.rn.f32x2 d,a,b,c; mov.b64 {$0,$1},d; }"
+    return code, "=f,=f,f,f,f,f,f,f", (*a, *b, *c)
+
+
+@asm((Float32, Float32), pure=True)
+def mul2(a, b):
+    """Lanewise a * b of FP32 pairs in one packed instruction (IEEE per lane)."""
+    code = "{ .reg .b64 a,b,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; "
+    code += "mul.rn.f32x2 d,a,b; mov.b64 {$0,$1},d; }"
+    return code, "=f,=f,f,f,f,f", (*a, *b)
+
+
+@asm(Uint16, pure=True)
+def gate_bf16_bits(value: Float32):
+    return "cvt.rn.bf16.f32 $0, $1;", "=h,f", (value,)
+
+
+@asm(Uint32, pure=True)
+def pack_pair(low, high, *, policy="plain"):
+    """BF16x2 of (low, high); ``policy=None`` is the bare conversion."""
+    code = "cvt.rn.bf16x2.f32 $0,$2,$1; "
+    if policy == "stage":
+        # NaN halves become canonical, so no word is the 0xffffffff staging sentinel.
+        code += (
+            "and.b32 bits,$0,0x7fff; setp.gt.u32 p,bits,0x7f80; "
+            "@p and.b32 $0,$0,0xffff0000; @p or.b32 $0,$0,0x7fc0; "
+            "and.b32 bits,$0,0x7fff0000; setp.gt.u32 p,bits,0x7f800000; "
+            "@p and.b32 $0,$0,0xffff; @p or.b32 $0,$0,0x7fc00000; "
+        )
+    if policy == "latent":
+        # Negative BF16 zero is the downstream mailbox's unpublished sentinel.
+        code += (
+            "and.b32 bits,$0,0xffff; setp.eq.u32 p,bits,0x8000; "
+            "@p and.b32 $0,$0,0xffff0000; shr.u32 bits,$0,16; "
+            "setp.eq.u32 p,bits,0x8000; @p and.b32 $0,$0,0xffff; "
+        )
+    if policy is not None:
+        code = "{ .reg .b32 bits; .reg .pred p; " + code + "}"
+    return code, "=r,f,f", (Float32(low), Float32(high))
+
+
+bf16x2 = partial(pack_pair, policy=None)
+
+
+@cute.jit
+def no_neg_zero(word):
+    """``word`` with each negative-zero BF16 half made +0."""
+    result = Uint32(word)
+    if (result & 0xFFFF) == 0x8000:
+        result &= Uint32(0xFFFF0000)
+    if (result >> 16) == 0x8000:
+        result &= Uint32(0xFFFF)
+    return result
+
+
+def published(words, ready=True, *, sentinel=0x80000000):
+    """``ready`` and no word is the mailbox ``sentinel``."""
+    for word in words:
+        ready &= word != Uint32(sentinel)
+    return ready
+
+
+@cute.jit
+def front_dirty(word):
+    return ((word & 0xFFFF) == 0xFFFF) | ((word >> 16) == 0xFFFF)
+
+
+def fp8_bits(value: Float32):
+    code = (
+        "{ .reg .b16 pair; cvt.rn.satfinite.e4m3x2.f32 pair, 0f00000000, $1; "
+        "cvt.u32.u16 $0, pair; }"
+    )
+    return Uint8(ptx(code, "=r,f", (value,), Uint32, effects=False))
+
+
+@asm()
+def store_fp8_pair(pointer, low: Float32, high: Float32):
+    """FP8 E4M3 (satfinite) of ``low`` and ``high`` as the two bytes at ``pointer``."""
+    code = "{ .reg .b16 pair; cvt.rn.satfinite.e4m3x2.f32 pair,$2,$1; "
+    code += "st.global.b16 [$0],pair; }"
+    return code, "l,f,f,~{memory}", (pointer, low, high)
+
+
+@asm(Uint32, pure=True)
+def pack_fp8(values):
+    code = "{ .reg .b16 lo,hi; cvt.rn.satfinite.e4m3x2.f32 lo,$2,$1; "
+    code += "cvt.rn.satfinite.e4m3x2.f32 hi,$4,$3; mov.b32 $0,{lo,hi}; }"
+    return code, "=r,f,f,f,f", values
+
+
+@cute.jit
+def mx_exponent(amax: Float32, floor: cutlass.Constexpr[bool] = False):
+    bits = Uint32(float_bits(amax))
+    e = Int32(bits >> 23)
+    result = Int32(0)
+    if cutlass.const_expr(floor):
+        if e == 255:
+            result = 254
+        elif e > 8:
+            result = e - 8
+    else:
+        if e != 0:
+            result = min(max(e - 8 + Int32(bits & 8388607 > 6291456), 0), 254)
+    return result
+
+
+@cute.jit
+def mx_rescale(scale: Int32):
+    e = 254 - scale
+    return bits_float(Uint32(e) << 23) if e >= 1 else Float32(0)
+
+
+@cute.jit
+def warp_sum(value: Float32, lanes: cutlass.Constexpr[int] = 32):
+    for shift in cutlass.range_constexpr(lanes.bit_length() - 2, -1, -1):
+        value += cute.arch.shuffle_sync_bfly(value, 1 << shift)
+    return value
+
+
+@asm(Uint32, pure=True)
+def redux(value, op="min"):
+    return f"redux.sync.{op}.u32 $0,$1,0xffffffff;", "=r,r", (Uint32(value),)
+
+
+def warp_max_abs(value):
+    return bits_float(redux(float_bits(abs(value)), "max"))
+
+
+@asm((Float32,) * 4, pure=True)
+def mma(a, b, accumulator, *, fp8=False, bf16=False):
+    """m16n8k16 (FP8: k32) mma.sync."""
+    kind = "e4m3" if fp8 else "bf16" if bf16 else "f16"
+    code = f"mma.sync.aligned.m16n8k{32 if fp8 else 16}.row.col.f32.{kind}.{kind}.f32 "
+    code += "{$0,$1,$2,$3},{$4,$5,$6,$7},{$8,$9},{$10,$11,$12,$13};"
+    return code, "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", (*a, *b, *accumulator)
+
+
+def mma_into(acc, a, b, **kwargs):
+    """``acc += a @ b`` with one ``mma``; ``acc`` is a four-element register view."""
+    result = mma(a, b, tuple(acc[k] for k in range(4)), **kwargs)
+    for k in range(4):
+        acc[k] = result[k]
+
+
+@asm()
+def ldmatrix(pointer, *, trans=False, count=4):
+    """``count`` 8x8 b16 shared matrices; lanes 8i..8i+7 address matrix i's rows."""
+    names = ",".join(f"${i}" for i in range(count))
+    code = f"ldmatrix.sync.aligned.m8n8.x{count}{'.trans' * trans}.shared.b16 "
+    code += f"{{{names}}},[${count}];"
+    constraints = ",".join(["=r"] * count + ["r"])
+    return code, constraints, (pointer,), (Uint32,) * count
+
+
+def split(total, i, n):
+    """(first, count) of part ``i`` of ``total`` items split ``n`` ways."""
+    first = total * i // n
+    return first, total * (i + 1) // n - first
+
+
+def thread():
+    """(thread, warp, lane, CTA) indices."""
+    tid, _, _ = cute.arch.thread_idx()
+    cta, _, _ = cute.arch.block_idx()
+    return tid, tid // 32, tid % 32, cta
+
+
+def shared_address(pointer):
+    return Uint32(pointer.toint())
+
+
+def at(tensor, *coord):
+    """Pointer to ``tensor[coord]``."""
+    return tensor.iterator + cute.crd2idx(coord, tensor.layout)
+
+
+def view(base, offset, dtype, shape, stride=None):
+    """``dtype`` tensor at byte ``offset`` of pointer ``base``."""
+    pointer = cute.recast_ptr(base + offset, dtype=dtype)
+    return cute.make_tensor(pointer, cute.make_layout(shape, stride=stride))
+
+
+def zeros(shape, dtype):
+    tensor = cute.make_rmem_tensor(shape, dtype)
+    tensor.fill(0)
+    return tensor
+
+
+def alloc(allocator, dtype, shape, align=16):
+    """Row-major ``shape`` tensor from a ``SmemAllocator``."""
+    stride = None
+    if isinstance(shape, tuple):
+        stride = tuple(math.prod(shape[i + 1 :]) for i in range(len(shape)))
+    layout = cute.make_layout(shape, stride=stride)
+    return allocator.allocate_tensor(dtype, layout, align)
+
+
+@asm()
+def load_words(
+    pointer,
+    *,
+    count=2,
+    shared=False,
+    evict_first=False,
+    mutable=False,
+    volatile=False,
+    scope=None,
+    cache_policy="cg",
+    word_type="b32",
+    result_type=Uint32,
+):
+    """``count`` words (a tuple); ``scope`` is a relaxed load at that scope.
+    ``mutable`` (a side-effecting load) for data written during the step (activations,
+    tables, refills). Only static weights use ``evict_first``."""
+    space = "global" + (f".{cache_policy}" if cache_policy else "")
+    if shared:
+        space = "shared"
+    if volatile:
+        space = "volatile.global"
+    if scope:
+        space = f"relaxed.{scope}.global"
+    policy = hint = tail = ""
+    if evict_first:
+        assert not shared
+        space = "global.nc.L1::no_allocate"
+        policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; "
+        hint, tail = ".L2::cache_hint", ", pol"
+    vector, registers = _vector(count)
+    address = f"[${count}]{tail}"
+    code = f"{{ {policy}ld.{space}{hint}{vector}.{word_type} {registers}, {address}; }}"
+    constraints = ["=f" if result_type is Float32 else "=r"] * count
+    constraints += ["r" if shared else "l"] + (["~{memory}"] if scope else [])
+    # Shared data is written by other threads between reads: never pure.
+    pure = not (shared or mutable or volatile or scope)
+    return code, ",".join(constraints), (pointer,), (result_type,) * count, pure
+
+
+smem_words = partial(load_words, count=4, shared=True, mutable=True)
+load_volatile = partial(load_words, count=1, volatile=True, word_type="u32")
+mailbox_load = partial(load_words, count=1, scope="gpu", word_type="u32")
+load_f32x4 = partial(load_words, count=4, result_type=Float32, word_type="f32")
+early_weights = partial(
+    load_words, cache_policy="nc.L1::no_allocate", word_type="u32", mutable=True
+)
+
+
+@asm()
+def load_bf16x8(pointer, *, shared=False, mutable=False):
+    opcode = "ld.shared.v4.b32" if shared else "ld.global.cg.v4.b32"
+    # Shifting BF16 bits into the high FP32 half preserves values exactly.
+    code = "{ .reg .b32 x<4>; " + opcode + " {x0,x1,x2,x3}, [$8]; "
+    for j in range(4):
+        code += f"shl.b32 ${2 * j}, x{j}, 16; and.b32 ${2 * j + 1}, x{j}, 0xffff0000; "
+    constraints = ",".join(["=f"] * 8 + ["r" if shared else "l"])
+    return code + "}", constraints, (pointer,), (Float32,) * 8, not (shared or mutable)
+
+
+@asm()
+def load_vector(pointer, *, cache="cg"):
+    """``cache`` "weight": a streaming L2 evict_first load of static weights; else
+    ``ld.global.<cache>``. Side-effecting in every mode."""
+    opcode, policy, tail = "ld.global." + cache, "", ""
+    if cache == "weight":
+        policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; "
+        opcode = "ld.global.L1::no_allocate.L2::cache_hint.L2::256B"
+        tail = ",pol"
+    if cache == "volatile":
+        opcode = "ld.volatile.global"
+    code = "{ " + policy + opcode + ".v4.b32 {$0,$1,$2,$3},[$4]" + tail + "; }"
+    clobber = ",~{memory}" if cache == "volatile" else ""
+    return code, "=r,=r,=r,=r,l" + clobber, (pointer,), (Uint32,) * 4
+
+
+@asm(Float32, pure=True)
+def table_load(pointer):
+    return "ld.global.nc.f32 $0,[$1];", "=f,l", (pointer,)
+
+
+@asm()
+def store_words(pointer, words, *, shared=False, scope=None):
+    """``scope`` is a relaxed global store at that scope."""
+    vector, registers = _vector(len(words), 1)
+    if shared:
+        # Other threads read it: a memory side effect, like the global stores.
+        space, constraints = (
+            f"shared{vector}.b32",
+            "r" + ",r" * len(words) + ",~{memory}",
+        )
+    else:
+        space = (
+            f"relaxed.{scope}.global{vector}.u32" if scope else f"global{vector}.b32"
+        )
+        constraints = "l" + ",r" * len(words) + ",~{memory}"
+        words = tuple(Uint32(v) for v in words)
+    return f"st.{space} [$0], {registers};", constraints, (pointer, *words)
+
+
+store_shared_words = partial(store_words, shared=True)
+store_relaxed = partial(store_words, scope="sys")
+mailbox_store = partial(store_words, scope="gpu")
+
+
+@asm()
+def store_f32x4(pointer, values, when=None):
+    """``when``: a predicate, so callers need no branch."""
+    code = "st.global.v4.f32 [$0], {$1, $2, $3, $4};"
+    if when is None:
+        return code, "l,f,f,f,f", (pointer, *values)
+    code = "{ .reg .pred p; setp.ne.b32 p,$5,0; @p " + code + " }"
+    return code, "l,f,f,f,f,r", (pointer, *values, Int32(when))
+
+
+@asm()
+def store_gate_mailbox(pointer, bits: Uint16):
+    return "st.relaxed.gpu.global.u16 [$0], $1;", "l,h,~{memory}", (pointer, bits)
+
+
+@asm()
+def store_bf16_pair(pointer, a, b):
+    code = "{ .reg .b32 pair; cvt.rn.bf16x2.f32 pair,$2,$1; "
+    return code + "st.relaxed.gpu.global.b32 [$0],pair; }", "l,f,f", (pointer, a, b)
+
+
+@asm()
+def store_pair(pointer, a, b):
+    """BF16x2 (a, b) with negative-zero halves made +0."""
+    code = "{ .reg .b32 pair,half; .reg .pred p; cvt.rn.bf16x2.f32 pair,$2,$1; "
+    code += "and.b32 half,pair,0xffff; setp.eq.b32 p,half,0x8000; "
+    code += "@p and.b32 pair,pair,0xffff0000; shr.u32 half,pair,16; "
+    code += "setp.eq.b32 p,half,0x8000; @p and.b32 pair,pair,0x0000ffff; "
+    return code + "st.global.b32 [$0],pair; }", "l,f,f", (pointer, a, b)
+
+
+@asm()
+def push_float(pointer, value, barrier, peer, when=None):
+    """DSMEM store of one FP32 ``value`` (or a 4-tuple) completing on ``barrier``
+    (``when``: predicated on it)."""
+    values = tuple(map(Float32, value if isinstance(value, tuple) else (value,)))
+    vector, registers = _vector(len(values), 1)
+    b, p, w = (f"${len(values) + i}" for i in (1, 2, 3))
+    code = f"{{ .reg .b32 dst,bar; mapa.shared::cluster.u32 dst,$0,{p}; "
+    code += f"mapa.shared::cluster.u32 bar,{b},{p}; "
+    guard = "" if when is None else f".reg .pred w; setp.ne.b32 w,{w},0; @w "
+    code += (
+        guard + f"st.async.shared::cluster.mbarrier::complete_tx::bytes{vector}.b32 "
+    )
+    code += f"[dst],{registers},[bar]; }}"
+    constraints = "r" + ",f" * len(values) + ",r,r" + ",r" * (when is not None)
+    operands = (pointer, *values, barrier, Int32(peer))
+    operands += () if when is None else (Int32(when),)
+    return code, constraints + ",~{memory}", operands
+
+
+@asm()
+def multicast_words(address, words, *, kind="f32"):
+    """Relaxed multimem store of FP32 ("f32") or integer ("b32", "bf16x2") words."""
+    vector, registers = _vector(len(words), 1)
+    f32 = kind == "f32"
+    values = (bits_float(v) if f32 else Uint32(v) for v in words)
+    constraints = "l" + (",f" if f32 else ",r") * len(words) + ",~{memory}"
+    code = f"multimem.st.relaxed.sys.global{vector}.{kind} [$0],{registers};"
+    return code, constraints, (Int64(address), *values)
+
+
+@asm()
+def publish_hint(pointer):
+    return "red.relaxed.gpu.global.add.u32 [$0],1;", "l,~{memory}", (pointer,)
+
+
+@asm()
+def copy_g2s(dst, src, size=16, barrier=None, *, evict_first=False):
+    """16-byte cp.async, or a static/dynamic ``size`` bulk copy on ``barrier``."""
+    operands = [dst, src]
+    policy = hint = tail = ""
+    if evict_first:
+        policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; "
+        hint, tail = ".L2::cache_hint", ", pol"
+    if barrier is None:
+        assert size == 16
+        code = f"cp.async.cg.shared.global{hint} [$0], [$1], 16{tail};"
+    else:
+        operands.append(barrier)
+        if not isinstance(size, int):
+            operands.append(size)
+            size = "$3"
+        code = (
+            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+            f"{hint} [$0], [$1], {size}, [$2]{tail};"
+        )
+    constraints = ",".join(["r", "l"] + ["r"] * (len(operands) - 2))
+    return "{ " + policy + code + " }", constraints, operands
+
+
+copy_evict = partial(copy_g2s, evict_first=True)
+
+
+@asm()
+def copy_bulk(dst, src, barrier, size):
+    code = "{ .reg .b64 p; createpolicy.fractional.L2::evict_first.b64 p,1.0; "
+    code += "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+    code += ".L2::cache_hint [$0],[$1],$3,[$2],p; }"
+    return code, "r,l,r,r,~{memory}", (dst, src, barrier, Int32(size))
+
+
+@asm()
+def copy_s2cluster(dst, src, barrier, peer, *, size):
+    code = "{ .reg .b32 dst,bar; mapa.shared::cluster.u32 dst,$0,$3; "
+    code += "mapa.shared::cluster.u32 bar,$2,$3; cp.async.bulk.shared::cluster"
+    code += f".shared::cta.mbarrier::complete_tx::bytes [dst],[$1],{size},[bar]; }}"
+    return code, "r,r,r,r", (dst, src, barrier, Int32(peer))
+
+
+@asm()
+def stage_weight(dst, src):
+    code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; "
+    code += "cp.async.cg.shared.global.L2::cache_hint.L2::256B [$0],[$1],16,pol; }"
+    return code, "r,l,~{memory}", (dst, src)
+
+
+@asm()
+def load_tma2(dst, descriptor, barrier, column, row):
+    code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; "
+    code += "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier"
+    code += "::complete_tx::bytes.L2::cache_hint [$0],[$1,{$3,$4}],[$2],pol; }"
+    return code, "r,l,r,r,r", (dst, descriptor, barrier, Int32(column), Int32(row))
+
+
+@asm()
+def load_tma3(dst, descriptor, barrier, row, chunk):
+    code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; "
+    code += "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier"
+    code += "::complete_tx::bytes.L2::cache_hint [$0],[$1,{0,$3,$4}],[$2],pol; }"
+    return code, "r,l,r,r,r", (dst, descriptor, barrier, Int32(row), Int32(chunk))
+
+
+@asm()
+def prefetch_bytes(src, size):
+    return "cp.async.bulk.prefetch.L2.global [$0], $1;", "l,r", (src, Int32(size))
+
+
+@dsl_user_op
+def prefetch_bulk(address, size, *, loc=None, ip=None):
+    """``prefetch_bytes`` of a 64-bit address through the NVVM op."""
+    pointer = cute.make_ptr(cutlass.Uint8, address, cute.AddressSpace.gmem)
+    nvvm.cp_async_bulk_prefetch(pointer, size, loc=loc, ip=ip)
+
+
+@asm()
+def prefetch_l2(pointer):
+    return "prefetch.global.L2 [$0];", "l", (pointer,)
+
+
+@asm()
+def prefetch_map(descriptor, row, chunk):
+    code = "cp.async.bulk.prefetch.tensor.3d.L2.global.tile [$0,{0,$1,$2}];"
+    return code, "l,r,r", (descriptor, Int32(row), Int32(chunk))
+
+
+def acquire_tensormap(pointer):
+    generic = cute.make_ptr(cutlass.Uint8, pointer.toint(), cute.AddressSpace.generic)
+    nvvm.fence_proxy_acquire("sys", generic, 128)
+    nvvm.prefetch_tensormap(generic)
+
+
+@asm()
+def acquire_tensormap_asm(pointer):
+    """``acquire_tensormap`` as one asm block on the unfolded address."""
+    code = "fence.proxy.tensormap::generic.acquire.sys [$0],128; "
+    return code + "prefetch.tensormap [$0];", "l", (pointer,)
+
+
+@asm(Int32)
+def try_wait(barrier, phase=0, *, test=False, scope="cluster"):
+    """Acquire parity probe at ``scope`` ("cluster", or "cta" when every arrival is
+    CTA-local); ``test`` is the non-blocking probe of phase 0."""
+    op, parity = ("test_wait", "0") if test else ("try_wait", "$2")
+    code = f"{{ .reg .pred p; mbarrier.{op}.parity.acquire.{scope}.shared::cta.b64 "
+    code += f"p,[$1],{parity}; selp.u32 $0,1,0,p; }}"
+    if test:
+        return code, "=r,r", (barrier,)
+    return code, "=r,r,r,~{memory}", (barrier, Int32(phase))
+
+
+@cute.jit
+def wait(barrier, phase=0, test: cutlass.Constexpr[bool] = False, scope="cluster"):
+    done = try_wait(barrier, phase, test=test, scope=scope)
+    while done == 0:
+        done = try_wait(barrier, phase, test=test, scope=scope)
+
+
+@asm(Int32)
+def try_wait_phase_zero(barrier):
+    code = "{ .reg .pred ready; mbarrier.try_wait.parity.acquire.cluster"
+    code += ".shared::cta.b64 ready,[$1],0; selp.u32 $0,1,0,ready; }"
+    return code, "=r,r,~{memory}", (barrier,)
+
+
+@cute.jit
+def wait_phase_zero_cluster(barrier: cute.Pointer):
+    done = try_wait_phase_zero(barrier)
+    while done == 0:
+        done = try_wait_phase_zero(barrier)
+
+
+def init_barriers(barriers, count, *, async_proxy=False):
+    for i in range(count):
+        cute.arch.mbarrier_init(barriers + i, 1)
+    cute.arch.mbarrier_init_fence()
+    if async_proxy:
+        cute.arch.fence_proxy("async.shared", space="cta")
+
+
+@asm()
+def named_barrier(barrier_id=1, count=128, *, arrive=False, all_groups=False):
+    if all_groups:
+        barrier_id, count = 2, 512
+    op = "arrive" if arrive else "sync"
+    return f"bar.{op} {barrier_id},{count};", "~{memory}", ()
+
+
+def group_barrier(group):
+    cute.arch.barrier(barrier_id=1 + 3 * group, number_of_threads=128)
+
+
+@asm()
+def sleep_ns(ns):
+    """``ns`` is an immediate if a Python int, else a register."""
+    if isinstance(ns, int):
+        return f"nanosleep.u32 {ns};", "", ()
+    return "nanosleep.u32 $0;", "r", (Uint32(ns),)
+
+
+@cute.jit
+def smem_desc(pointer: cute.Pointer, swizzle: cutlass.Constexpr[bool]):
+    address = shared_address(pointer)
+    value = Uint64((address & 262143) >> 4) | Uint64(1) << 46
+    if cutlass.const_expr(swizzle):
+        value |= Uint64(1) << 16 | Uint64(64) << 32 | Uint64(2) << 61
+    else:
+        value |= Uint64(8) << 32
+    return value
+
+
+def fence_tmem(*, before=False):
+    nvvm.tcgen05_fence("before_thread_sync" if before else "after_thread_sync")
+
+
+def wait_tmem(barrier, phase=0, scope="cluster"):
+    """``wait``, then order later tcgen05 operations after it."""
+    wait(barrier, phase, scope=scope)
+    fence_tmem()
+
+
+@asm()
+def free_tmem(address, columns):
+    code = f"tcgen05.dealloc.cta_group::1.sync.aligned.b32 $0,{columns};"
+    return code, "r", (Uint32(address),)
+
+
+@cute.jit
+def release_tmem(tmem, columns: cutlass.Constexpr[int], warp, owner):
+    """CTA teardown: warp ``owner`` frees the TMEM after every tcgen05 operation."""
+    fence_tmem(before=True)
+    cute.arch.sync_threads()
+    if warp == owner:
+        fence_tmem()
+        free_tmem(tmem, columns)
+
+
+@asm()
+def require_zero_tmem(address):
+    return "{ .reg .pred p; setp.ne.u32 p,$0,0; @p trap; }", "r", (Uint32(address),)
+
+
+@asm()
+def load_tmem(address, count=1):
+    """``count`` FP32 TMEM columns of this lane, waiting for the load."""
+    registers = ",".join(f"${i}" for i in range(count))
+    code = f"tcgen05.ld.sync.aligned.32x32b.x{count}.b32 {{{registers}}},[${count}]; "
+    code += "tcgen05.wait::ld.sync.aligned;"
+    constraints = ",".join(["=f"] * count + ["r", "~{memory}"])
+    result = Float32 if count == 1 else (Float32,) * count
+    return code, constraints, (Uint32(address),), result
+
+
+@asm()
+def copy_scale(column, descriptor):
+    code = "tcgen05.cp.cta_group::1.32x128b.warpx4 [$0],$1;"
+    return code, "r,l", (Uint32(column), Uint64(descriptor))
+
+
+commit = nvvm.tcgen05_commit
+
+
+@asm()
+def commit_asm(barrier):
+    """``commit`` as an asm block on the unfolded address."""
+    code = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [$0];"
+    return code, "r", (barrier,)
+
+
+@asm()
+def mma_bf16_block(d, a, b, instruction, accumulate, empty, release):
+    """Warp-uniform issue of one 64-K block (K-major 128B-swizzled A and B): the
+    elected lane runs the four K16 BF16 MMAs into ``d`` (the first accumulates iff
+    ``accumulate`` is nonzero) and, if ``release``, commits ``empty``."""
+    mma = "tcgen05.mma.cta_group::1.kind::f16"
+    commit = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64"
+    code = "{ .reg .pred e,p0,pt,pc; .reg .b64 a1,a2,a3,b1,b2,b3; "
+    code += "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 p0,$4,0; "
+    code += "setp.ne.b32 pc,$6,0; and.pred pc,pc,e; "
+    code += "add.s64 a1,$1,2; add.s64 a2,$1,4; add.s64 a3,$1,6; "
+    code += "add.s64 b1,$2,2; add.s64 b2,$2,4; add.s64 b3,$2,6; "
+    code += f"@e {mma} [$0],$1,$2,$3,p0; @e {mma} [$0],a1,b1,$3,pt; "
+    code += f"@e {mma} [$0],a2,b2,$3,pt; @e {mma} [$0],a3,b3,$3,pt; "
+    code += f"@pc {commit} [$5]; }}"
+    operands: tuple = (Uint32(d), Uint64(a), Uint64(b), Uint32(instruction))
+    operands += (Uint32(accumulate), empty, Uint32(release))
+    return code, "r,l,l,r,r,r,r", operands
+
+
+@asm()
+def mma_mx(d, a, b, instruction, sfa, sfb, accumulate):
+    """One block-scaled MXF8 tcgen05 MMA; ``accumulate`` is a runtime flag."""
+    words = (Uint32(v) for v in (instruction, sfa, sfb, accumulate))
+    code = "{ .reg .pred p; setp.ne.b32 p,$6,0; tcgen05.mma.cta_group::1"
+    code += ".kind::mxf8f6f4.block_scale [$0],$1,$2,$3,[$4],[$5],p; }"
+    return code, "r,l,l,r,r,r,r", (Uint32(d), Uint64(a), Uint64(b), *words)
+
+
+@cute.jit
+def mma_steps(d, a, b, sfa, sfb, accumulate, fc2: cutlass.Constexpr[bool]):
+    for k in cutlass.range_constexpr(6 if fc2 else 8):
+        a_offset: cutlass.Constexpr = k // 4 * 1024 + k % 4 * 2 + 4 * (fc2 and k >= 4)
+        instruction: cutlass.Constexpr = 142738048 | (k % 4 << 29) | (k % 4 << 4)
+        ops: tuple = (d, a + a_offset, b + k // 4 * 64 + k % 4 * 2, instruction)
+        ops += (sfa + k // 4 * 4, sfb + (k * 4 if fc2 else k // 4 * 4))
+        mma_mx(*ops, accumulate if k == 0 else 1)
+
+
+@asm()
+def chunk_mma(*args):
+    """(d, adesc, bdesc, sfa, sfa_src, sfb, sfb_src, copy_sfb, accum, empty_bar,
+    acc_bar): four K32 MMAs of one K128 chunk, with optional completion arrivals."""
+    constraints = "r,l,l,r,l,r,l,r,r,r,r"
+    types = [Uint32 if kind == "r" else Uint64 for kind in constraints.split(",")]
+    cp, mma = "tcgen05.cp.cta_group::1.32x128b.warpx4", "tcgen05.mma.cta_group::1"
+    commit = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64"
+    code = (
+        "{\n.reg .pred e,p0,pt,pc,pa,pe; .reg .b32 i<4>; .reg .b64 xa<4>,xb<4>;\n"
+        "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 pc,$7,0;\n"
+        "and.pred pc,pc,e; setp.ne.b32 pa,$10,0; and.pred pa,pa,e;\n"
+        "setp.ne.b32 pe,$9,0; and.pred pe,pe,e; setp.ne.b32 p0,$8,0;\n"
+        "mov.b32 i0,142738048; or.b32 i1,i0,536870928; or.b32 i2,i0,1073741856;\n"
+        f"or.b32 i3,i0,1610612784; @pc {cp} [$5],$6; @e {cp} [$3],$4;\n"
+        "add.s64 xa0,$1,0; add.s64 xb0,$2,0; add.s64 xa1,$1,2; add.s64 xb1,$2,2;\n"
+        "add.s64 xa2,$1,4; add.s64 xb2,$2,4; add.s64 xa3,$1,6; add.s64 xb3,$2,6;\n"
+    )
+    for k in range(4):
+        code += f"@e {mma}.kind::mxf8f6f4.block_scale [$0],\n"
+        code += f"xa{k},xb{k},i{k},[$3],[$5],{'pt' if k else 'p0'};\n"
+    code += f"@pe {commit} [$9]; @pa {commit} [$10];\n}}"
+    return code, constraints, tuple(t(v) for t, v in zip(types, args))
+
+
+def launch(kernel, args, grid, block, cluster=None, pdl=True, **options):
+    """Launch kernel(args[:-1]) on stream args[-1]; an int dimension means (n, 1, 1)."""
+    dim = lambda d: d if isinstance(d, tuple) else (d, 1, 1)
+    if cluster is not None:
+        options["cluster"] = dim(cluster)
+    kernel(args[:-1]).launch(
+        grid=dim(grid), block=dim(block), stream=args[-1], use_pdl=pdl, **options
+    )
+
+
+@cute.jit
+def scatter_reduce(values, lane, count: cutlass.Constexpr[int]):
+    """Warp-reduce ``count`` values; lane groups end up owning one value each."""
+    selected = Int32(0)
+    for stage in cutlass.range_constexpr(5):
+        offset = 16 >> stage
+        remaining = max(count >> stage, 1)
+        if cutlass.const_expr(remaining > 1):
+            half = remaining // 2
+            upper = (lane & offset) != 0
+            for i in cutlass.range_constexpr(half):
+                send = values[i] if upper else values[i + half]
+                keep = values[i + half] if upper else values[i]
+                values[i] = keep + cute.arch.shuffle_sync_bfly(send, offset)
+            if upper:
+                selected += half
+        else:
+            values[0] += cute.arch.shuffle_sync_bfly(values[0], offset)
+    return values[0], selected

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/twoshot.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Two-shot Lamport all-reduce of BF16 rows [B, width] over TP ranks.
+
+Rank ``o`` owns columns [o W, (o + 1) W), W = width / owners. Every buffer is
+Lamport-armed and only its local reader re-arms it:
+
+- ``push``: a rank stores its 8-column chunk of row ``token`` into the owner's
+  receive slot [rank] (a peer store; the owner's ``recv`` [TP, cap, W]; callers with a
+  runtime parity pass the whole [parities x TP, cap, W] and slot parity x TP + rank).
+- ``reduce``: the owner polls the TP slots of an owned chunk in its own memory, sums
+  them in FP32 in rank order, rounds once and multicasts the BF16 sum into every
+  rank's gather row [cap, width]; then re-arms the slots.
+- every rank then polls its gather row (one slot) and re-arms it.
+
+Each GPU moves ~2 x B x width x 2 bytes instead of TP x B x width x 2 (one-shot), and
+every rank gets the same bits. ``owners`` = 1 is a one-GPU stand-in: rank 0 owns every
+column and its TP local slots play the ranks. The o-projection reduces with ``reduce``
+and re-arms at once (two parities); the MoE tail reduces in WaveMoeTail and re-arms the
+previous MoE layer's rows (three parities).
+"""
+
+import cutlass
+import cutlass.cute as cute
+import torch
+import torch.distributed._symmetric_memory as symm_mem
+from cutlass import Float32, Int32, Int64, Uint32
+
+from .mailbox import rearm
+from .primitives import (
+    add_rn,
+    at,
+    load_vector,
+    multicast_words,
+    pack_pair,
+    published,
+    store_relaxed,
+    unpack_bf16,
+    zeros,
+)
+
+SENTINEL = -0x80000000
+
+
+class TwoShot:
+    """Symmetric buffers per layer parity: ``recv`` [TP, cap, W] (this rank's owned
+    columns, one slot per source rank), ``gather`` [1, cap, width] (multicast) and the
+    ranks' ``recv`` base addresses [parity, TP]."""
+
+    def __init__(self, group, device, tp, rank, width=7168, cap=64, parities=2):
+        self.tp, self.rank, self.width = tp, rank, width
+        shape = (parities, tp, cap, width // tp)
+        bf16 = dict(dtype=torch.bfloat16, device=device)
+        self.recv = symm_mem.empty(shape, **bf16)
+        self.gather = symm_mem.empty((parities, 1, cap, width), **bf16)
+        handles = [symm_mem.rendezvous(t, group) for t in (self.recv, self.gather)]
+        size = self.recv[0].numel() * 2
+        ptrs = handles[0].buffer_ptrs
+        bases = [[int(b) + p * size for b in ptrs] for p in range(parities)]
+        self.peers = torch.tensor(bases, dtype=torch.int64, device=device)
+        self.gather_mc = int(handles[1].multicast_ptr or 0)
+        if not self.gather_mc:  # views() offsets this address
+            raise RuntimeError("Kimi K3 two-shot all-reduce needs multicast memory")
+        for t in (self.recv, self.gather):
+            t.view(torch.int32).fill_(SENTINEL)
+        torch.accelerator.synchronize(device)
+        torch.distributed.barrier(group=group)
+
+    def views(self, parity):
+        """(recv, peers, gather, gather multicast address) of layer parity."""
+        gather_mc = self.gather_mc + parity * self.gather[0].numel() * 2
+        return self.recv[parity], self.peers[parity], self.gather[parity], gather_mc
+
+
+@cute.jit
+def push(recv, peers, slot, token, column, words):
+    """Store this rank's BF16 chunk (token, column .. column + 7) in receive slot
+    ``slot`` of its owner."""
+    width: cutlass.Constexpr = recv.shape[2]
+    owner = column // width
+    offset = cute.crd2idx((slot, token, column - owner * width), recv.layout)
+    store_relaxed(Int64(peers[owner]) + Int64(offset) * 2, words)  # polled remotely
+
+
+@cute.jit
+def reduce(recv, gather, gather_mc, rank, token, column):
+    """Owner of (token, column): the TP slots summed in FP32 in rank order, rounded
+    once (no negative zero, the sentinel), multicast to every gather row; slots
+    re-armed. Returns nothing; callers then poll ``gather`` like a one-slot mailbox."""
+    tp: cutlass.Constexpr = recv.shape[0]
+    local = column - rank * recv.shape[2]
+    raw = zeros((tp, 4), Uint32)
+    pending = Int32((1 << tp) - 1)
+    while pending != 0:
+        for slot in cutlass.range_constexpr(tp):
+            if (pending & (1 << slot)) != 0:
+                words = load_vector(at(recv, slot, token, local), cache="volatile")
+                for j in cutlass.range_constexpr(4):
+                    raw[slot, j] = words[j]
+                if published(words):
+                    pending &= ~(1 << slot)
+    values = cute.make_rmem_tensor(8, Float32)
+    for j in cutlass.range_constexpr(4):
+        low, high = unpack_bf16(raw[0, j])
+        for slot in cutlass.range_constexpr(1, tp):
+            a, b = unpack_bf16(raw[slot, j])
+            low, high = add_rn(low, a), add_rn(high, b)
+        values[2 * j], values[2 * j + 1] = low, high
+    words = tuple(
+        pack_pair(values[2 * j], values[2 * j + 1], policy="latent") for j in range(4)
+    )
+    offset = cute.crd2idx((0, token, column), gather.layout)
+    multicast_words(gather_mc + Int64(offset) * 2, words)
+    for slot in cutlass.range_constexpr(tp):
+        rearm(at(recv, slot, token, local))

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/projection.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Split-K projection on the FlashInfer CuTe core vendored in ``splitk.py`` (with an
+evict_first weight TMA; policies map padded rows to (M tile, N tile, split K,
+stages); attention owns live slot masks)."""
+
+from functools import cache
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
+
+from ..common.host import architecture, tensor_maps, to_cute
+from ..common.host import stream as current_stream
+from .splitk import SplitKGemm, SplitKTactic
+
+
+def matrix_view(tensor, *, output=False):
+    """Keep packed row pitches and FlashInfer's singleton-safe dynamic major axis."""
+    rows, cols = tensor.shape
+    shape = (cols, rows, 1) if output else (rows, cols, 1)
+    stride = (1, tensor.stride(0), 0) if output else (tensor.stride(0), 1, 0)
+    result = to_cute(tensor.as_strided(shape, stride))
+    return result.mark_layout_dynamic(leading_dim=0 if output else 1)
+
+
+@cache
+def compile_projection(tactic, output_dtype, arch):
+    """Code-only cache; weights, descriptors and activation pointers stay private."""
+
+    def fake(leading, dtype=cutlass.BFloat16):
+        shape = tuple(cute.sym_int() for _ in range(3))
+        stride = (1, cute.sym_int(64), 0) if leading == 0 else (cute.sym_int(64), 1, 0)
+        return make_fake_tensor(dtype, shape, stride, assumed_align=16)
+
+    y = fake(0, cutlass.Float32 if output_dtype == torch.float32 else cutlass.BFloat16)
+    split = SplitKTactic(*tactic)
+    kernel = SplitKGemm(tactic=split)
+    args = (fake(1), fake(1), y, y, cutlass.Int64(0), make_fake_stream())
+    return cute.experimental.compile(kernel, *args, options=f"--gpu-arch={arch}")
+
+
+def weight_map(weight, rows):
+    """Device TMA map of ``weight`` [N, K] (row pitch kept): box {64, rows}, SW128."""
+    types = cuda.CUtensorMapDataType
+    dtype = {
+        torch.bfloat16: types.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
+        torch.float16: types.CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
+    }[weight.dtype]
+    swizzle = cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B
+    l2 = cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B
+    n, k = weight.shape
+    pitch = weight.stride(0) * weight.element_size()
+    spec = (weight.data_ptr(), dtype, (k, n), (pitch,), (64, rows), swizzle, l2)
+    return tensor_maps([spec], weight.device)
+
+
+class PreparedProjection:
+    """Split-K projection with a per-batch tactic table and caller-owned outputs."""
+
+    def __init__(self, weight, policies, *, output_dtype=torch.bfloat16):
+        if weight.dim() != 2 or weight.dtype not in (torch.bfloat16, torch.float16):
+            raise ValueError(
+                f"projection weight must be a BF16/FP16 [N, K] matrix, got "
+                f"{tuple(weight.shape)} {weight.dtype}"
+            )
+        pitch = weight.stride(0) * weight.element_size()
+        if weight.stride(1) != 1 or pitch % 16 or weight.data_ptr() % 16:
+            raise ValueError(
+                "projection weight rows need unit stride and a 16 B TMA pitch"
+            )
+        if output_dtype not in (torch.bfloat16, torch.float32):
+            raise ValueError(f"unsupported projection output dtype {output_dtype}")
+        self.weight, self.device = weight, weight.device
+        self.output_dtype, self.policies = output_dtype, dict(policies)
+        self.weight_view = matrix_view(weight)
+        self.kernels, self.extra, self.maps = {}, {}, {}
+
+    def prepare(self):
+        arch, dtype = architecture(self.device), self.output_dtype
+        with torch.accelerator.device_index(self.device.index):
+            for batch, tactic in self.policies.items():
+                code = compile_projection(tactic, dtype, arch)
+                self.kernels[batch] = code
+                self.extra[batch] = self.weight_map(tactic[0])
+        return self
+
+    def weight_map(self, rows):
+        """The kernel's extra argument: A's map for ``rows``-row tiles."""
+        if rows not in self.maps:
+            self.maps[rows] = weight_map(self.weight, rows)
+        return (cutlass.Int64(self.maps[rows].data_ptr()),)
+
+    def run(self, x, out):
+        y = matrix_view(out, output=True)
+        b, stream = x.shape[0], current_stream(self.device)
+        if b not in self.kernels or x.shape[1] != self.weight.shape[1]:
+            raise ValueError(
+                f"projection input {tuple(x.shape)}: prepared batches "
+                f"{sorted(self.kernels)}, K {self.weight.shape[1]}"
+            )
+        code, extra = self.kernels[b], self.extra[b]
+        code(self.weight_view, matrix_view(x), y, y, *extra, stream)
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/splitk.py

@@ -0,0 +1,380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from FlashInfer v0.7.0.post1 (Apache-2.0):
+# flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py
+"""FlashInfer 0.7.0.post1 ``SplitKDenseGemmKernel`` (``dense_bf16_gemm_sm100_splitk``)
+reduced to our use (PDL, no bias) with the weight (kernel-A) TMA issued through a
+device tensor map with an L2 evict_first policy, and its synchronization fixed: the
+owner waits for the peers' ``st.async`` partials at ``.acquire.cluster``; the MMA warp
+waits for TMEM-pointer phase 0 before its release arrival (else its phase-1 wait can
+pass before the epilogue has read TMEM); ``tcgen05.fence::*_thread_sync`` around
+every mbarrier hand-off of tcgen05 work.
+
+Each cluster rank accumulates an exact K slice in FP32; peers publish partials to
+rank 0 through DSMEM, which reduces, casts and stores once. Kernel-M carries public N
+and kernel-N public M."""
+
+import cuda.bindings.driver as _cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass import Int32
+from cutlass._mlir.dialects import llvm
+from cutlass.cute import experimental as cute_ext
+from cutlass.cute.nvgpu import tcgen05
+from cutlass.cutlass_dsl import dsl_user_op
+from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import (
+    OWNER_RANK,
+    SplitKTactic,
+    _store_shared_remote_v4,
+)
+
+CTA_K, ATOM_K = 128, 64  # K per tile; K per 128-byte swizzle atom (one A box)
+
+__all__ = ["SplitKTactic", "SplitKGemm"]
+
+
+def _asm(code, constraints, operands, loc, ip):
+    llvm.inline_asm(
+        None, operands, code, constraints, has_side_effects=True,
+        is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, loc=loc, ip=ip,
+    )  # fmt: skip
+
+
+@dsl_user_op
+def _acquire_tensor_map(tensor_map, *, loc=None, ip=None):
+    code = "fence.proxy.tensormap::generic.acquire.sys [$0], 128; "
+    _asm(code + "prefetch.tensormap [$0];", "l", [tensor_map.ir_value()], loc, ip)
+
+
+@dsl_user_op
+def _tma_load_evict_first(smem, mbar, tensor_map, k0, row0, *, loc=None, ip=None):
+    """One {64, rows} box of A at (k0, row0), L2 evict_first."""
+    code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; "
+    code += "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx"
+    code += "::bytes.L2::cache_hint [$0], [$1, {$2, $3}], [$4], pol; }"
+    operands = [smem.toint().ir_value(), tensor_map.ir_value(), k0.ir_value()]
+    operands += [row0.ir_value(), mbar.toint().ir_value()]
+    _asm(code, "r,l,r,r,r", operands, loc, ip)
+
+
+@dsl_user_op
+def _wait_cluster(mbar, phase, *, loc=None, ip=None):
+    """``cute.arch.mbarrier_wait`` at ``.acquire.cluster`` (remote ``st.async``)."""
+    code = "{ .reg .pred P1; LAB_WAIT: "
+    code += "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [$0], $1, "
+    code += "10000000; @P1 bra.uni DONE; bra.uni LAB_WAIT; DONE: }"
+    _asm(code, "r,r", [mbar.toint().ir_value(), phase.ir_value()], loc, ip)
+
+
+@dsl_user_op
+def _fence_tcgen05(*, before=False, loc=None, ip=None):
+    side = "before" if before else "after"
+    _asm(f"tcgen05.fence::{side}_thread_sync;", "", [], loc, ip)
+
+
+def _smem(dtype, size, alignment, layout=None):
+    layout = cute.make_layout(size) if layout is None else layout
+    return cute_ext.allocate(dtype, cute.AddressSpace.smem, layout, alignment=alignment)
+
+
+class SplitKGemm:
+    """BF16 GEMM with a cluster-local split-K reduction; ``a_map`` is the device
+    address of A's tensor map (box {64, mma_m}, 128-byte swizzle)."""
+
+    def __init__(self, *, tactic: SplitKTactic):
+        self.acc_dtype = cutlass.Float32
+        self.cta_m, self.cta_n, self.cta_k = tactic.mma_m, tactic.mma_n, CTA_K
+        self.num_ab_stage, self.split_k = tactic.ab_stages, tactic.split_k
+        self.threads_per_cta, self.epilog_threads = 256, 128
+        self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n)
+        self.cta_group = tcgen05.CtaGroup.ONE
+        self.cluster_shape = (1, tactic.split_k, 1)
+        values_per_thread = tactic.mma_m * tactic.mma_n // self.epilog_threads
+        assert values_per_thread % 4 == 0, "remote stores move 4 values"
+        self.mailbox_elements = (
+            (tactic.split_k - 1) * self.epilog_threads * values_per_thread
+        )
+        self.expected_transaction_bytes = self.mailbox_elements * 4
+
+    @cute.experimental.jit
+    def __call__(
+        self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, bias: cute.Tensor,
+        a_map: cutlass.Int64, stream: _cuda.CUstream,
+    ):  # fmt: skip
+        # Grid-y packs output-N tile and cluster rank.
+        self.kernel(a, b, c, bias, a_map).launch(
+            grid=(
+                cute.ceil_div(c.layout.shape[0], self.cta_m),
+                cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k,
+                c.layout.shape[2],
+            ),
+            block=(self.threads_per_cta, 1, 1),
+            cluster=self.cluster_shape,
+            smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")),
+            stream=stream,
+            use_pdl=True,
+        )
+
+    @cute.experimental.kernel
+    def kernel(
+        self, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor, mBias: cute.Tensor,
+        a_map: cutlass.Int64,
+    ):  # fmt: skip
+        """A (M, K, L) and B (N, K, L) K-major, C (M, N, L) M-major; bias unused."""
+        stages = self.num_ab_stage
+        ab_dtype = mA.element_type
+        majors = (utils.LayoutEnum.from_tensor(t).mma_major_mode() for t in (mA, mB))
+        tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            ab_dtype, ab_dtype, *majors, self.acc_dtype, self.cta_group,
+            self.mma_tiler_mn,
+        )  # fmt: skip
+
+        mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k)
+        block_idx = cute.arch.block_idx()
+        bidx = block_idx[0]
+        split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        n_idx = block_idx[1] // self.split_k
+        l_idx = block_idx[2]
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+
+        shape = (tiled_mma, mnk_tiler, ab_dtype, stages)
+        sA = _smem(ab_dtype, None, 1024, sm100_utils.make_smem_layout_a(*shape))
+        sB = _smem(ab_dtype, None, 1024, sm100_utils.make_smem_layout_b(*shape))
+        acc_layout = cute_ext.make_tmem_layout_acc(
+            tiled_mma, self.mma_tiler_mn, acc_stage=1
+        )
+        bar_full = _smem(cutlass.Int64, stages, 8).iterator
+        bar_empty = _smem(cutlass.Int64, stages, 8).iterator
+        bar_tma_epilog = _smem(cutlass.Int64, 1, 8).iterator  # FlashInfer's bias
+        bar_mma_epilog = _smem(cutlass.Int64, 1, 8).iterator
+        bar_tmem_alloc = _smem(cutlass.Int64, 1, 8).iterator
+        tmem_base_ptr = _smem(cutlass.Int32, 1, 4).iterator
+        mailbox, bar_reduce = sA, bar_mma_epilog  # split 1: compile-time dead
+        if cutlass.const_expr(self.split_k > 1):
+            mailbox = _smem(cutlass.Float32, self.mailbox_elements, 128)
+            bar_reduce = _smem(cutlass.Int64, 1, 8).iterator
+
+        if warp_idx == 0:
+            with cute.arch.elect_one():
+                for i in range(stages):
+                    cute.arch.mbarrier_init(bar_full + i, 2)
+                    cute.arch.mbarrier_init(bar_empty + i, 1)
+                cute.arch.mbarrier_init(bar_tma_epilog, 32)
+                cute.arch.mbarrier_init(bar_mma_epilog, 1)
+                cute.arch.mbarrier_init(bar_tmem_alloc, 160)
+                if cutlass.const_expr(self.split_k > 1):
+                    # Owner arrival plus peer transaction-byte credits.
+                    cute.arch.mbarrier_init(bar_reduce, 1)
+
+        cute.arch.mbarrier_init_fence()
+        if cutlass.const_expr(self.split_k > 1):
+            cute.arch.cluster_arrive_relaxed()  # peers' barriers before their stores
+        else:
+            cute.arch.barrier()
+
+        # K splits evenly over the cluster (tactics are validated for the shapes).
+        k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k
+        k_tile_start = split_rank * k_tile_count
+
+        if cutlass.const_expr(self.split_k > 1):
+            cute.arch.cluster_wait()
+
+        # Warp 3 is idle; warps 4-7 run the epilogue.
+        tiles = (k_tile_start, k_tile_count)
+        if warp_idx == 0:
+            self.weight_warp(bar_full, bar_empty, sA, a_map, bidx * self.cta_m, *tiles)
+        elif warp_idx == 1:
+            b_tile = cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx))
+            b_map = cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B")
+            self.activation_warp(bar_full, bar_empty, b_tile, sB, b_map, *tiles)
+        elif warp_idx == 2:
+            k_blocks = self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2])
+            self.mma_warp(
+                bar_full, bar_empty, bar_mma_epilog, bar_tmem_alloc, tiled_mma, sA, sB,
+                tmem_base_ptr, acc_layout, k_blocks, k_tile_count,
+            )  # fmt: skip
+        elif warp_idx >= 4:
+            tile = cute.local_tile(mC, (self.cta_m, self.cta_n), (bidx, n_idx, l_idx))
+            self.epilog_warp(
+                bar_mma_epilog, bar_tmem_alloc, tmem_base_ptr, acc_layout, tile,
+                cute.arch.thread_idx()[0] - 128, mC.element_type,
+                utils.LayoutEnum.from_tensor(mC), mailbox, bar_reduce, split_rank,
+            )  # fmt: skip
+
+    @cute.experimental.jit
+    def weight_warp(
+        self, bar_full, bar_empty, s_tile: cute.Tensor, a_map: cutlass.Int64,
+        row0: cutlass.Int32, k_tile_start: cutlass.Int32, k_tile_count: cutlass.Int32,
+    ):  # fmt: skip
+        """A: two evict_first boxes per K tile, before the PDL wait (weights)."""
+        stages = self.num_ab_stage
+        stage_bytes = cute.size_in_bytes(
+            s_tile.element_type, cute.slice_(s_tile.layout, (None, None, None, 0))
+        )
+        _acquire_tensor_map(a_map)
+        empty_phase = cutlass.Int32(1)
+        for k_tile in cutlass.range(k_tile_count, unroll=1):
+            stage = k_tile % stages
+            cute.arch.mbarrier_wait(bar_empty + stage, empty_phase)
+            with cute.arch.elect_one():
+                cute.arch.mbarrier_arrive_and_expect_tx(bar_full + stage, stage_bytes)
+            with cute.arch.elect_one():
+                k0 = (k_tile_start + k_tile) * self.cta_k
+                tile = s_tile[None, None, None, stage].iterator
+                for atom in cutlass.range_constexpr(self.cta_k // ATOM_K):
+                    box, mbar = tile + atom * self.cta_m * ATOM_K, bar_full + stage
+                    k = k0 + atom * ATOM_K
+                    _tma_load_evict_first(box, mbar, a_map, k, row0)
+            if stage == stages - 1:
+                empty_phase = empty_phase ^ 1
+
+        cute.arch.griddepcontrol_launch_dependents()
+        self.drain(bar_empty, empty_phase, k_tile_count)
+
+    @cute.experimental.jit
+    def activation_warp(
+        self, bar_full, bar_empty, g_tile: cute.Tensor, s_tile: cute.Tensor,
+        cta_v_map: cute.Layout, k_tile_start: cutlass.Int32,
+        k_tile_count: cutlass.Int32,
+    ):  # fmt: skip
+        """B (the activations): after the PDL wait."""
+        stages = self.num_ab_stage
+        stage_bytes = cute.size_in_bytes(
+            s_tile.element_type, cute.slice_(s_tile.layout, (None, None, None, 0))
+        )
+        cute.arch.griddepcontrol_wait()
+        empty_phase = cutlass.Int32(1)
+        for k_tile in cutlass.range(k_tile_count, unroll=1):
+            stage = k_tile % stages
+            cute.arch.mbarrier_wait(bar_empty + stage, empty_phase)
+            with cute.arch.elect_one():
+                cute.arch.mbarrier_arrive_and_expect_tx(bar_full + stage, stage_bytes)
+            cute_ext.tma_load(
+                g_tile[None, None, k_tile_start + k_tile],
+                s_tile[None, None, None, stage],
+                (bar_full + stage).value,
+                cta_v_map=cta_v_map,
+                tma_operation_type=cute_ext.OperationTypeEnum.SM90_TMA_LOAD,
+                update_expect_tx=False,
+            )
+            if stage == stages - 1:
+                empty_phase = empty_phase ^ 1
+        self.drain(bar_empty, empty_phase, k_tile_count)
+
+    @cute.experimental.jit
+    def drain(self, bar_empty, empty_phase: cutlass.Int32, k_tile_count: cutlass.Int32):
+        stages = self.num_ab_stage
+        for tail in cutlass.range(stages, unroll=1):
+            stage = (tail + k_tile_count) % stages
+            cute.arch.mbarrier_wait(bar_empty + stage, empty_phase)
+            if stage == stages - 1:
+                empty_phase = empty_phase ^ 1
+
+    @cute.experimental.jit
+    def mma_warp(
+        self, bar_full, bar_empty, bar_mma_epilog, bar_tmem_alloc,
+        tiled_mma: cute.TiledMma, sA: cute.Tensor, sB: cute.Tensor, tmem_base_ptr,
+        acc_layout: cutlass.Constexpr, mma_inst_tile_k: cutlass.Constexpr,
+        k_tile_count: cutlass.Int32,
+    ):  # fmt: skip
+        num_tmem_cols = 256
+        cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False)
+        _fence_tcgen05(before=True)
+        cute.arch.mbarrier_arrive(bar_tmem_alloc)
+        cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False)
+
+        tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr)
+        accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0]
+        mma_atom = cute.make_mma_atom(tiled_mma.op)
+        full_phase = cutlass.Int32(0)
+        for k_tile in cutlass.range(k_tile_count, unroll=1):
+            stage = k_tile % self.num_ab_stage
+            cute.arch.mbarrier_wait(bar_full + stage, full_phase)
+            _fence_tcgen05()
+            for k_block in range(mma_inst_tile_k):
+                accumulate = k_tile != 0 if k_block == 0 else True
+                mma_atom.set(tcgen05.Field.ACCUMULATE, accumulate)
+                cute_ext.dot(
+                    mma_atom,
+                    cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3),
+                    cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3),
+                    accumulator,
+                )
+            with cute.arch.elect_one():
+                tcgen05.commit(bar_empty + stage, None, self.cta_group)
+            if stage == self.num_ab_stage - 1:
+                full_phase = full_phase ^ 1
+
+        with cute.arch.elect_one():
+            tcgen05.commit(bar_mma_epilog, None, self.cta_group)
+        cute.arch.mbarrier_wait(bar_tmem_alloc, 0)  # the release lands in phase 1
+        cute.arch.mbarrier_arrive(bar_tmem_alloc)
+        cute.arch.mbarrier_wait(bar_tmem_alloc, 1)
+        _fence_tcgen05()
+        cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False)
+
+    @cute.experimental.jit
+    def epilog_warp(
+        self, bar_mma_epilog, bar_tmem_alloc, tmem_base_ptr,
+        acc_layout: cutlass.Constexpr, gD_tile: cute.Tensor, epi_tid: cutlass.Int32,
+        c_dtype: cutlass.Constexpr, d_layout: cutlass.Constexpr, mailbox, bar_reduce,
+        split_rank: cutlass.Int32,
+    ):  # fmt: skip
+        # Wait until MMA publishes the TMEM base pointer.
+        cute.arch.mbarrier_arrive(bar_tmem_alloc)
+        cute.arch.mbarrier_wait(bar_tmem_alloc, 0)
+        _fence_tcgen05()
+
+        acc_view = cute.make_tensor(
+            cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr),
+            acc_layout,
+        )[((None, None), 0, 0, 0)]
+        epi_tile = (self.cta_m, self.cta_n)
+        load_op = sm100_utils.get_tmem_load_op(
+            (self.cta_m, self.cta_n, self.cta_k), d_layout, c_dtype, self.acc_dtype,
+            epi_tile, False,
+        )  # fmt: skip
+        tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(load_op, acc_view)
+        gD_epi = cute.flat_divide(gD_tile, epi_tile)
+
+        # Match each epilogue thread's TMEM partition in RMEM.
+        rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid)
+        rmem = cute.AddressSpace.rmem
+        rAcc = cute_ext.allocate(self.acc_dtype, rmem, rmem_layout, alignment=32)
+        rD = cute_ext.allocate(c_dtype, rmem, rmem_layout, alignment=32)
+        thr_t2r = tiled_copy_t2r.get_slice(epi_tid)
+
+        cute.arch.mbarrier_wait(bar_mma_epilog, 0)
+        _fence_tcgen05()
+        cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc)
+        # Make tcgen05.ld visible before TMEM release and RMEM use.
+        cute.arch.fence_view_async_tmem_load()
+        _fence_tcgen05(before=True)
+        cute.arch.mbarrier_arrive(bar_tmem_alloc)
+
+        # Peers publish FP32 partials; only rank 0 reduces and stores.
+        if cutlass.const_expr(self.split_k > 1):
+            vpt = cutlass.const_expr(cute.size(rmem_layout))  # values per thread
+            vpp = cutlass.const_expr(self.epilog_threads * vpt)  # values per peer
+            if split_rank != OWNER_RANK:
+                for i in cutlass.range_constexpr(0, vpt, 4):
+                    words = (rAcc[i], rAcc[i + 1], rAcc[i + 2], rAcc[i + 3])
+                    slot = mailbox.iterator + (split_rank - Int32(1)) * vpp
+                    slot = slot + epi_tid * vpt + i
+                    _store_shared_remote_v4(*words, slot, bar_reduce, Int32(OWNER_RANK))
+            else:
+                if epi_tid == 0:
+                    expected = self.expected_transaction_bytes
+                    cute.arch.mbarrier_arrive_and_expect_tx(bar_reduce, expected)
+                _wait_cluster(bar_reduce, Int32(0))
+                for peer in cutlass.range_constexpr(self.split_k - 1):
+                    for i in cutlass.range_constexpr(vpt):
+                        rAcc[i] = rAcc[i] + mailbox[peer * vpp + epi_tid * vpt + i]
+
+        if split_rank == OWNER_RANK:
+            rD.store(rAcc.load().to(c_dtype))
+            # Preserve TMEM coordinates; the copy predicates output tails.
+            cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0])

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/ws_gemm.py

@@ -0,0 +1,293 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Weight-stationary skinny GEMM for the MoE front: y[M, N] = x[M, K] @ W[N, K]^T.
+
+Each CTA owns ``rt`` 16-row N tiles; its ``rt * kq`` warps split K into ``kq`` ranges.
+Weights enter registers and cp.async SMEM stages before the PDL wait. The epilogue
+publishes the first ``valid`` of ``split`` rows (router shard) as FP32 words at BF16
+column 3584 + 2 * (``slot`` + row) of a multicast [1, capacity, width] mailbox and the
+rest (latent down shard) as sanitized BF16 at column ``base``.
+
+``tg`` 0 is the single pass: warps load x from global and the cluster leader sums the
+K partials over DSMEM. ``tg`` > 0 streams tokens in groups of ``tg``: warp 0 bulk-copies
+the CTA's K slice of x once (one mbarrier per group) for all row tiles, and after each
+group's MMAs every warp pushes its 8-token x 16-row partial blocks to their owner CTAs
+(block % ks) with st.async; owners sum the ks * kq parts in K order and publish one
+group behind, so only the last group's blocks follow the MMA loop.
+"""
+
+from typing import Any
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Int64, Uint32, Uint64
+
+from ..common.primitives import (
+    alloc,
+    asm,
+    at,
+    bf16x2,
+    bits_float,
+    copy_g2s,
+    float_bits,
+    launch,
+    load_words,
+    mma_into,
+    multicast_words,
+    push_float,
+    smem_words,
+    wait,
+)
+
+
+@asm(Uint32, pure=True)
+def positive_zero(low, high):
+    """BF16 (low, high) with -0.0 halves made +0.0, so it never equals the sentinel."""
+    code = "{ .reg .b32 t; .reg .pred p,q; and.b32 t,$1,0xFFFF; "
+    code += "setp.eq.u32 p,t,0x8000; shr.u32 t,$1,16; setp.eq.u32 q,t,0x8000; "
+    code += "mov.b32 $0,$1; @p and.b32 $0,$0,0xFFFF0000; @q and.b32 $0,$0,0x0000FFFF; }"
+    return code, "=r,r", (Uint32(bf16x2(low, high)),)
+
+
+@asm(Float32)
+def peer_load(pointer, peer):
+    """FP32 at the same shared-memory offset in cluster CTA ``peer``."""
+    code = "{ .reg .b32 a; mapa.shared::cluster.u32 a,$1,$2; "
+    return code + "ld.shared::cluster.f32 $0,[a]; }", "=f,r,r", (pointer, Int32(peer))
+
+
+@asm(Uint32, pure=True)
+def positive_zero32(value):
+    """FP32 bits of ``value`` with -0.0 made +0.0, so it never equals the sentinel."""
+    code = "{ .reg .pred p; setp.eq.u32 p,$1,0x80000000; selp.b32 $0,0,$1,p; }"
+    return code, "=r,r", (Uint32(float_bits(value)),)
+
+
+class WsGemm:
+    def __init__(
+        self, n, k, m, rt, kq, staged, split, valid, ks=1, tg=0, ch=0
+    ):  # fmt: skip
+        self.n, self.k, self.m, self.rt, self.kq = n, k, m, rt, kq
+        # ks: a cluster of ks CTAs splits K, so each CTA preloads 1/ks of the tile's
+        # weights; ch: accumulator chains per 8-token tile (0: 4 / tiles per group).
+        self.split, self.valid, self.ks = split, valid, ks
+        self.staged, self.kb = staged, k // (32 * kq * ks)
+        self.rb, self.nt = self.kb - staged, (m + 7) // 8
+        self.grouped, self.ng = tg > 0, min(self.nt, tg // 8) if tg else self.nt
+        if self.grouped and ks <= 1:  # token groups push partials to K-split peers
+            raise ValueError(f"WsGemm token groups (tg={tg}) need a K-split cluster")
+        self.groups, self.chains = -(-self.nt // self.ng), ch or max(1, 4 // self.ng)
+        self.grid, self.threads = n // (16 * rt) * ks, 32 * rt * kq
+        self.blocks, self.parts = self.nt * rt, ks * kq
+        self.slots = -(-self.blocks // ks)  # owned blocks s * ks + kpart
+        # Bulk x rows: the CTA's K slice plus 64 B, so token pairs hit disjoint banks.
+        self.slice = k // ks
+        self.pitch = self.slice // 2 + 16
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, self.grid, self.threads, self.ks, min_blocks_per_mp=1)
+
+    @cute.kernel
+    def kernel(self, args):
+        x, w = args[0:2]
+        smem = cutlass.utils.SmemAllocator()
+        block, _, _ = cute.arch.block_idx()
+        cta, kpart = block // self.ks, block % self.ks
+        tid, _, _ = cute.arch.thread_idx()
+        lane, warp = tid % 32, tid // 32
+        group, part = lane // 4, lane % 4
+        tile, quarter = warp // self.kq, warp % self.kq
+        row = (cta * self.rt + tile) * 16 + group
+        kofs = (kpart * self.kq + quarter) * self.kb * 32 + part * 8
+        shape = (self.nt, self.rt, self.kq, 8, 17)
+        red = smem.allocate_tensor(Float32, cute.make_layout(shape), 16)
+        stage = alloc(smem, Uint32, (self.rt * self.kq, max(self.staged, 1), 2, 32, 4))
+        xs = recv = None
+        bars: Any = None
+        owned: Any = None
+        if cutlass.const_expr(self.grouped):
+            xs = alloc(smem, Uint32, (self.nt * 8, self.pitch))
+            bars = alloc(smem, Uint64, (self.groups,), 8)
+            recv = alloc(smem, Uint32, (self.slots, self.parts, 32, 4))
+            owned = alloc(smem, Uint64, (self.slots,), 8)
+            if tid == 0:
+                for g in cutlass.range_constexpr(self.groups):
+                    cute.arch.mbarrier_init(bars.iterator + g, 1)
+                for s in cutlass.range_constexpr(self.slots):
+                    if s * self.ks + kpart < self.blocks:
+                        cute.arch.mbarrier_init(owned.iterator + s, 1)
+                        nbytes = self.parts * 512
+                        cute.arch.mbarrier_arrive_and_expect_tx(
+                            owned.iterator + s, nbytes
+                        )
+                cute.arch.mbarrier_init_fence()
+            cute.arch.sync_threads()
+            cute.arch.cluster_arrive_relaxed()
+        wp = w.iterator + row * self.k + kofs
+        for j in cutlass.range_constexpr(self.staged):
+            for h in cutlass.range_constexpr(2):
+                dst = at(stage, warp, j, h, lane, 0)
+                copy_g2s(dst, wp + h * 8 * self.k + (self.rb + j) * 32)
+        cute.arch.cp_async_commit_group()
+        wr = {}
+        for j in cutlass.range_constexpr(self.rb):
+            for h in cutlass.range_constexpr(2):
+                src = wp + h * 8 * self.k + j * 32
+                wr[j, h] = load_words(src, count=4, evict_first=True)
+        cute.arch.griddepcontrol_wait()
+        cute.arch.griddepcontrol_launch_dependents()
+        if cutlass.const_expr(self.grouped):
+            if warp == 0:
+                self.stage_x(x, xs, bars, kpart, lane)
+            cute.arch.cp_async_wait_group(0)
+        else:
+            cute.arch.cp_async_wait_group(0)
+            cute.arch.sync_threads()
+        for g in cutlass.range_constexpr(self.groups):
+            acc = cute.make_rmem_tensor((self.ng, self.chains, 4), Float32)
+            acc.fill(0)
+            if cutlass.const_expr(self.grouped):
+                wait(bars.iterator + g)
+            for j in cutlass.range_constexpr(self.kb):
+                if cutlass.const_expr(j < self.rb):
+                    wa, wb = wr[j, 0], wr[j, 1]
+                else:
+                    sa = at(stage, warp, j - self.rb, 0, lane, 0)
+                    wa, wb = smem_words(sa), smem_words(sa + 128)
+                for t in cutlass.range_constexpr(self.tiles(g)):
+                    token = (g * self.ng + t) * 8 + group
+                    xv = (Uint32(0), Uint32(0), Uint32(0), Uint32(0))
+                    if cutlass.const_expr(self.grouped):
+                        col = (quarter * self.kb + j) * 16 + part * 4
+                        xv = smem_words(at(xs, token, col))
+                    else:
+                        xg = x.iterator + token * x.stride[0] + kofs + j * 32
+                        if (g * self.ng + t) * 8 + 7 < self.m or token < self.m:
+                            xv = load_words(xg, count=4, mutable=True)
+                    for h in cutlass.range_constexpr(2):
+                        c = (j * 2 + h) % self.chains
+                        a = (wa[2 * h], wb[2 * h], wa[2 * h + 1], wb[2 * h + 1])
+                        b = (xv[2 * h], xv[2 * h + 1])
+                        mma_into(acc[t, c, None], a, b, bf16=True)
+            if cutlass.const_expr(self.grouped and g == 0):
+                cute.arch.cluster_wait()  # owners' barriers are armed
+            for t in cutlass.range_constexpr(self.tiles(g)):
+                values = []
+                for v in cutlass.range_constexpr(4):
+                    value = acc[t, 0, v]
+                    for c in cutlass.range_constexpr(1, self.chains):
+                        value = value + acc[t, c, v]
+                    tt, col = g * self.ng + t, group + (v // 2) * 8
+                    if cutlass.const_expr(self.grouped):
+                        values.append(value)
+                    else:
+                        red[tt, tile, quarter, 2 * part + v % 2, col] = value
+                if cutlass.const_expr(self.grouped):
+                    blk = (g * self.ng + t) * self.rt + tile
+                    s, peer = blk // self.ks, blk % self.ks
+                    dst = at(recv, s, kpart * self.kq + quarter, lane, 0)
+                    push_float(dst, tuple(values), owned.iterator + s, peer)
+            if cutlass.const_expr(self.grouped and g > 0):
+                self.settle(args, red, recv, owned, cta, kpart, warp, lane, g - 1)
+        if cutlass.const_expr(self.grouped):
+            self.settle(args, red, recv, owned, cta, kpart, warp, lane, self.groups - 1)
+        else:
+            cute.arch.sync_threads()
+            if cutlass.const_expr(self.ks > 1):
+                cute.arch.cluster_arrive()
+                cute.arch.cluster_wait()
+            self.gather(args, red, cta, kpart, tid, lane)
+            if cutlass.const_expr(self.ks > 1):
+                cute.arch.cluster_arrive()
+                cute.arch.cluster_wait()
+
+    def tiles(self, g):
+        """8-token tiles in token group ``g``."""
+        return min(self.ng, self.nt - g * self.ng)
+
+    @cute.jit
+    def stage_x(self, x, xs, bars, kpart, lane):
+        """Warp 0 bulk-copies every token's K slice; one barrier per token group."""
+        if lane == 0:
+            for g in cutlass.range_constexpr(self.groups):
+                tokens = min(self.m, (g + 1) * self.ng * 8) - g * self.ng * 8
+                nbytes = tokens * self.slice * 2
+                cute.arch.mbarrier_arrive_and_expect_tx(bars.iterator + g, nbytes)
+        cute.arch.sync_warp()
+        for i in cutlass.range_constexpr(-(-self.m // 32)):
+            token = i * 32 + lane
+            if i * 32 + 31 < self.m or token < self.m:
+                src = x.iterator + token * x.stride[0] + kpart * self.slice
+                bar = bars.iterator + token // (8 * self.ng)
+                copy_g2s(at(xs, token, 0), src, self.slice * 2, bar)
+
+    @cute.jit
+    def settle(
+        self, args, red, recv, owned, cta, kpart, warp, lane, g: cutlass.Constexpr
+    ):
+        """Owned blocks of token group ``g``: a warp per block waits for every K part,
+        sums them in K order, transposes through ``red`` and publishes."""
+        for t in cutlass.range_constexpr(self.tiles(g)):
+            for tile in cutlass.range_constexpr(self.rt):
+                blk = (g * self.ng + t) * self.rt + tile
+                s = blk // self.ks
+                if (kpart == blk % self.ks) & (warp == s % (self.rt * self.kq)):
+                    wait(owned.iterator + s)
+                    total = cute.make_rmem_tensor(4, Float32)
+                    for p in cutlass.range_constexpr(self.parts):
+                        words = smem_words(at(recv, s, p, lane, 0))
+                        for v in cutlass.range_constexpr(4):
+                            value = bits_float(words[v])
+                            total[v] = value if p == 0 else total[v] + value
+                    tt = g * self.ng + t
+                    for v in cutlass.range_constexpr(4):
+                        col = lane // 4 + v // 2 * 8
+                        red[tt, tile, 0, 2 * (lane % 4) + v % 2, col] = total[v]
+                    cute.arch.sync_warp()
+                    token = tt * 8 + lane // 2
+                    if (lane < 16) & (token < self.m):
+                        half = lane % 2 * 8
+                        values = [
+                            red[tt, tile, 0, lane // 2, half + j] for j in range(8)
+                        ]
+                        col = (cta * self.rt + tile) * 16 + half
+                        self.epilogue(args, token, col, values)
+
+    @cute.jit
+    def gather(self, args, red, cta, kpart, tid, lane):
+        """Each thread sums one column's K parts in order (the leader also its peers');
+        each eight-lane group then gathers the sums for its first lane's 16 B store."""
+        total = self.m * self.rt * 16
+        for start in range(0, total, self.threads):
+            item = start + tid
+            v8, half = item % 8, (item // 8) % 2
+            rtile, token = (item // 16) % self.rt, item // (16 * self.rt)
+            mine = Float32(0)
+            if item < total:
+                mine = red[token // 8, rtile, 0, token % 8, half * 8 + v8]
+                for q in cutlass.range_constexpr(1, self.kq):
+                    mine = mine + red[token // 8, rtile, q, token % 8, half * 8 + v8]
+                if kpart == 0:
+                    for peer in cutlass.range_constexpr(1, self.ks):
+                        for q in cutlass.range_constexpr(self.kq):
+                            idx = (token // 8, rtile, q, token % 8, half * 8 + v8)
+                            mine = mine + peer_load(at(red, *idx), peer)
+            values = [cute.arch.shuffle_sync(mine, (lane & 24) + v) for v in range(8)]
+            if (item < total) & (v8 == 0) & (kpart == 0):
+                col = (cta * self.rt + rtile) * 16 + half * 8
+                self.epilogue(args, token, col, values)
+
+    @cute.jit
+    def epilogue(self, args, token, col, values):
+        out, mc, slot, base = args[2:6]
+        if col < self.valid:
+            for v in cutlass.range_constexpr(2):
+                words = tuple(positive_zero32(values[4 * v + i]) for i in range(4))
+                offset = token * out.shape[2] + 3584 + 2 * (slot + col + 4 * v)
+                multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2")
+        elif col >= self.split:
+            column = base + col - self.split
+            words = tuple(positive_zero(*values[i : i + 2]) for i in range(0, 8, 2))
+            offset = token * out.shape[2] + column
+            multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2")

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/kernel.py

@@ -0,0 +1,715 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cluster-split KDA recurrence; an optional fifth warp projects F_B and multicasts the
+decays while the main warps run the convolution and query/key normalization."""
+
+from functools import partial
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32
+from cutlass.cute.arch import store_async_dsmem
+
+from ..common.primitives import (
+    alloc,
+    asm,
+    at,
+    copy_g2s,
+    dot,
+    float_bits,
+    fma,
+    launch,
+    load_f32x4,
+    load_words,
+    mma_into,
+    push_float,
+    sigmoid,
+    smem_words,
+    store_f32x4,
+    unpack_bf16,
+    wait,
+    warp_sum,
+    zeros,
+)
+
+load_bf16x4 = partial(load_words, count=2, cache_policy="nc", word_type="u32")
+load_weight_bf16x8 = partial(load_words, count=4, cache_policy="nc", word_type="u32")
+load_activation_bf16x4 = partial(load_words, count=2, word_type="u32", mutable=True)
+load_activation_bf16x8 = partial(load_words, count=4, word_type="u32", mutable=True)
+load_staged = partial(
+    load_words, count=4, shared=True, mutable=True, result_type=Float32, word_type="f32"
+)
+
+
+@asm((Float32, Int32), pure=True)
+def reciprocal_guess(denominator):
+    """``sigmoid_reciprocal``'s refined fast path (IEEE mode) and whether it holds."""
+    code = "{ .reg .f32 r,error,neg_error; .reg .b32 bits; .reg .pred regular; "
+    code += "mov.b32 bits,$2; add.u32 bits,bits,0x01800000; "
+    code += "and.b32 bits,bits,0x7f800000; setp.gt.u32 regular,bits,0x01ffffff; "
+    code += "rcp.approx.ftz.f32 r,$2; fma.rn.f32 error,$2,r,0fbf800000; "
+    code += "neg.f32 neg_error,error; fma.rn.f32 $0,r,neg_error,r; "
+    return code + "selp.u32 $1,1,0,regular; }", "=f,=r,f", (denominator,)
+
+
+@asm(Float32, pure=True)
+def reciprocal_rn(value):
+    return "div.rn.f32 $0,0f3f800000,$1;", "=f,f", (value,)
+
+
+@cute.jit
+def sigmoids(values):
+    """``sigmoid(v, False)`` of every value, bit for bit, with all fast paths ahead
+    of the division fallbacks' branches (verify tokens get ILP)."""
+    count: cutlass.Constexpr = len(values)
+    den = [Float32(1) + cute.math.exp(-Float32(v), fastmath=True) for v in values]
+    result, regular = zeros(count, Float32), zeros(count, Int32)
+    irregular = Int32(0)
+    for i in cutlass.range_constexpr(count):
+        result[i], regular[i] = reciprocal_guess(den[i])
+        irregular |= regular[i] ^ 1
+    if irregular != 0:  # rare: one branch for the whole batch
+        for i in cutlass.range_constexpr(count):
+            if regular[i] == 0:
+                result[i] = reciprocal_rn(den[i])
+    return result
+
+
+def advance(base, n):
+    """``base + n``, or ``base`` itself for a static zero ``n`` (keeps q = 1 code)."""
+    return base + n if n else base
+
+
+@asm()
+def store_evict_f32x4(pointer, values):
+    code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; "
+    code += "st.global.L2::cache_hint.v4.f32 [$0],{$1,$2,$3,$4},pol; }"
+    return code, "l,f,f,f,f", (pointer, *values)
+
+
+class KdaSplit:
+    """128 main threads and C value partitions. Conv state [slot, 3 * heads * 128, 3]
+    and recurrent state [slot, head, value, key] keep their strides; slots <= 0 are
+    padding (never written) and active rows own distinct slots. Weights must be ready
+    at launch; each run waits for prior cache restoration and slot publication.
+
+    ``schedule`` = (early, lanes per state row, minimum CTAs per SM, CTAs per head);
+    zeros keep the batch defaults. Two CTAs per head give each CTA 64 state rows and
+    each gate lane two decays. Early issues every load the predecessor does not
+    produce (slot, state and conv rows, conv / F_B / norm weights) before the PDL
+    wait, converts operands at first use, and syncs the cluster only before each
+    role's first DSMEM send. Safe because those loads read only weights and data a
+    graph or non-PDL boundary orders before this launch (states and conv rows of this
+    layer's previous call or a fresh prompt's non-PDL fill, the host's slot table),
+    never same-step producer output, which is read after the wait; PDL-chained calls
+    without such a boundary must not share state slots. Early 2 also bulk-copies
+    the CTA's (contiguous) state rows to shared memory and updates four passes at a
+    time, which frees the registers that held them across the input phase; early 3
+    (two CTAs per head) also stages the gate warp's F_B rows.
+
+    ``tokens`` q > 1 (speculative verify; ``batch`` counts requests, rows b * q + t)
+    runs each request's q tokens in order with vLLM's spec cache semantics: slots
+    [B, >= q] and accepted [B] (an extra last tensor, read early like the slots); the
+    state comes from slot[b, accepted - 1] and token t's state goes to slot[b, t] if
+    > 0; conv slot[b, 0] holds [dim, q + 2] with taps at accepted - 1 ..., rewritten to
+    [c_accepted, c_accepted + 1, x_0, ..., x_q-1]. Token t's F_B row is MMA column t;
+    CTA t % C normalizes token t."""
+
+    def __init__(self, heads, batch, gate_mode, has_bias, schedule=(0, 0, 0), tokens=1):
+        # The gate stage sends decays for tokens 2 * quarter + column (at most 8), while
+        # the decay barrier expects 512 bytes per token: more tokens would never arrive.
+        if not 1 <= tokens <= 8:
+            raise ValueError(f"KdaSplit supports 1-8 tokens per request; got {tokens}")
+        self.heads, self.gate_mode, self.has_bias = heads, gate_mode, has_bias
+        early, lanes, self.occupancy, cluster, evict = (*schedule, 0, 0)[:5]
+        self.early, self.staged = bool(early), early >= 2
+        self.gate_staged, self.tokens = early == 3, tokens
+        assert gate_mode != 2 or not (early or cluster)
+        assert tokens == 1 or not (self.staged or gate_mode == 2 or cluster == 16)
+        # TP8 (12 heads) uses fast math; B1 releases dependents at entry.
+        self.fast_math, self.trigger_at_entry = heads == 12, batch == 1
+        self.cluster = cluster or (8 if batch == 1 else 4)
+        self.lanes_per_row = lanes or (8 if batch <= (4 if heads == 6 else 2) else 16)
+        self.rows, self.groups = 128 // self.cluster, 128 // self.lanes_per_row
+        self.passes = (self.rows + self.groups - 1) // self.groups
+        self.wave = min(4, self.passes) if self.staged else self.passes
+        self.chunks = 32 // self.lanes_per_row
+        assert not self.gate_staged or self.rows >= 64
+        self.evict = bool(evict)  # staged state traffic marked L2 evict-first
+        assert self.staged or not self.evict
+
+    @cute.jit
+    def __call__(self, *args):
+        grid = (self.cluster, self.heads, args[0].shape[0] // self.tokens)
+        bound = dict(min_blocks_per_mp=self.occupancy) if self.occupancy else {}
+        launch(self.kernel, args, grid, 160, self.cluster, **bound)
+
+    @cute.kernel
+    def kernel(self, args):
+        x, weights, bias, conv_state, raw_gate, beta, a_log, dt_bias, slots = args[:9]
+        state, output, gate, norm_weight, lower_bound, scale, eps, f_b_weight = args[
+            9:17
+        ]
+        T: cutlass.Constexpr = self.tokens
+        if cutlass.const_expr(not self.early):
+            cute.arch.griddepcontrol_wait()
+        thread, _, _ = cute.arch.thread_idx()
+        partition, head, token = cute.arch.block_idx()
+        if cutlass.const_expr(T == 1):
+            slot = slots[token]
+            conv_slot, offset, writes = slot, 0, [slot]
+        else:
+            request, token = token, token * T
+            offset = args[17][request] - 1
+            slot, conv_slot = slots[request, offset], slots[request, 0]
+            writes = [slots[request, t] for t in range(T)]
+        if slot <= 0:
+            cute.arch.griddepcontrol_wait()
+            if (partition == 0) & (thread < 128):
+                for t in cutlass.range_constexpr(T):
+                    output[advance(token, t), head * 128 + thread] = BFloat16(0)
+            cute.arch.griddepcontrol_launch_dependents()
+        else:
+            lane, warp = thread % 32, thread // 32
+            group = thread // self.lanes_per_row
+            group_lane = thread % self.lanes_per_row
+            dim: cutlass.Constexpr = self.heads * 128
+            value_base, qk_channel = partition * self.rows, head * 128 + thread
+            value_channel = head * 128 + value_base + thread
+
+            smem = cutlass.utils.SmemAllocator()
+            shared_q, shared_k, shared_decay, shared_v, shared_o, reductions = (
+                alloc(smem, Float32, size * T)
+                for size in (128, 128, 128, self.rows, 128, 8)
+            )
+            barrier, gate_barrier = (smem.allocate(cutlass.Uint64, 8) for _ in range(2))
+            owned: cutlass.Constexpr = -(-T // self.cluster)  # tokens a CTA normalizes
+            if cutlass.const_expr(owned > 1):  # one output barrier per owned token
+                barrier = smem.allocate(8 * owned, 8)
+                barrier = cute.recast_ptr(barrier, dtype=cutlass.Uint64)
+            if cutlass.const_expr(self.staged):
+                staged = alloc(smem, Float32, (self.rows, 128), align=128)
+                state_barrier = smem.allocate(cutlass.Uint64, 8)
+            gate_rows = None
+            if cutlass.const_expr(self.gate_staged):  # [copy, lane] 16-byte slots
+                gate_rows = alloc(smem, Uint32, (self.rows // 2, 32, 4), align=128)
+            if cutlass.const_expr(self.trigger_at_entry):
+                cute.arch.griddepcontrol_launch_dependents()
+            if cutlass.const_expr(not self.early):
+                self.setup(barrier, gate_barrier, partition, thread)
+                cute.arch.cluster_wait()
+            if thread >= 128:
+                if cutlass.const_expr(self.early):
+                    cute.arch.cluster_arrive_relaxed()
+                coords = (head, token, value_base, lane)
+                self.gate_warp(args, shared_decay, gate_barrier, coords, gate_rows)
+
+            if thread < 128:
+                if cutlass.const_expr(self.early):  # thread 0 arms; barrier 1 orders it
+                    self.setup(barrier, gate_barrier, partition, thread)
+                # 64-bit page bases (slot x page stride passes Int32), 32-bit offsets.
+                page = lambda s: state[Int64(s), None, None, None]
+                state_page = page(slot)
+                conv_page = conv_state[Int64(conv_slot), None, None]
+                state_row = lambda r, c, s=state_page: at(s, head, value_base + r, c)
+                history = zeros((self.passes, self.chunks, 4), Float32)
+                if cutlass.const_expr(self.staged):  # noqa: SIM102
+                    if thread == 0:
+                        cute.arch.mbarrier_init(state_barrier, 1)
+                        cute.arch.mbarrier_init_fence()
+                        size = self.rows * 512
+                        cute.arch.mbarrier_arrive_and_expect_tx(state_barrier, size)
+                        src, policy = state_row(0, 0), dict(evict_first=self.evict)
+                        copy_g2s(staged.iterator, src, size, state_barrier, **policy)
+                for p in cutlass.range_constexpr(0 if self.staged else self.passes):
+                    row = p * self.groups + group
+                    if row < self.rows:
+                        for j in cutlass.range_constexpr(self.chunks):
+                            column = 4 * (group_lane + self.lanes_per_row * j)
+                            values = load_f32x4(state_row(row, column))
+                            for e in cutlass.range_constexpr(4):
+                                history[p, j, e] = values[e]
+                taps = zeros((3, 3), BFloat16)
+                coeff = zeros((3, 4), weights.element_type if self.early else Float32)
+                accum = zeros(3 * T, Float32)
+                for s in cutlass.range_constexpr(3):
+                    channel = value_channel if s == 2 else qk_channel
+                    if (s != 2) | (thread < self.rows):
+                        for t in cutlass.range_constexpr(3):
+                            taps[s, t] = conv_page[s * dim + channel, offset + t]
+                        for t in cutlass.range_constexpr(4):
+                            coeff[s, t] = weights[s, t, channel]
+                        if cutlass.const_expr(self.has_bias):
+                            accum[s] = bias[s * dim + channel]
+                            for t in cutlass.range_constexpr(1, T):
+                                accum[3 * t + s] = accum[s]
+                nw = norm_weight[thread]
+
+                cute.arch.griddepcontrol_wait()
+                inputs = cute.make_rmem_tensor(3 * T, BFloat16)
+                betas, gates = [], []
+                for t in cutlass.range_constexpr(T):
+                    row = advance(token, t)
+                    inputs[3 * t], inputs[3 * t + 1] = (
+                        x[row, qk_channel],
+                        x[row, dim + qk_channel],
+                    )
+                    inputs[3 * t + 2] = BFloat16(0)
+                    if cutlass.const_expr(t == T - 1):  # noqa: SIM102 (all tokens)
+                        if thread < self.rows:
+                            for u in cutlass.range_constexpr(T):
+                                channel = 2 * dim + value_channel
+                                inputs[3 * u + 2] = x[advance(token, u), channel]
+                    if cutlass.const_expr(self.early):  # converted at first use
+                        betas.append(beta[row, head])
+                        gates.append(gate[row, qk_channel])
+                    else:
+                        betas.append(Float32(beta[row, head]))
+                        gates.append(Float32(gate[row, qk_channel]))
+                cute.arch.griddepcontrol_launch_dependents()
+
+                vs, squares, qks = [], [], []
+                for t in cutlass.range_constexpr(T):  # window: 3 taps, then the inputs
+                    for s in cutlass.range_constexpr(3):
+                        for i in cutlass.range_constexpr(4):
+                            w = t + i
+                            term = taps[s, w] if w < 3 else inputs[3 * (w - 3) + s]
+                            term = Float32(term)
+                            cf = Float32(coeff[s, i]) if self.early else coeff[s, i]
+                            acc = accum[3 * t + s]
+                            accum[3 * t + s] = fma(term, cf, acc, ftz=self.fast_math)
+                act = lambda i, value: sigmoid(value, self.fast_math)
+                if cutlass.const_expr(T > 1 and not self.fast_math):
+                    raw = [accum[i] for i in range(3 * T)] + betas + gates
+                    act = partial(lambda batch, i, value: batch[i], sigmoids(raw))
+                for t in cutlass.range_constexpr(T):
+                    q = accum[3 * t] * act(3 * t, accum[3 * t])
+                    k = accum[3 * t + 1] * act(3 * t + 1, accum[3 * t + 1])
+                    v = accum[3 * t + 2] * act(3 * t + 2, accum[3 * t + 2])
+                    betas[t] = act(3 * T + t, Float32(betas[t]))
+                    gates[t] = act(4 * T + t, Float32(gates[t]))
+                    shared_q[advance(thread, 128 * t)] = q
+                    shared_k[advance(thread, 128 * t)] = k
+                    vs.append(v)
+                    qks += [q, k]
+                    if cutlass.const_expr(t == T - 1):  # noqa: SIM102
+                        if thread < self.rows:
+                            for u in cutlass.range_constexpr(T):
+                                shared_v[advance(thread, self.rows * u)] = vs[u]
+                    squares += [warp_sum(q * q), warp_sum(k * k)]
+                    if cutlass.const_expr(t == T - 1):  # noqa: SIM102
+                        if lane == 0:
+                            for u in cutlass.range_constexpr(T):
+                                reductions[advance(warp, 8 * u)] = squares[2 * u]
+                                reductions[advance(warp + 4, 8 * u)] = squares[
+                                    2 * u + 1
+                                ]
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                def factors(t):
+                    r = lambda i: reductions[8 * t + i]
+                    qn = (r(0) + r(2)) + (r(1) + r(3))
+                    kn = (r(4) + r(6)) + (r(5) + r(7))
+                    fq = cute.math.rsqrt(qn + 1.0e-6, fastmath=True) * scale
+                    return fq, cute.math.rsqrt(kn + 1.0e-6, fastmath=True)
+
+                if cutlass.const_expr(T > 1):  # normalize each token's q, k once
+                    for t in cutlass.range_constexpr(T):
+                        fq, fk = factors(t)
+                        shared_q[128 * t + thread] = qks[2 * t] * fq
+                        shared_k[128 * t + thread] = qks[2 * t + 1] * fk
+                    cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                def operands(t):
+                    fq, fk = factors(t) if T == 1 else (None, None)
+                    qr = cute.make_rmem_tensor((self.chunks, 4), Float32)
+                    kr = cute.make_rmem_tensor((self.chunks, 4), Float32)
+                    for j in cutlass.range_constexpr(self.chunks):
+                        column = 4 * (group_lane + self.lanes_per_row * j)
+                        q4 = load_f32x4(
+                            advance(shared_q.iterator, 128 * t) + column, shared=True
+                        )
+                        k4 = load_f32x4(
+                            advance(shared_k.iterator, 128 * t) + column, shared=True
+                        )
+                        for e in cutlass.range_constexpr(4):
+                            qr[j, e] = q4[e] * fq if T == 1 else q4[e]
+                            kr[j, e] = k4[e] * fk if T == 1 else k4[e]
+                    if cutlass.const_expr(t == 0):
+                        wait(gate_barrier)
+                    decays = []
+                    for j in cutlass.range_constexpr(self.chunks):
+                        column = 4 * (group_lane + self.lanes_per_row * j)
+                        decay_ptr = advance(shared_decay.iterator, 128 * t) + column
+                        decays.append(load_f32x4(decay_ptr, shared=True))
+                    values = shared_v
+                    if cutlass.const_expr(T > 1):
+                        pointer = shared_v.iterator + self.rows * t
+                        values = cute.make_tensor(pointer, cute.make_layout(self.rows))
+                    return kr, qr, decays, values, betas[t]
+
+                ready = operands(0)
+                for t in cutlass.range_constexpr(T):
+                    current, ready = ready, operands(t + 1) if t + 1 < T else None
+                    # Butterfly order: interleave rows per level, publish before state.
+                    dots = cute.make_rmem_tensor(self.passes, Float32)
+                    if cutlass.const_expr(self.staged):
+                        wait(state_barrier)
+                    source = staged if self.staged else None
+                    owner, bar = t % self.cluster, advance(barrier, t // self.cluster)
+                    for start in cutlass.range_constexpr(0, self.passes, self.wave):
+                        stop = min(start + self.wave, self.passes)
+                        lanes = (group, group_lane)
+                        self.update(
+                            history, dots, current, lanes, (start, stop), source
+                        )
+                        if cutlass.const_expr(self.early and start == 0 and t == 0):
+                            cute.arch.cluster_wait()
+                        for p in cutlass.range_constexpr(start, stop):
+                            row = p * self.groups + group
+                            send = (row < self.rows) & (group_lane == 0)
+                            if cutlass.const_expr(T > 1):  # predicated: no branches
+                                position = self.output_position(value_base + row)
+                                dst = advance(shared_o.iterator, 128 * t) + position
+                                push_float(dst, dots[p], bar, owner, when=send)
+                            elif send:
+                                position = self.output_position(value_base + row)
+                                dst = shared_o.iterator + position
+                                store_async_dsmem(dst, float_bits(dots[p]), bar, owner)
+                        rows = partial(state_row, s=page(writes[t]))
+                        when = None if T == 1 else writes[t] > 0
+                        self.store_rows(history, rows, (start, stop), lanes, when)
+                if thread < self.rows:
+                    channel = 2 * dim + value_channel
+                    window = (
+                        taps[2, 1],
+                        taps[2, 2],
+                        *(inputs[3 * t + 2] for t in range(T)),
+                    )
+                    for i, value in enumerate(window):
+                        conv_page[channel, i] = value
+
+                for t in cutlass.range_constexpr(T):
+                    if partition == t % self.cluster:
+                        wait(advance(barrier, t // self.cluster))
+                        raw_o = shared_o[advance(self.output_position(thread), 128 * t)]
+                        sums = cute.make_rmem_tensor(4, Float32)
+                        for g in cutlass.range_constexpr(4):
+                            base = lane % 8 * 16 + g * 4
+                            o_row = advance(shared_o.iterator, 128 * t)
+                            o4 = load_f32x4(o_row + base, shared=True)
+                            s0, s1, s2, s3 = (o * o for o in o4)
+                            sums[g] = warp_sum((s0 + s2) + (s1 + s3), 8)
+                        sumsq = (sums[0] + sums[2]) + (sums[1] + sums[3])
+                        rstd = cute.math.rsqrt(sumsq * (1.0 / 128) + eps, fastmath=True)
+                        result = raw_o * rstd * nw * gates[t]
+                        output[advance(token, t), qk_channel] = BFloat16(result)
+                        if cutlass.const_expr(t == 0):
+                            for section in cutlass.range_constexpr(2):
+                                channel = section * dim + qk_channel
+                                window = (taps[section, 1], taps[section, 2])
+                                window += tuple(
+                                    inputs[3 * u + section] for u in range(T)
+                                )
+                                for i, value in enumerate(window):
+                                    conv_page[channel, i] = value
+
+    @cute.jit
+    def store_rows(self, history, state_row, span, lanes, when=None):
+        """Store the passes in ``span`` of the CTA's state rows (``when``: predicated
+        on it, without branches)."""
+        (start, stop), (group, group_lane) = span, lanes
+        for p in cutlass.range_constexpr(start, stop):
+            row = p * self.groups + group
+            if cutlass.const_expr(when is not None):
+                for j in cutlass.range_constexpr(self.chunks):
+                    column = 4 * (group_lane + self.lanes_per_row * j)
+                    values = [history[p, j, e] for e in range(4)]
+                    pointer = state_row(row % self.rows, column)
+                    store_f32x4(pointer, values, when=when & (row < self.rows))
+            elif row < self.rows:
+                for j in cutlass.range_constexpr(self.chunks):
+                    column = 4 * (group_lane + self.lanes_per_row * j)
+                    pointer = state_row(row, column)
+                    row_values = [history[p, j, e] for e in range(4)]
+                    if cutlass.const_expr(self.evict):
+                        store_evict_f32x4(pointer, row_values)
+                    else:
+                        store_f32x4(pointer, row_values)
+
+    @cute.jit
+    def update(self, history, dots, operands, lanes, span, staged):
+        """Decay, delta-rule update and query dot of the passes in ``span``, whose rows
+        ``staged`` (if given) first supplies from shared memory."""
+        kr, qr, decays, shared_v, beta = operands
+        (group, lane), (start, stop) = lanes, span
+        if cutlass.const_expr(staged is not None):
+            for p in cutlass.range_constexpr(start, stop):
+                row = (p * self.groups + group) % self.rows
+                for j in cutlass.range_constexpr(self.chunks):
+                    column = 4 * (lane + self.lanes_per_row * j)
+                    values = load_staged(staged.iterator + (row * 128 + column))
+                    for e in cutlass.range_constexpr(4):
+                        history[p, j, e] = values[e]
+        for p in cutlass.range_constexpr(start, stop):
+            for j in cutlass.range_constexpr(self.chunks):
+                for e in cutlass.range_constexpr(4):
+                    history[p, j, e] *= decays[j][e]
+            dots[p] = self.row_dot(history, kr, p)
+        for shift in reversed(range(self.lanes_per_row.bit_length() - 1)):
+            for p in cutlass.range_constexpr(start, stop):
+                dots[p] += cute.arch.shuffle_sync_bfly(dots[p], 1 << shift)
+        for p in cutlass.range_constexpr(start, stop):
+            safe_row = (p * self.groups + group) % self.rows
+            v_new = (shared_v[safe_row] - dots[p]) * beta
+            for j in cutlass.range_constexpr(self.chunks):
+                for e in cutlass.range_constexpr(4):
+                    history[p, j, e] = fma(
+                        kr[j, e], v_new, history[p, j, e], ftz=self.fast_math
+                    )
+            dots[p] = self.row_dot(history, qr, p)
+        for shift in reversed(range(self.lanes_per_row.bit_length() - 1)):
+            for p in cutlass.range_constexpr(start, stop):
+                dots[p] += cute.arch.shuffle_sync_bfly(dots[p], 1 << shift)
+
+    @cute.jit
+    def gate_warp(self, args, shared_decay, barrier, coords, staged_rows):
+        """Project one value partition and multicast its decays to every CTA. Modes fix
+        the F_B reduction order: 0 = one chain of eight MMAs, 1 = four chains of two,
+        2 = sixteen 8-FMA partials summed in a butterfly. Projection math stays IEEE.
+        ``staged_rows``: shared [copy, lane, 4] words; each lane cp.asyncs its own
+        16-byte F_B operands there before the PDL wait and reads them back after it."""
+        f_a, a_log, dt_bias, lower_bound, weight = (args[i] for i in (4, 6, 7, 13, 16))
+        head, token, value_base, lane = coords
+        rows, mode = self.rows, self.gate_mode
+        per: cutlass.Constexpr = max(1, rows // 32)  # decays per lane
+        a = a_log[head]
+        if cutlass.const_expr(mode == 2):
+            feature = lane % rows
+            row = head * 128 + value_base + feature
+            # uint4 weights before the wait, two uint2 activation loads per 8-FMA
+            # partial after it; keep that granularity and the low-then-high FMA order.
+            wr = [load_weight_bf16x8(at(weight, row, 8 * i)) for i in range(16)]
+            dt = dt_bias[row]
+            cute.arch.griddepcontrol_wait()
+            partial = cute.make_rmem_tensor(16, Float32)
+            for i in cutlass.range_constexpr(16):
+                a0 = load_activation_bf16x4(at(f_a, token, 8 * i))
+                a1 = load_activation_bf16x4(at(f_a, token, 8 * i + 4))
+                av = (*a0, *a1)
+                acc = Float32(0)
+                for e in cutlass.range_constexpr(4):
+                    x0, x1 = unpack_bf16(av[e])
+                    w0, w1 = unpack_bf16(wr[i][e])
+                    acc = fma(x0, w0, acc)
+                    acc = fma(x1, w1, acc)
+                partial[i] = acc
+            for shift in cutlass.range_constexpr(3, -1, -1):
+                for i in cutlass.range_constexpr(1 << shift):
+                    partial[i] += partial[i | (1 << shift)]
+            raw_g = [Float32(BFloat16(partial[0]))]
+            dt = [dt]
+        else:
+            tiles: cutlass.Constexpr = rows // 16
+            chains: cutlass.Constexpr = 4 if mode == 1 else 1
+            group, quarter = lane >> 2, lane & 3
+            feature = 16 * (quarter >> 1) + group + 8 * (quarter & 1)
+            dt = [Float32(0) for _ in range(per)]
+            if cutlass.const_expr(self.tokens > 1):  # token column, its features
+                token += group % self.tokens
+                features = [
+                    16 * (i // 2) + group + 8 * (i % 2) for i in range(rows // 8)
+                ]
+                dt = [dt_bias[head * 128 + value_base + f] for f in features]
+            for r in cutlass.range_constexpr(per if self.tokens == 1 else 0):
+                if feature + 32 * r < rows:
+                    dt[r] = dt_bias[head * 128 + value_base + 32 * r + feature]
+            # Two decays per lane load 16 bytes per row: chunk c then covers columns
+            # 32 (c // 2) + 8 quarter + 4 (c % 2) + [0, 4) instead of 16 c + 4 quarter.
+            wide: cutlass.Constexpr = per >= 2
+            wr = cute.make_rmem_tensor((tiles, 8, 4), Uint32)
+            if cutlass.const_expr(staged_rows is not None):
+                for tile in cutlass.range_constexpr(tiles):
+                    row = head * 128 + value_base + 16 * tile + group
+                    for chunk in cutlass.range_constexpr(4):
+                        col = 32 * chunk + 8 * quarter
+                        for half in cutlass.range_constexpr(2):
+                            slot = staged_rows[
+                                (tile * 4 + chunk) * 2 + half, lane, None
+                            ]
+                            source = at(weight, row + 8 * half, col)
+                            copy_g2s(slot.iterator, source)
+                cute.arch.cp_async_commit_group()
+            else:
+                for tile in cutlass.range_constexpr(tiles):
+                    row = head * 128 + value_base + 16 * tile + group
+                    for chunk in cutlass.range_constexpr(4 if wide else 8):
+                        if cutlass.const_expr(wide):
+                            col = 32 * chunk + 8 * quarter
+                            low = load_weight_bf16x8(at(weight, row, col))
+                            high = load_weight_bf16x8(at(weight, row + 8, col))
+                            self.split_pairs(wr, tile, chunk, low, high)
+                        else:
+                            col = 16 * chunk + 4 * quarter
+                            low = load_bf16x4(at(weight, row, col))
+                            high = load_bf16x4(at(weight, row + 8, col))
+                            wr[tile, chunk, 0], wr[tile, chunk, 2] = low
+                            wr[tile, chunk, 1], wr[tile, chunk, 3] = high
+            cute.arch.griddepcontrol_wait()
+            if cutlass.const_expr(wide):
+                pairs = [
+                    load_activation_bf16x8(at(f_a, token, 32 * c + 8 * quarter))
+                    for c in range(4)
+                ]
+                ar = [pairs[c // 2][2 * (c % 2) : 2 * (c % 2) + 2] for c in range(8)]
+            else:
+                ar = [
+                    load_activation_bf16x4(at(f_a, token, 16 * c + 4 * quarter))
+                    for c in range(8)
+                ]
+            if cutlass.const_expr(staged_rows is not None):
+                cute.arch.cp_async_wait_group(0)
+                for tile in cutlass.range_constexpr(tiles):
+                    for chunk in cutlass.range_constexpr(4):
+                        copy = (tile * 4 + chunk) * 2
+                        low = smem_words(staged_rows[copy, lane, None].iterator)
+                        high = smem_words(staged_rows[copy + 1, lane, None].iterator)
+                        self.split_pairs(wr, tile, chunk, low, high)
+            acc = zeros((chains, tiles, 4), Float32)
+            for chunk in cutlass.range_constexpr(8):
+                chain: cutlass.Constexpr = chunk // 2 if mode == 1 else 0
+                for tile in cutlass.range_constexpr(tiles):
+                    wa = tuple(wr[tile, chunk, e] for e in range(4))
+                    mma_into(acc[chain, tile, None], wa, ar[chunk], bf16=True)
+            if cutlass.const_expr(self.tokens > 1):
+                gates = (a, lower_bound, dt, features)
+                targets = (shared_decay, barrier, value_base)
+                return self.multicast(acc, gates, targets, lane)
+            raw_g = []
+            for r in cutlass.range_constexpr(per):
+                partial = cute.make_rmem_tensor(chains, Float32)
+                for chain in cutlass.range_constexpr(chains):
+                    selected = acc[chain, 2 * r, 0]
+                    if quarter & 1:
+                        selected = acc[chain, 2 * r, 2]
+                    if cutlass.const_expr(tiles > 2 * r + 1):  # noqa: SIM102
+                        if quarter >> 1:
+                            selected = acc[chain, 2 * r + 1, 0]
+                            if quarter & 1:
+                                selected = acc[chain, 2 * r + 1, 2]
+                    partial[chain] = selected
+                value = partial[0]
+                if cutlass.const_expr(mode == 1):
+                    value = Float32(0) + (
+                        ((value + partial[1]) + partial[2]) + partial[3]
+                    )
+                raw_g.append(Float32(BFloat16(value)))
+
+        decay = []
+        for r in cutlass.range_constexpr(per):
+            gate = cute.math.exp(a, fastmath=True) * (raw_g[r] + dt[r])
+            gate = sigmoid(gate, self.fast_math)
+            decay.append(cute.math.exp(lower_bound * gate, fastmath=True))
+        packets: cutlass.Constexpr = rows // 4
+        packet, destination = lane % packets, lane // packets
+        values = cute.make_rmem_tensor(4, Float32)
+        for e in cutlass.range_constexpr(4):
+            source = 4 * packet + e
+            if cutlass.const_expr(per == 2):
+                feature = source & 31
+                source = (feature & 7) << 2 | (feature >> 4) << 1 | feature >> 3 & 1
+                low = cute.arch.shuffle_sync(decay[0], source)
+                high = cute.arch.shuffle_sync(decay[1], source)
+                values[e] = low
+                if packet >= 8:
+                    values[e] = high
+            elif cutlass.const_expr(per > 2):  # one CTA per head: 128 features
+                feature = source & 31
+                source = (feature & 7) << 2 | (feature >> 4) << 1 | feature >> 3 & 1
+                shuffled = [cute.arch.shuffle_sync(d, source) for d in decay]
+                values[e] = shuffled[0]
+                for r in cutlass.range_constexpr(1, per):
+                    if packet >> 3 == r:
+                        values[e] = shuffled[r]
+            else:
+                if cutlass.const_expr(mode != 2):
+                    source = (source & 7) << 2 | (source >> 4) << 1 | source >> 3 & 1
+                values[e] = cute.arch.shuffle_sync(decay[0], source)
+        offset = shared_decay.iterator + value_base + 4 * packet
+        pointer = cute.make_ptr(
+            Float32, offset.llvm_ptr, cute.AddressSpace.smem, assumed_align=16
+        )
+        words = tuple(float_bits(values[e]) for e in range(4))
+        if cutlass.const_expr(self.early):
+            cute.arch.cluster_wait()
+        store_async_dsmem(pointer, words, barrier, destination)
+
+    @cute.jit
+    def multicast(self, acc, gates, targets, lane):
+        """Decays (q > 1) of the lane's D fragments (feature rows, token columns 2
+        quarter + {0, 1}) to every CTA's shared [token, 128] decays."""
+        (a, lower_bound, dt, features), (shared_decay, barrier, value_base) = (
+            gates,
+            targets,
+        )
+        quarter = lane & 3
+        if cutlass.const_expr(self.early):
+            cute.arch.cluster_wait()
+        for i in cutlass.range_constexpr(len(features)):
+            tile, e = i // 2, (i % 2) * 2
+            for column in cutlass.range_constexpr(2):
+                value = acc[0, tile, e + column]
+                if cutlass.const_expr(self.gate_mode == 1):
+                    parts = [acc[c, tile, e + column] for c in range(4)]
+                    value = Float32(0) + (((value + parts[1]) + parts[2]) + parts[3])
+                raw_g = Float32(BFloat16(value))
+                gate = cute.math.exp(a, fastmath=True) * (raw_g + dt[i])
+                gate = sigmoid(gate, self.fast_math)
+                decay = float_bits(cute.math.exp(lower_bound * gate, fastmath=True))
+                token = 2 * quarter + column
+                offset = token * 128 + value_base + features[i]
+                pointer = shared_decay.iterator + offset
+                if token < self.tokens:
+                    for rank in cutlass.range_constexpr(self.cluster):
+                        store_async_dsmem(pointer, decay, barrier, rank)
+
+    @staticmethod
+    def split_pairs(wr, tile, chunk, low, high):
+        """16 bytes of rows r and r + 8 hold the A fragments of chunks 2c and 2c + 1."""
+        for h in range(2):
+            c = 2 * chunk + h
+            wr[tile, c, 0], wr[tile, c, 2] = low[2 * h : 2 * h + 2]
+            wr[tile, c, 1], wr[tile, c, 3] = high[2 * h : 2 * h + 2]
+
+    @cute.jit
+    def setup(self, barrier, gate_barrier, partition, thread):
+        """Thread 0 arms the output (leader) and decay barriers; the CTA arrives on the
+        cluster barrier peers wait on before their sends."""
+        tokens, cluster = self.tokens, self.cluster
+        owned = [(barrier, partition == 0, 512)]  # token t's outputs: CTA t % cluster
+        if cutlass.const_expr(tokens > 1):
+            owned = [
+                (barrier + k, partition + k * cluster < tokens, 512)
+                for k in range(-(-tokens // cluster))
+            ]
+        for bar, cta, size in (*owned, (gate_barrier, True, 512 * tokens)):
+            if cta & (thread == 0):
+                cute.arch.mbarrier_init(bar, 1)
+                cute.arch.mbarrier_arrive_and_expect_tx(bar, size)
+                cute.arch.mbarrier_init_fence()
+        cute.arch.cluster_arrive_relaxed()
+
+    @cute.jit
+    def row_dot(self, history, vector, p):
+        partial = cute.make_rmem_tensor(self.chunks, Float32)
+        for j in cutlass.range_constexpr(self.chunks):
+            row = (history[p, j, e] for e in range(4))
+            partial[j] = dot(row, (vector[j, e] for e in range(4)), ftz=self.fast_math)
+        for shift in cutlass.range_constexpr(self.chunks.bit_length() - 2, -1, -1):
+            for j in cutlass.range_constexpr(1 << shift):
+                partial[j] += partial[j + (1 << shift)]
+        return partial[0]
+
+    @cute.jit
+    def output_position(self, row: Int32):
+        return ((row & 7) << 4) | ((row >> 5) << 2) | ((row & 31) >> 3)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Prepared KDA core. One packed projection supplies Q/K/V, the output gate, F_A and
+beta; per-batch policies select the gate schedule and the projection row pitch."""
+
+from functools import cache
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import fake, layout, stream, to_cute
+from ..common.policy import batch_table
+from .kernel import KdaSplit
+
+# KdaSplit schedule "early.lanes.occupancy.cluster.evict" per batch by heads (TP8 12,
+# TP16 6), from GB300 sweeps: early predecessor-independent loads (MMA gate modes only;
+# 2 also stages the state rows in shared memory, 3 also the F_B operands), lanes per
+# state row, minimum CTAs per SM, CTAs per head (0: defaults) and L2 evict-first state
+# traffic (staged only).
+SCHEDULE = {12: "16-32:3.16.3.2.1,64:3.16.4.2.1", 6: "16:1.0.0,32-64:3.16.3.2"}
+# Verify (q > 1) schedules by request count B, from GB300
+# sweeps over q 4-8: no staging; eight-CTA clusters, then fewer as B grows.
+SPEC_SCHEDULE = {
+    12: "1:1.8.2.8,2-4:1.8.3.8,5-12:1.8.3.4,13-16:1.8.3.2",
+    6: "1:1.8.2.8,2:1.8.0.8,3-12:1.8.3.8,13-16:1.8.3.4",
+}
+
+
+@cache
+def executable(specs, heads, gate_mode, has_bias, device, tokens=1):
+    args = tuple(fake(shape, stride, dtype) for shape, stride, dtype in specs)
+    batch = specs[0][0][0]
+    table = SCHEDULE if tokens == 1 else SPEC_SCHEDULE
+    schedule = batch_table(table.get(heads, ""))
+    schedule = schedule.get(batch // tokens, (0, 0, 0)) if gate_mode != 2 else (0,) * 3
+    if schedule[0] >= 2 and specs[9][1][-2:] != (128, 1):  # staging needs whole rows
+        schedule = (1, *schedule[1:4])  # no evict hint without staging
+    requests = batch // tokens
+    kernel = KdaSplit(heads, requests, gate_mode, has_bias, schedule, tokens)
+    scalars = (cutlass.Float32(v) for v in (-5, 128**-0.5, 1e-6))
+    with torch.accelerator.device_index(device):
+        return cute.compile(
+            kernel, *args[:13], *scalars, *args[13:], make_fake_stream()
+        )
+
+
+def decode_conv_weight(weight):
+    """The [3, width, D] (Q, K, V) decode layout of the [3 * D, 1, width] conv
+    weight."""
+    width = weight.shape[-1]
+    return weight.detach().view(3, -1, width).transpose(1, 2).contiguous()
+
+
+class PreparedKda:
+    def __init__(self, attention, policies, conv_dim_first):
+        bias, dt = attention.conv1d.bias, attention.dt_bias
+        self.heads = attention.A_log.numel()
+        self.dim, self.has_bias = self.heads * 128, bias is not None
+        conv_w = decode_conv_weight(attention.conv1d.weight)
+        bias, fb = dt if bias is None else bias, attention.f_b_proj.weight
+        norm = attention.o_norm.weight.detach().float()  # BF16 slows the fused kernel
+        self.weights = (conv_w, bias, fb, attention.A_log, dt, norm)
+        self.device, self.policies = dt.device, policies
+        scalars = (attention.gate_lower_bound, 128**-0.5, attention.o_norm.eps)
+        self.scalars = tuple(map(cutlass.Float32, scalars))
+        self.conv_dim_first, self.kernels = conv_dim_first, {}
+
+    def _operands(self, caches, projected, slots, out, accepted=None):
+        """Kernel tensors in launch order: 13 before the scalars, then F_B (and the
+        verify step's accepted counts)."""
+        conv, state = caches[:2]
+        conv = conv if self.conv_dim_first else conv.transpose(-1, -2)
+        w, bias, fb, a_log, dt, norm = self.weights
+        d, h = self.dim, self.heads
+        p = lambda start, width: projected[:, start : start + width]
+        ops = (p(0, 3 * d), w, bias, conv, p(4 * d, 128), p(4 * d + 128, h), a_log, dt)
+        tail = (fb,) if accepted is None else (fb, accepted)
+        return ops + (slots, state, out, p(3 * d, d), norm) + tail
+
+    def prepare_caches(self, caches):
+        """Policy keys: a batch, or (requests, tokens) for verify steps."""
+        for key, (pitch, mode) in self.policies.items():
+            requests, tokens = key if isinstance(key, tuple) else (key, 1)
+            batch, meta = requests * tokens, dict(device="meta")
+            projected = torch.empty((batch, pitch), dtype=torch.bfloat16, **meta)
+            out = torch.empty(batch, self.dim, dtype=torch.bfloat16, **meta)
+            accepted = torch.empty(requests, dtype=torch.int32, **meta)
+            shape = (requests, tokens) if tokens > 1 else (batch,)
+            slots = torch.empty(shape, dtype=torch.int32, **meta)
+            accepted = accepted if tokens > 1 else None
+            operands = self._operands(caches, projected, slots, out, accepted)
+            self.kernels[key] = executable(
+                tuple(map(layout, operands)),
+                *(self.heads, mode, self.has_bias, self.device.index, tokens),
+            )
+        return self
+
+    def run(self, projected, out, *, caches, slots, accepted=None):
+        """``accepted`` [>= B] (verify steps): ``slots`` is [>= B, q] contiguous."""
+        batch = projected.shape[0]
+        key, rows = batch, batch
+        if accepted is not None:
+            rows = batch // slots.shape[1]
+            key, accepted = (rows, slots.shape[1]), accepted[:rows]
+            if slots.stride() != (slots.shape[1], 1):  # e.g. a block-table slice
+                slots = slots.contiguous()
+        tensors = self._operands(caches, projected, slots[:rows], out, accepted)
+        views = tuple(to_cute(t, 16) for t in tensors)
+        self.kernels[key](*views[:13], *self.scalars, *views[13:], stream(self.device))
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/dcp.py

@@ -0,0 +1,378 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MLA decode under decode context parallelism (DCP = TP, rank d owns head block d)
+without collectives, three kernels per layer. Each rank's KV cache holds 1 / DCP of
+every context.
+
+  MlaQuery     FP8 query rows [B, H, 576] (+ the KV insert where this rank owns the
+               slot); each CTA multicasts its bytes into every rank's Lamport q
+               mailbox [B, DCP*H, 576], then polls its share until all landed
+  tokenspeed   all DCP*H heads over the local KV shard -> BF16 rows + LSE (log2)
+  DcpValue     re-arms the q mailbox; rows pushed to their head owner's Lamport
+               mailbox (peer stores); the owner LSE-merges its H heads over the DCP
+               sources, then W_UV and the gate (padding from the global lengths)
+"""
+
+import math
+from functools import cache
+
+import cutlass
+import cutlass.cute as cute
+import torch
+import torch.distributed._symmetric_memory as symm_mem
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import fake
+from ..common.primitives import (
+    alloc,
+    at,
+    bf16x2,
+    bits_float,
+    float_bits,
+    launch,
+    load_volatile,
+    load_words,
+    mma_into,
+    no_neg_zero,
+    pack_pair,
+    published,
+    sigmoid,
+    store_relaxed,
+    store_shared_words,
+    store_words,
+    unpack_bf16,
+    warp_sum,
+    zeros,
+)
+
+LSE_SENT = 0xFFFFFFFF  # a NaN pattern the merge never sends
+# The LSE of an empty or invalid source: finite, so the merge never forms inf - inf.
+EMPTY_LSE = -3.0e38
+Q_SENT = 0xFFFFFFFF  # four FP8 0xFF bytes: never written (satfinite, canonical NaN)
+PAIR_SENT = 0x80000000  # BF16 pair (+0, -0); sent pairs have no negative zero
+
+
+class DcpBox:
+    """Model-owned DCP state over ``group``, Lamport mailboxes (sentinel-filled,
+    re-armed by their reader): q [2, CAP, DCP*H, 576] FP8 bytes (multicast stores),
+    partials [2, DCP, CAP, H, 512] BF16 and LSE [2, DCP, CAP, H] FP32 (peer stores),
+    and the local tokenspeed rows / LSE. Parity alternates by layer. ``interleave``:
+    the KV cache's DCP interleave (global masking of spec verify rows)."""
+
+    def __init__(self, group, device, size, rank, heads, cap=64, interleave=1):
+        self.size, self.rank, self.heads, self.cap = size, rank, heads, cap
+        self.interleave = interleave
+        empty = symm_mem.empty
+        self.q = empty((2, cap, size * heads, 576), dtype=torch.uint8, device=device)
+        shape = (2, size, cap, heads)
+        self.partial = empty((*shape, 512), dtype=torch.bfloat16, device=device)
+        self.lse = empty(shape, dtype=torch.float32, device=device)
+        tensors = (self.q, self.partial, self.lse)
+        handles = [symm_mem.rendezvous(t, group) for t in tensors]
+        self.q_mc = int(handles[0].multicast_ptr or 0)
+        if not self.q_mc:
+            raise RuntimeError("Kimi K3 specialized MLA DCP needs multicast memory")
+        i64 = dict(dtype=torch.int64, device=device)
+        self.peers = torch.tensor(handles[1].buffer_ptrs, **i64)
+        self.lse_peers = torch.tensor(handles[2].buffer_ptrs, **i64)
+        self.q.fill_(0xFF)
+        self.partial.view(torch.int32).fill_(-0x80000000)
+        self.lse.view(torch.int32).fill_(-1)
+        rows = (cap, 1, size * heads)
+        self.rows = torch.empty((*rows, 512), dtype=torch.bfloat16, device=device)
+        self.rows_lse = torch.empty(rows, dtype=torch.float32, device=device)
+        torch.accelerator.synchronize(device)
+        torch.distributed.barrier(group=group)
+
+    def mc(self, parity):
+        """The q mailbox multicast address of ``parity``."""
+        return Int64(self.q_mc + parity * self.q[0].numel())
+
+
+class DcpValue:
+    """The DCP tail of one MLA layer, one 512-thread CTA per (head, 16 requests).
+    Every rank re-arms this parity's q mailbox (tokenspeed has read it) and pushes all
+    its local tokenspeed rows (+ LSE) to their head owner's Lamport mailbox (empty
+    local shards: zeros, LSE -inf). Warp r then LSE-weights request t0 + r of this
+    head over the DCP sources (exp2(lse - max) / sum in FP32), the CTA merges the
+    polled rows in FP32
+    (one BF16 rounding, re-arming what it read) into shared memory, and W_UV (BF16
+    MMA, FP32 accumulation) and the sigmoid gate follow as in MlaValue: warp w owns
+    outputs 32 (w // 4) + 8 (w % 4) .. + 7. Rows are tokens, ``q`` per request."""
+
+    def __init__(self, heads, size, batch, sms=148, q=1):
+        self.heads, self.size, self.batch = heads, size, batch
+        self.q = q
+        # Push-only CTAs (z > 0) spread the NVLink stores over more SMs: up to one
+        # wave, at least 4 16 B chunks per thread.
+        merge = heads * ((batch + 15) // 16)
+        chunks = batch * size * heads * 64
+        pushers = min(max(1, sms // merge), -(-chunks // (merge * 512 * 4)))
+        self.grid = (heads, (batch + 15) // 16, pushers)
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, self.grid, 512)
+
+    @cute.kernel
+    def kernel(self, args):
+        rows, rows_lse, lens, box, lbox, peers, lpeers, mailbox = args[0:8]
+        w_uv, gate, totals, output, parity, rank = args[8:14]
+        H, D, B, cap = self.heads, self.size, self.batch, box.shape[1]
+        smem = cutlass.utils.SmemAllocator()
+        latent = alloc(smem, BFloat16, (16, 544), 128)
+        scales = alloc(smem, Float32, (16, D))
+        tid, _, _ = cute.arch.thread_idx()
+        head, tile, pusher = cute.arch.block_idx()
+        lane, warp = tid % 32, tid // 32
+        group, quarter = lane // 4, lane % 4
+        v = 32 * (warp // 4) + 8 * (warp % 4) + group
+        # B fragments: 8 latent values per 32-wide chunk; A uses the same permutation.
+        wb = [
+            load_words(at(w_uv, head, 32 * c + 8 * quarter, v), count=4)
+            for c in range(16)
+        ]
+        cute.arch.griddepcontrol_wait()
+        threads = H * self.grid[1] * self.grid[2] * 512
+        me = ((pusher * self.grid[1] + tile) * H + head) * 512 + tid
+        for c in range(me, B * D * H * 36, threads):  # 16 B q mailbox chunks
+            store_words(mailbox.iterator + c * 16, (Uint32(Q_SENT),) * 4)
+        for i in range(me, B * D * H * 64, threads):
+            r, k = i // 64, i % 64  # r = b * D * H + global head
+            b, gh = r // (D * H), r % (D * H)
+            words = (Uint32(0),) * 4
+            if lens[b // self.q] > 0:
+                words = load_words(at(rows, b, 0, gh, 8 * k), count=4, mutable=True)
+            words = tuple(no_neg_zero(w) for w in words)
+            slot = ((parity * D + rank) * cap + b) * H + gh % H
+            store_relaxed(peers[gh // H] + (Int64(slot) * 512 + 8 * k) * 2, words)
+        for i in range(me, B * D * H, threads):
+            b, gh = i // (D * H), i % (D * H)
+            value, n = rows_lse[b, 0, gh], lens[b // self.q]
+            bad = (value != value) | (value <= Float32(EMPTY_LSE))  # NaN, -inf
+            if (n <= 0) | bad | (value >= Float32(-EMPTY_LSE)):  # or +inf
+                value = Float32(EMPTY_LSE)
+            slot = ((parity * D + rank) * cap + b) * H + gh % H
+            store_relaxed(lpeers[gh // H] + Int64(slot) * 4, (float_bits(value),))
+        cute.arch.griddepcontrol_launch_dependents()
+        if pusher == 0:
+            self.merge(args, latent, scales, wb, tid)
+
+    @cute.jit
+    def merge(self, args, latent, scales, wb, tid):
+        """The owner's merge of request rows t0.. of this head, then W_UV + gate."""
+        box, lbox, gate, totals, output = args[3], args[4], args[9], args[10], args[11]
+        D, B = self.size, self.batch
+        head, tile, _ = cute.arch.block_idx()
+        lane, warp = tid % 32, tid // 32
+        group, quarter = lane // 4, lane % 4
+        t0 = tile * 16
+        col = head * 128 + 32 * (warp // 4) + 8 * (warp % 4) + 2 * quarter
+        gates = []
+        for r in cutlass.range_constexpr(2):
+            word = Uint32(0)
+            if t0 + group + 8 * r < B:
+                word = load_words(
+                    at(gate, t0 + group + 8 * r, col), count=1, mutable=True
+                )[0]
+            gates.append(unpack_bf16(word))
+        if t0 + warp < B:  # warp r: the source weights of request t0 + r
+            lse = Float32(EMPTY_LSE)
+            if lane < D:
+                source = at(lbox, lane, t0 + warp, head)
+                word = load_volatile(source)[0]
+                while word == Uint32(LSE_SENT):
+                    word = load_volatile(source)[0]
+                store_words(source, (Uint32(LSE_SENT),))
+                lse = bits_float(word)
+            top = lse
+            for shift in cutlass.range_constexpr(4, -1, -1):
+                top = max(top, cute.arch.shuffle_sync_bfly(top, 1 << shift))
+            weight = Float32(0)
+            if lse > Float32(EMPTY_LSE):
+                weight = cute.math.exp2(lse - top, fastmath=True)
+            total = warp_sum(weight)
+            scale = Float32(0)
+            if total > 0:
+                scale = weight / total
+            if lane < D:
+                scales[warp, lane] = scale
+        cute.arch.sync_threads()
+        chunk = tid % 64
+        for half in cutlass.range_constexpr(2):  # rows r and r + 8, 16 B chunk
+            r = tid // 64 + 8 * half
+            merged = zeros(8, Float32)
+            if t0 + r < B:
+                for s0 in cutlass.range_constexpr(0, D, 8):  # 8 polls in flight
+                    sources = [
+                        at(box, s, t0 + r, head, 8 * chunk)
+                        for s in range(s0, min(D, s0 + 8))
+                    ]
+                    loads = [load_volatile(source, count=4) for source in sources]
+                    for i in cutlass.range_constexpr(len(sources)):
+                        words = loads[i]
+                        while not published(words):
+                            words = load_volatile(sources[i], count=4)
+                        store_words(sources[i], (Uint32(PAIR_SENT),) * 4)
+                        w = scales[r, s0 + i]
+                        for e in cutlass.range_constexpr(4):
+                            lo, hi = unpack_bf16(words[e])
+                            merged[2 * e] += w * lo
+                            merged[2 * e + 1] += w * hi
+            pairs = tuple(
+                pack_pair(merged[2 * e], merged[2 * e + 1], policy=None)
+                for e in range(4)
+            )
+            store_shared_words(latent.iterator + r * 544 + 8 * chunk, pairs)
+        cute.arch.sync_threads()
+        parts = zeros((2, 4), Float32)  # two accumulation chains
+        for c in cutlass.range_constexpr(16):
+            x, y = (
+                load_words(
+                    latent.iterator + r * 544 + 32 * c + 8 * quarter,
+                    count=4,
+                    shared=True,
+                )
+                for r in (group, group + 8)
+            )
+            for s in cutlass.range_constexpr(2):
+                a = (x[2 * s], y[2 * s], x[2 * s + 1], y[2 * s + 1])
+                b = (wb[c][2 * s], wb[c][2 * s + 1])
+                mma_into(parts[s, None], a, b, bf16=True)
+        acc = [parts[0, k] + parts[1, k] for k in range(4)]
+        for r in cutlass.range_constexpr(2):
+            t = t0 + group + 8 * r
+            if t < B:
+                # Padded rows may hold a NaN gate; keep inactive rows zero.
+                lo, hi = Float32(0), Float32(0)
+                if totals[t // self.q] > 0:
+                    lo, hi = (
+                        Float32(BFloat16(acc[2 * r + e])) * sigmoid(gates[r][e], False)
+                        for e in range(2)
+                    )
+                store_words(at(output, t, col), (bf16x2(lo, hi),))
+
+
+@cache
+def compile_value(heads, size, batch, cap, uv, pitch, device, sms=148, q=1):
+    """DcpValue of one batch (``q`` rows per request); ``uv``: the (shape, stride,
+    dtype) of W_UV, ``pitch``: the projected rows' (the gate follows the 2112 query /
+    KV columns)."""
+    f32, i32 = torch.float32, torch.int32
+    w_uv, gate = fake(*uv), fake((batch, heads * 128), (pitch, 1))
+    rows = fake((batch, 1, size * heads, 512))
+    rows_lse = fake((batch, 1, size * heads), None, f32, 4)
+    lens = fake((batch // q,), None, i32, 4)
+    box = fake((size, cap, heads, 512))
+    lbox = fake((size, cap, heads), None, f32, 4)
+    peers = fake((size,), None, torch.int64, 8)
+    mailbox = fake((cap, size * heads, 576), None, torch.uint8)
+    output = fake((batch, heads * 128), (heads * 128, 1))
+    args: tuple = (rows, rows_lse, lens, box, lbox, peers, peers, mailbox, w_uv)
+    args += (gate, lens, output, Int32(0), Int32(0), make_fake_stream())
+    with torch.accelerator.device_index(device):
+        return cute.compile(DcpValue(heads, size, batch, sms, q), *args)
+
+
+_WORKSPACE: dict[tuple[torch.device, int], torch.Tensor] = {}
+
+
+def stock_decode(
+    query,
+    cache,
+    block_table,
+    lengths,
+    max_seq,
+    scales,
+    out,
+    lse=None,
+    min_split=1,
+    pdl=True,
+    cp=(1, 0, 1),
+    totals=None,
+    pack=False,
+):
+    """tokenspeed_mla_decode's FP8 path with caller-owned ``out`` and ``lse`` (log2)
+    buffers and no per-call allocation: the split count follows from (B, max_seq,
+    min_split); ``lengths`` bound the work, empty splits are skipped by its reducer.
+    The q rows of a request mask causally; under DCP in global positions (``cp``:
+    world, rank, interleave; ``totals``: the global lengths). ``pack``: M128 tiles hold
+    consecutive (query row, head) rows across the q rows."""
+    import tvm_ffi
+    from tokenspeed_mla import mla_decode as md
+
+    batch, q_len, heads, _ = query.shape
+    fp8, device = torch.float8_e4m3fn, query.device
+    arch = torch.cuda.get_device_capability(device)
+    tiler, _ = md.select_mla_decode_tilers(
+        heads, q_len, is_fp8=True, compute_capability=arch
+    )
+    pack = pack and tiler[0] == 128
+    fold = 1 if pack else md.get_mla_decode_fold_sq_factor(heads, q_len, tiler[0])
+    rows = (q_len // fold, heads * fold)
+    if pack:
+        rows = (md.compute_q_tile_layout(heads, q_len, 128)[1], 128)
+    sms = md.get_num_sm(device)
+    split, size = md._get_split_kv_and_workspace_size(
+        batch, *rows, 512, sms, max_seq, fp8, tiler, min_split
+    )
+    # One split-KV workspace per (device, heads), shared by every MLA layer: the calls
+    # are stream ordered. PreparedMla.prepare_cache runs each prepared (batch, q) before
+    # any capture, so the buffer reaches its final size first; growing it during a
+    # capture would leave the graphs captured before with the freed buffer.
+    key = device, heads
+    if key not in _WORKSPACE or _WORKSPACE[key].numel() < size:
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "the MLA split-KV workspace must be sized before capture"
+            )
+        bound = max(size, sms * heads * 513 * 4)
+        _WORKSPACE[key] = torch.empty(bound, dtype=torch.int8, device=device)
+    kernel = md._get_compiled_mla_kernel(
+        torch_dtype=fp8,
+        page_size=cache.shape[1],
+        kv_lora_rank=512,
+        qk_rope_head_dim=64,
+        is_persistent=False,
+        is_var_seq=True,
+        is_var_split_kv=False,
+        compute_capability=arch,
+        has_local_visible_lens=False,
+        skip_correction_threshold=8.0,
+        is_workspace_size_zero=size == 0,
+        fold_sq_factor=fold,
+        causal_mask=True,
+        num_heads=heads,
+        seq_len_q=q_len,
+        window_left=-1,
+        cp_world=cp[0],
+        cp_rank=cp[1],
+        cp_interleave_size=cp[2],
+        use_pdl=pdl,
+        return_lse=lse is not None,
+        reducer_d_tiles=md._get_reducer_d_tiles(batch, q_len, heads, sms, split),
+        reducer_max_splits=md._get_reducer_max_splits(split),
+        pack_q=pack,
+        partial_fp16=False,
+    )
+    cache = cache.view(fp8)
+    softmax, output_scale = scales
+    args = [query[..., :512], query[..., 512:], cache[..., :512], cache[..., 512:]]
+    args += [block_table, out, lse, _WORKSPACE[key][:size] if size else None]
+    args += [Int32(split), lengths, lengths if totals is None else totals, None]
+    args += [Float32(softmax), Float32(output_scale)]
+    with tvm_ffi.use_torch_stream():
+        kernel(*args)
+
+
+def fill_splits(batch, device):
+    """Splits filling the SMs (one CTA pair each); tokenspeed caps itself at 32."""
+    from tokenspeed_mla import get_num_sm
+
+    return min(256, max(1, get_num_sm(device) // (2 * batch)))
+
+
+def local_max(max_context, size):
+    """Longest local shard of a ``max_context`` sequence (block-interleaved)."""
+    return math.ceil(max_context / (size * 64)) * 64

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/prepared.py

@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Prepared MLA core: the packed input projection feeds the query kernel, then either
+tokenspeed_mla_decode (core "tokenspeed") or FlashInfer's trtllm-gen MLA decode (core
+"trtllm"; FP8 Q and KV as tokenspeed, FP32 softmax and accumulation) into the BF16
+latent, and MlaValue (W_UV / gate). Under DCP (``dcp``: a DcpBox) the tokenspeed core
+runs over the local KV shard for every DCP rank's heads between the query's q
+publication and ``DcpValue``'s merge (see mla.dcp). Spec verify: ``q`` query tokens
+per request (rows grouped by request, causal among them)."""
+
+from functools import cache
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import fake, layout, prepare_shared_carveout, stream, to_cute
+from .dcp import (
+    compile_value,
+    fill_splits,
+    local_max,
+    stock_decode,
+)
+from .query import MlaQuery
+from .value import MlaValue
+
+# tokenspeed (SM-filling splits) streams long contexts fast; trtllm-gen wins from B16.
+TRTLLM_CORE_BATCH, CORES = 16, ("tokenspeed", "trtllm")
+_FLASHINFER: dict[tuple[torch.device, int], tuple[torch.Tensor, torch.Tensor]] = {}
+
+
+@cache
+def compile_core(batch, heads, pitch, specs, device, dcp=0, sms=148, q=1):
+    """CuTe kernels (query, value) of one batch of rows (``q`` per request); under DCP
+    (``dcp``: its size) the query publishes and the value is DcpValue."""
+    qnorm, kvnorm, qb, uk, uv, qinv, kinv, kv = (
+        fake(shape, stride, dtype, 4 if dtype == torch.float32 else 16)
+        for shape, stride, dtype in specs
+    )
+    query = fake((batch, heads, 576), (heads * 576, 576, 1), torch.uint8)
+    projected = fake((batch, 2112), (pitch, 1))
+    slots = fake((batch,), (1,), torch.int64, 8)
+    lengths = fake((batch // q,), (1,), torch.int32, 4)
+    gate = fake((batch, heads * 128), (pitch, 1))
+    output = fake((batch, heads * 128), (heads * 128, 1))
+    latent = fake((batch, heads, 512))
+    f32, s = cutlass.Float32, make_fake_stream()
+    cp_async = heads == 12 and batch == 1
+    mla_query = MlaQuery(heads, batch, q_b_cp_async=cp_async, dcp=dcp)
+    q_args: tuple = (projected, qnorm, kvnorm, qb, uk, kv, slots, qinv, kinv)
+    q_args += (f32(1e-6), query)
+    if dcp:  # this parity's q mailbox, its multicast address, rank
+        mailbox = fake((64, dcp * heads, 576), None, torch.uint8)
+        q_args += (mailbox, cutlass.Int64(0), cutlass.Int32(0))
+    with torch.accelerator.device_index(device):
+        kernels = [cute.compile(mla_query, *q_args, s)]
+        if batch > 1:
+            kernels = [prepare_shared_carveout(kernels[0], device, 100)]
+        if dcp:
+            uv_spec = specs[4]
+            kernels.append(
+                compile_value(heads, dcp, batch, 64, uv_spec, pitch, device, sms, q)
+            )
+        else:
+            v_args = (latent, uv, gate, lengths, output, s)
+            value = MlaValue(heads, batch, q)
+            kernels.append(cute.compile(value, *v_args))
+    return tuple(kernels)
+
+
+def core_for(batch, core="auto", q=1):
+    """``core`` unless "auto": trtllm-gen from B16, else tokenspeed (spec verify,
+    q > 1: tokenspeed, whose SM-filling splits stream long contexts far faster)."""
+    if core == "auto":
+        core = "trtllm" if batch >= TRTLLM_CORE_BATCH and q == 1 else "tokenspeed"
+    if core not in CORES:
+        raise ValueError(
+            f"Kimi K3 MLA core {core!r}: expected auto, tokenspeed or trtllm"
+        )
+    return core
+
+
+def pad_requests(batch, block_table, lengths, totals=None):
+    """Request metadata for ``batch`` rows: PIECEWISE graphs pad the tokens but not
+    the requests, so missing rows become empty requests (length 0 on page 0)."""
+    rows = min(lengths.shape[0], block_table.shape[0])
+    if rows >= batch:
+        return block_table, lengths, totals
+    pad = torch.nn.functional.pad
+    block_table = pad(block_table[:rows], (0, 0, 0, batch - rows))
+    lengths = pad(lengths[:rows], (0, batch - rows))
+    totals = None if totals is None else pad(totals[:rows], (0, batch - rows))
+    return block_table, lengths, totals
+
+
+def trtllm_decode(query, cache, block_table, lengths, max_seq, scales, out):
+    """FlashInfer trtllm-gen MLA decode of FP8 ``query`` [B, q, H, 576] (causal among
+    the q rows) over the FP8 paged cache into BF16 ``out`` [B, q, H, 512], PDL chained;
+    zeroed workspace and multi-CTA counter (self-resetting) owned here, one per (device,
+    heads)."""
+    from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+    from flashinfer.utils import get_trtllm_gen_multi_ctas_kv_counter_bytes
+
+    device, heads = query.device, query.shape[2]
+    if (device, heads) not in _FLASHINFER:
+        sms = torch.cuda.get_device_properties(device).multi_processor_count
+        count = get_trtllm_gen_multi_ctas_kv_counter_bytes(64, heads, sms)
+        workspace = torch.zeros(128 << 20, dtype=torch.int8, device=device)
+        counter = torch.zeros(count, dtype=torch.uint8, device=device)
+        _FLASHINFER[device, heads] = workspace, counter
+    workspace, counter = _FLASHINFER[device, heads]
+    softmax, output_scale = scales
+    trtllm_batch_decode_with_kv_cache_mla(
+        query=query,
+        kv_cache=cache.view(torch.float8_e4m3fn).unsqueeze(1),
+        workspace_buffer=workspace,
+        qk_nope_head_dim=128,
+        kv_lora_rank=512,
+        qk_rope_head_dim=64,
+        block_tables=block_table,
+        seq_lens=lengths,
+        max_seq_len=max_seq,
+        out=out,
+        bmm1_scale=softmax,
+        bmm2_scale=output_scale,
+        enable_pdl=True,
+        backend="trtllm-gen",
+        multi_ctas_kv_counter_buffer=counter,
+    )
+
+
+class PreparedMla:
+    """``max_context`` (the deployment's longest sequence) bounds the split count;
+    ``core`` forces "tokenspeed" or "trtllm". ``dcp``: the model's DcpBox, with this
+    layer's mailbox ``parity``. ``queries``: the query lengths q (spec verify: 1 + the
+    draft length) prepared for every batch of rows divisible by q."""
+
+    def __init__(
+        self,
+        attention,
+        pitches,
+        query_scratch,
+        *,
+        max_context=None,
+        core="auto",
+        dcp=None,
+        parity=0,
+        queries=(1,),
+    ):
+        self.heads = attention.W_UK_T.shape[0]
+        norms = attention.q_a_layernorm.weight, attention.kv_a_layernorm.weight
+        self.weights: tuple = (*norms, attention.q_b_proj.weight, attention.W_UK_T)
+        self.weights += (attention.W_UV, attention._q_scale_inv, attention._k_scale_inv)
+        self.device, self.pitches = attention.q_b_proj.weight.device, pitches
+        k_scale = attention._k_scale_float
+        softmax = attention.scale * attention._q_scale_float * k_scale
+        self.eps = cutlass.Float32(attention.rms_norm_eps)
+        self.stock_scales = softmax, k_scale
+        self.views = tuple(
+            to_cute(w, 4 if w.dtype == torch.float32 else 16) for w in self.weights
+        )
+        self.query_scratch, self.kernels = query_scratch, {}
+        shapes = [(b, q) for b in pitches for q in queries if b % q == 0]
+        self.cores = {(b, q): core_for(b, core, q) for b, q in shapes}
+        # The split count only depends on this bound; seq_lens bound the work.
+        self.max_seq = max_context or 1 << 17
+        self.dcp, self.parity = dcp, parity
+        if dcp is not None:
+            self.cores = dict.fromkeys(shapes, "tokenspeed")
+            self.max_seq = local_max(self.max_seq, dcp.size)
+        shape, bf16 = (max(pitches, default=0), self.heads, 512), torch.bfloat16
+        self.latent = torch.empty(shape, dtype=bf16, device=self.device)
+
+    def prepare_cache(self, kv):
+        specs = tuple(layout(w) for w in (*self.weights, kv.view(torch.uint8)))
+        for batch, q in self.cores:
+            device = self.device.index
+            dcp = self.dcp.size if self.dcp is not None else 0
+            sms = torch.cuda.get_device_properties(self.device).multi_processor_count
+            pitch = self.pitches[batch]
+            args = batch, self.heads, pitch, specs, device, dcp, sms, q
+            query = to_cute(self.query_scratch[:batch])
+            latent = to_cute(self.latent[:batch])
+            self.kernels[batch, q] = (*compile_core(*args), query, latent)
+            if self.dcp is not None:
+                self.prepare_dcp(batch)
+            if kv.is_cuda:  # compile before graph capture
+                table = torch.zeros(batch, 1, dtype=torch.int32, device=kv.device)
+                ones = torch.full((batch,), q, dtype=torch.int32, device=kv.device)
+                self.stock_core(kv, table, ones, batch, q, ones)
+        return self
+
+    def prepare_dcp(self, batch):
+        """(MlaQuery's publish arguments, DcpValue's mailbox arguments, parity and
+        rank) of one batch."""
+        box, p = self.dcp, self.parity
+        mailbox, ids = to_cute(box.q[p]), (cutlass.Int32(p), cutlass.Int32(box.rank))
+        rows = to_cute(box.rows[:batch]), to_cute(box.rows_lse[:batch], 4)
+        boxes = to_cute(box.partial[p]), to_cute(box.lse[p], 4)
+        peers = to_cute(box.peers, 8), to_cute(box.lse_peers, 8)
+        self.dcp_args = getattr(self, "dcp_args", {})
+        self.dcp_args[batch] = (mailbox, box.mc(p), ids[1]), (rows, boxes, peers), ids
+
+    def stock_core(self, cache, block_table, lengths, batch, q=1, totals=None):
+        """tokenspeed_mla_decode or trtllm-gen into the BF16 latent (DCP: tokenspeed
+        for every rank's heads over the local shard into the DcpBox rows + LSE; q > 1
+        masks causally in global positions, ``totals``). tokenspeed: a fixed split
+        count per batch (one CTA pair per split filling the SMs); both PDL chained
+        between our kernels."""
+        box, reqs = self.dcp, batch // q
+        query = self.query_scratch[:batch].view(torch.float8_e4m3fn)
+        out, lse, cp = self.latent[:batch], None, {}
+        if box is not None:
+            query, out = box.q[self.parity, :batch].view(query.dtype), box.rows[:batch]
+            lse = box.rows_lse[:batch].view(reqs, q, -1)
+            if q > 1:  # D*H heads: packed (row, head) tiles stream the KV once per tile
+                cp = dict(cp=(box.size, box.rank, box.interleave), totals=totals[:reqs])
+                cp.update(pack=True)
+        query, out = query.view(reqs, q, *query.shape[1:]), out.view(reqs, q, -1, 512)
+        args = cache, block_table[:reqs], lengths[:reqs], self.max_seq
+        if self.cores[batch, q] == "trtllm":
+            trtllm_decode(query, *args, self.stock_scales, out)
+            return
+        split = fill_splits(batch, self.device)
+        stock_decode(query, *args, self.stock_scales, out, lse, split, **cp)
+
+    def run(
+        self, projected, out, *, cache, slots, block_table, lengths, totals=None, q=1
+    ):
+        """``lengths``/``block_table``: this rank's KV (the local shard under DCP);
+        ``totals``: under DCP, the requests' full lengths. ``q`` rows per request (spec
+        verify; ``slots`` per row). Rows of ``projected`` beyond the requests run as
+        empty requests (zero output)."""
+        batch, s = projected.shape[0], stream(self.device)
+        reqs = batch // q
+        block_table, lengths, totals = pad_requests(reqs, block_table, lengths, totals)
+        kv = to_cute(cache.view(torch.uint8))
+        mq, v, query, latent = self.kernels[batch, q]
+        qnorm, kvnorm, qb, uk, uv, qinv, kinv = self.views
+        q_slots, x = to_cute(slots[:batch], 8), to_cute(projected[:, :2112])
+        q_args = (x, qnorm, kvnorm, qb, uk, kv, q_slots, qinv, kinv, self.eps, query)
+        out_view, gate = to_cute(out), to_cute(projected[:, 2112:])
+        lens = to_cute(lengths[:reqs], 4)
+        if self.dcp is not None:  # DcpValue zeroes padded requests (global lengths)
+            publish, ((rows, rows_lse), boxes, peers), ids = self.dcp_args[batch]
+            mq(*q_args, *publish, s)
+            self.stock_core(cache, block_table, lengths, batch, q, totals)
+            v_args = (rows, rows_lse, lens, *boxes, *peers, publish[0], uv, gate)
+            v(*v_args, to_cute(totals[:reqs], 4), out_view, *ids, s)
+            return out
+        mq(*q_args, s)
+        self.stock_core(cache, block_table, lengths, batch, q)
+        v(latent, uv, gate, lens, out_view, s)
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/query.py

@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cluster MLA query RMSNorm, Q_B projection, W_UK absorption and paged FP8 KV insert.
+Static weights load before the PDL wait; current-token data is consumed after it."""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import BFloat16, Float32, Int64
+
+from ..common.primitives import (
+    alloc,
+    at,
+    bf16x2,
+    copy_evict,
+    copy_g2s,
+    copy_s2cluster,
+    float_bits,
+    fma,
+    fma2,
+    fp8_bits,
+    launch,
+    ldmatrix,
+    load_bf16x8,
+    load_volatile,
+    load_words,
+    mma_into,
+    mul2,
+    multicast_words,
+    prefetch_l2,
+    published,
+    query_norm_rsqrt,
+    smem_words,
+    store_fp8_pair,
+    store_shared_words,
+    unpack_bf16,
+    view,
+    wait,
+    warp_sum,
+    zeros,
+)
+
+KP = 1536 + 32  # padded Q_B / x row (elements): 3136 B, conflict-free 16 B fragments
+# The smallest batch on the tiled kernel (per head count):
+# its 16-request tiles beat FP32 FMAs from B4 at 12 heads and from B6 at 6 heads.
+TILED_FROM = {12: 4, 6: 6}
+Q_SENT = 0xFFFFFFFF  # q mailbox sentinel: four FP8 0xFF bytes, never written
+
+
+class MlaQuery:
+    """A cluster splits one head's 192 Q_B rows; each CTA then absorbs its share of the
+    512 latent columns. Small batches (below TILED_FROM) are served by one CTA per
+    cluster part with FP32 FMAs; wider batches tile 16-64 requests per CTA (one wave
+    of 96 CTAs at B64 for both head counts) and use BF16 MMAs (FP32 accumulation).
+
+    ``dcp`` (decode context parallel size; args then end with this parity's local q
+    mailbox, its multicast address and the rank): each CTA also multicasts the query
+    bytes it wrote into head block ``rank`` of every rank's Lamport q mailbox
+    (mla.dcp.DcpBox), then polls its share of the gathered mailbox until no word is
+    the sentinel, so the next kernel reads the gathered query."""
+
+    def __init__(
+        self,
+        heads: int,
+        batch: int,
+        *,
+        q_b_cp_async: bool = False,
+        dcp: int = 0,
+    ):
+        tiled = TILED_FROM[heads]
+        self.heads, self.batch, self.wide = heads, batch, batch >= tiled
+        self.dcp = dcp
+        self.cluster = 16 if heads == 6 and not self.wide else 8
+        self.rows, self.latent = 192 // self.cluster, 512 // self.cluster
+        self.groups, self.p_per_group = 256 // self.latent, self.latent // 2
+        self.tokens = 1 if batch == 2 else min(batch, 8)
+        if self.wide:
+            self.tokens = 16 if batch <= 16 else 32 if batch <= 32 or heads < 8 else 64
+        self.grid = (self.cluster, heads, (batch + self.tokens - 1) // self.tokens)
+        carveout = 100 if batch > 1 else None
+        self.options = dict(cluster=self.cluster, preferred_smem_carveout=carveout)
+        self.q_b_cp_async = q_b_cp_async
+
+    @cute.jit
+    def __call__(self, *args):
+        kernel = self.tiled if self.wide else self.kernel
+        launch(kernel, args, self.grid, 256, **self.options)
+
+    @cute.jit
+    def live(self, token):
+        return cutlass.const_expr(not self.wide) or token < self.batch
+
+    @cute.jit
+    def publish(self, args, tok0, tid):
+        """DCP: this CTA's query bytes (latent columns L*part.. as 16 B, the rope rows
+        of Q_B rows R*part.. as 4 B multicast stores) -> every rank's q mailbox (FP8
+        bytes are never 0xFF: no word is the sentinel)."""
+        output, (box_mc, rank) = args[10], args[12:14]
+        part, head, _ = cute.arch.block_idx()
+        R, L, D = self.rows, self.latent, self.dcp
+        chunks: cutlass.Constexpr = L // 16
+        items: cutlass.Constexpr = chunks + R // 4
+        cute.arch.sync_threads()
+        for i in range(tid, self.tokens * items, 256):
+            t, w = i // items, i % items
+            p = R * part + 4 * (w - chunks)  # Q_B row of a rope word
+            row = ((tok0 + t) * D + rank) * self.heads + head
+            if w < chunks:
+                byte = L * part + 16 * w
+                if self.live(tok0 + t):
+                    words = load_words(
+                        at(output, tok0 + t, head, byte), count=4, mutable=True
+                    )
+                    multicast_words(box_mc + Int64(row) * 576 + byte, words)
+            elif self.live(tok0 + t) & (p >= 128):
+                byte = 512 + p - 128
+                word = load_words(
+                    at(output, tok0 + t, head, byte), count=1, mutable=True
+                )
+                multicast_words(box_mc + Int64(row) * 576 + byte, word, kind="b32")
+
+    @cute.jit
+    def gathered(self, args, tid):
+        """DCP: poll this CTA's share of the gathered mailbox (16 B chunks of the
+        B x DCP*H x 576 bytes) until every rank's bytes have landed."""
+        mailbox = args[11]
+        part, head, tile = cute.arch.block_idx()
+        threads: cutlass.Constexpr = self.grid[0] * self.grid[1] * self.grid[2] * 256
+        chunks: cutlass.Constexpr = self.batch * self.dcp * self.heads * 36
+        first = ((tile * self.heads + head) * self.cluster + part) * 256 + tid
+        last = chunks - 1  # threads past the end re-poll it
+        for i0 in cutlass.range_constexpr(0, -(-chunks // threads), 8):
+            n: cutlass.Constexpr = min(8, -(-chunks // threads) - i0)
+            sources = [
+                mailbox.iterator + min(first + (i0 + i) * threads, last) * 16
+                for i in range(n)
+            ]
+            loads = [load_volatile(source, count=4) for source in sources]
+            for i in cutlass.range_constexpr(n):  # the loads above are in flight
+                words = loads[i]
+                while not published(words, sentinel=Q_SENT):
+                    words = load_volatile(sources[i], count=4)
+
+    @cute.kernel
+    def kernel(self, args):
+        qkv, w_qa, w_kva, w_qb, w_uk, cache, slots, q_scale_inv = args[0:8]
+        k_scale_inv, eps, output = args[8:11]
+        tid, _, _ = cute.arch.thread_idx()
+        part, head, tok0 = cute.arch.block_idx()
+        lane, warp = tid % 32, tid // 32
+        B, R, L, G = self.tokens, self.rows, self.latent, self.groups
+        smem = cutlass.utils.SmemAllocator()
+        qb = alloc(smem, BFloat16, (R, 1536), 128)
+        uk = alloc(smem, BFloat16, (128, L))
+        qa = alloc(smem, BFloat16, 1536)
+        xq = alloc(smem, BFloat16, (B, 1536))
+        qrow = alloc(smem, Float32, (R, B))
+        qn = alloc(smem, Float32, (B, 128))
+        red = alloc(smem, Float32, (G, B, L))
+        ss = alloc(smem, Float32, (8, max(B, 2)))
+        if cutlass.const_expr(not self.q_b_cp_async):
+            barw = smem.allocate(cutlass.Uint64, 8)
+        barq = smem.allocate(cutlass.Uint64, 8)
+        if tid == 0:
+            if cutlass.const_expr(not self.q_b_cp_async):
+                cute.arch.mbarrier_init(barw, 1)
+            cute.arch.mbarrier_init(barq, 1)
+            cute.arch.mbarrier_init_fence()
+            if cutlass.const_expr(not self.q_b_cp_async):
+                cute.arch.mbarrier_arrive_and_expect_tx(barw, R * 1536 * 2)
+                copy_evict(
+                    qb.iterator,
+                    w_qb.iterator + (head * 192 + part * R) * 1536,
+                    size=R * 1536 * 2,
+                    barrier=barw,
+                )
+            cute.arch.mbarrier_arrive_and_expect_tx(barq, B * 128 * 4)
+        cute.arch.cluster_arrive_relaxed()
+        for c in range(tid, 128 * L // 8, 256):
+            p, j = c // (L // 8), c % (L // 8)
+            copy_evict(uk.iterator + p * L + j * 8, at(w_uk, head, p, part * L + j * 8))
+        if tid < 192:
+            copy_evict(qa.iterator + tid * 8, w_qa.iterator + tid * 8)
+        cute.arch.cp_async_commit_group()
+        if cutlass.const_expr(self.q_b_cp_async):
+            for c in range(tid, R * 1536 // 8, 256):
+                dst = qb.iterator + c * 8
+                src = w_qb.iterator + (head * 192 + part * R) * 1536 + c * 8
+                copy_evict(dst, src)
+            cute.arch.cp_async_commit_group()
+        cute.arch.griddepcontrol_wait()
+        cute.arch.griddepcontrol_launch_dependents()
+        qsi = q_scale_inv[0]
+        xv = cute.make_rmem_tensor((B, 8), Float32)
+        for b in cutlass.range_constexpr(B):
+            total = Float32(0)
+            if (tid < 192) & self.live(tok0 + b):
+                values = load_bf16x8(at(qkv, tok0 + b, tid * 8), mutable=True)
+                for i in cutlass.range_constexpr(8):
+                    xv[b, i] = values[i]
+                    total = fma(values[i], values[i], total)
+            total = warp_sum(total)
+            if lane == 0:
+                ss[warp, b] = total
+        cute.arch.cp_async_wait_group(1 if self.q_b_cp_async else 0)
+        cute.arch.sync_threads()
+        if tid < 192:
+            weights = load_bf16x8(qa.iterator + tid * 8, shared=True)
+            for b in cutlass.range_constexpr(B):
+                total = Float32(0)
+                for w in cutlass.range_constexpr(8):
+                    total += ss[w, b]
+                rr = query_norm_rsqrt(total, eps)
+                for i in cutlass.range_constexpr(8):
+                    xq[b, tid * 8 + i] = BFloat16(xv[b, i] * rr * weights[i])
+        cute.arch.sync_threads()
+        if cutlass.const_expr(self.q_b_cp_async):
+            cute.arch.cp_async_wait_group(0)
+            cute.arch.sync_threads()
+        else:
+            wait(barw)
+        for row in range(warp, R, 8):
+            acc = zeros((B, 2), Float32)
+            for chunk in cutlass.range_constexpr(6):
+                k = (chunk * 32 + lane) * 8
+                weights = load_bf16x8(qb.iterator + row * 1536 + k, shared=True)
+                for b in cutlass.range_constexpr(B):
+                    values = load_bf16x8(xq.iterator + b * 1536 + k, shared=True)
+                    for i in cutlass.range_constexpr(8):
+                        acc[b, i % 2] = fma(weights[i], values[i], acc[b, i % 2])
+            for b in cutlass.range_constexpr(B):
+                value = warp_sum(acc[b, 0] + acc[b, 1])
+                if lane == 0:
+                    qrow[row, b] = Float32(BFloat16(value))
+        cute.arch.sync_threads()
+        cute.arch.cluster_wait()
+        nrows = min(R, max(0, 128 - R * part))
+        for i in range(tid, nrows * B * self.cluster, 256):
+            dst, jb = i % self.cluster, i // self.cluster
+            j, b = jb % nrows, jb // nrows
+            cute.arch.store_async_dsmem(
+                qn.iterator + b * 128 + R * part + j, float_bits(qrow[j, b]), barq, dst
+            )
+        for i in range(tid, R * B, 256):
+            j, b = i % R, i // R
+            p = R * part + j
+            if (p >= 128) & self.live(tok0 + b):
+                output[tok0 + b, head, 512 + p - 128] = fp8_bits(qrow[j, b] * qsi)
+        wait(barq)
+        latent_col, group = tid % L, tid // L
+        acc = zeros((B, 2), Float32)
+        for i in cutlass.range_constexpr(self.p_per_group):
+            p = group * self.p_per_group + i
+            w = Float32(uk[p, latent_col])
+            for b in cutlass.range_constexpr(B):
+                acc[b, i % 2] = fma(qn[b, p], w, acc[b, i % 2])
+        for b in cutlass.range_constexpr(B):
+            red[group, b, latent_col] = acc[b, 0] + acc[b, 1]
+        cute.arch.sync_threads()
+        for i in range(tid, L * B, 256):
+            b, ll = i // L, i % L
+            value = Float32(0)
+            for g in cutlass.range_constexpr(G):
+                value += red[g, b, ll]
+            if self.live(tok0 + b):
+                rounded = Float32(BFloat16(value))
+                output[tok0 + b, head, L * part + ll] = fp8_bits(rounded * qsi)
+        if cutlass.const_expr(self.dcp):
+            self.publish(args, tok0, tid)
+        valid = (head == 0) & (part < B) & self.live(tok0 + part)
+        self.insert(args, ss, tok0, part, valid, tid)
+        if cutlass.const_expr(self.dcp):
+            self.gathered(args, tid)
+
+    @cute.jit
+    def insert(self, args, ss, tok0, j, valid, tid):
+        """RMSNorm the current latent KV row of request tok0 + j and write it, with
+        k_pe, to its FP8 cache slot."""
+        qkv, w_kva, (cache, slots), (k_scale_inv, eps) = (
+            args[0],
+            args[2],
+            args[5:7],
+            args[8:10],
+        )
+        lane, warp = tid % 32, tid // 32
+        slot = cutlass.Int64(-1)
+        if valid:
+            slot = slots[tok0 + j]
+        if slot >= 0:
+            ksi = k_scale_inv[0]
+            xx = cute.make_rmem_tensor(8, Float32)
+            total = Float32(0)
+            if tid < 64:
+                values = load_bf16x8(at(qkv, tok0 + j, 1536 + tid * 8), mutable=True)
+                for i in cutlass.range_constexpr(8):
+                    xx[i] = values[i]
+                    total = fma(values[i], values[i], total)
+            total = warp_sum(total)
+            cute.arch.sync_threads()
+            if (lane == 0) & (warp < 2):
+                ss[warp, 0] = total
+            cute.arch.sync_threads()
+            page, offset = slot // cache.shape[1], slot % cache.shape[1]
+            if tid < 64:
+                mean = fma(ss[0, 0] + ss[1, 0], Float32(1 / 512), eps)
+                rr = cute.math.rsqrt(mean, fastmath=True)
+                weights = load_bf16x8(w_kva.iterator + tid * 8)
+                for i in cutlass.range_constexpr(8):
+                    normed = Float32(BFloat16(xx[i] * rr * weights[i]))
+                    cache[page, offset, tid * 8 + i] = fp8_bits(normed * ksi)
+            elif tid < 72:
+                c = tid - 64
+                values = load_bf16x8(at(qkv, tok0 + j, 2048 + c * 8), mutable=True)
+                for i in cutlass.range_constexpr(8):
+                    cache[page, offset, 512 + c * 8 + i] = fp8_bits(values[i] * ksi)
+
+    @cute.kernel
+    def tiled(self, args):
+        """B>8: T requests per CTA. Q_B fragments load to registers before the PDL
+        wait; in-place RMSNorm of the staged rows, split-K Q_B MMAs over the warps,
+        nope rows gathered from the six producing CTAs by bulk DSMEM copies (into the
+        dead x rows), then W_UK MMAs (warp w owns latent columns 8w..8w+7)."""
+        qkv, w_qa, w_qb, w_uk, q_scale_inv = (args[i] for i in (0, 1, 3, 4, 7))
+        eps, output = args[9], args[10]
+        tid, _, _ = cute.arch.thread_idx()
+        part, head, tok0 = cute.arch.block_idx()
+        tok0 *= self.tokens
+        smem = cutlass.utils.SmemAllocator()
+        T, R, L = self.tokens, self.rows, self.latent
+        MT, senders = T // 16, (128 + R - 1) // R
+        lane, warp = tid % 32, tid // 32
+        group, quarter = lane // 4, lane % 4
+        xq = alloc(smem, BFloat16, (T, KP), 128)
+        uk = alloc(smem, BFloat16, (128, L + 8), 128)  # 144 B rows for ldmatrix
+        qa = alloc(smem, BFloat16, 1536)
+        ss = alloc(smem, Float32, (8, 2))
+        barq = smem.allocate(cutlass.Uint64, 8)
+        # After the Q_B MMAs the x rows hold the split-K partials, then our nope rows
+        # and the six producers' (peers write them only after our cluster arrival).
+        red = view(xq.iterator, 0, Float32, (8, T * R // 32, 32), (T * R, 32, 1))
+        qblk = view(xq.iterator, 8 * T * R * 4, BFloat16, (T, R), (R, 1))
+        qn = view(xq.iterator, 9 * T * R * 4, BFloat16, (senders, T, R), (T * R, R, 1))
+        if tid == 0:
+            cute.arch.mbarrier_init(barq, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.mbarrier_arrive_and_expect_tx(barq, senders * T * R * 2)
+        # Split-K Q_B: warp w owns k = 192w..192w+191 (8 values per lane and 32-wide
+        # chunk; the x fragments use the same permutation). L2 now, registers later.
+        rows = head * 192 + part * R + group
+        qb_at = [
+            [
+                at(w_qb, rows + 8 * n, (6 * warp + cc) * 32 + 8 * quarter)
+                for n in range(3)
+            ]
+            for cc in range(6)
+        ]
+        for chunk in qb_at:
+            for address in chunk:
+                prefetch_l2(address)
+        for c in range(tid, 128 * L // 8, 256):
+            p, j = c // (L // 8), c % (L // 8)
+            dst = uk.iterator + p * (L + 8) + j * 8
+            copy_evict(dst, at(w_uk, head, p, part * L + j * 8))
+        if tid < 192:
+            copy_evict(qa.iterator + tid * 8, w_qa.iterator + tid * 8)
+        cute.arch.cp_async_commit_group()
+        cute.arch.griddepcontrol_wait()
+        cute.arch.griddepcontrol_launch_dependents()
+        for c in range(tid, T * 192, 256):
+            t, j = c // 192, c % 192
+            if self.live(tok0 + t):
+                copy_g2s(xq.iterator + t * KP + j * 8, at(qkv, tok0 + t, j * 8))
+        cute.arch.cp_async_commit_group()
+        qsi = q_scale_inv[0]
+        # The KV insert overlaps the staging of x (CTA j of the tile owns request j).
+        j = head * self.cluster + part
+        self.insert(args, ss, tok0, j, (j < T) & self.live(tok0 + j), tid)
+        cute.arch.cp_async_wait_group(0)
+        cute.arch.sync_threads()
+        # Packed FP32 pairs: the same IEEE ops in the same order as the scalar form.
+        gamma = []
+        for i in cutlass.range_constexpr(6):
+            gamma += list(load_bf16x8(qa.iterator + 256 * i + 8 * lane, shared=True))
+        for pair in cutlass.range_constexpr(T // 16):  # requests warp + 8k, two at once
+            xrows = [xq.iterator + (warp + 16 * pair + 8 * h) * KP for h in range(2)]
+            values: list = []
+            rr: list = []
+            for h in cutlass.range_constexpr(2):
+                values.append([])
+                for i in cutlass.range_constexpr(6):
+                    for word in smem_words(xrows[h] + 256 * i + 8 * lane):
+                        values[h] += list(unpack_bf16(word))
+            for h in cutlass.range_constexpr(2):
+                zero = Float32(0), Float32(0)
+                chains = [zero, zero]  # (0, 1), (2, 3): element e feeds chain e % 4
+                for e in cutlass.range_constexpr(0, 48, 2):
+                    x = values[h][e], values[h][e + 1]
+                    chains[e % 4 // 2] = fma2(x, x, chains[e % 4 // 2])
+                (c0, c1), (c2, c3) = chains
+                rr.append(query_norm_rsqrt(warp_sum((c0 + c1) + (c2 + c3)), eps))
+            for h in cutlass.range_constexpr(2):
+                for i in cutlass.range_constexpr(6):
+                    words = []
+                    for e in cutlass.range_constexpr(8 * i, 8 * i + 8, 2):
+                        x = values[h][e], values[h][e + 1]
+                        y = mul2(mul2(x, (rr[h], rr[h])), (gamma[e], gamma[e + 1]))
+                        words.append(bf16x2(*y))
+                    store_shared_words(xrows[h] + 256 * i + 8 * lane, tuple(words))
+        qbr = [[load_words(address, count=4) for address in chunk] for chunk in qb_at]
+        cute.arch.sync_threads()
+        acc = zeros((MT * 3, 4), Float32)
+        for cc in cutlass.range_constexpr(6):
+            k = (6 * warp + cc) * 32 + 8 * quarter
+            z = qbr[cc]
+            for m in cutlass.range_constexpr(MT):
+                x, y = (
+                    load_words(xq.iterator + r * KP + k, count=4, shared=True)
+                    for r in (16 * m + group, 16 * m + group + 8)
+                )
+                for n in cutlass.range_constexpr(3):
+                    for s in cutlass.range_constexpr(2):
+                        a = (x[2 * s], y[2 * s], x[2 * s + 1], y[2 * s + 1])
+                        b = (z[n][2 * s], z[n][2 * s + 1])
+                        mma_into(acc[m * 3 + n, None], a, b, bf16=True)
+        cute.arch.sync_threads()
+        cute.arch.cluster_arrive_relaxed()  # x is consumed: peers may fill qn
+        for i in cutlass.range_constexpr(MT * 12):
+            red[warp, i, lane] = acc[i // 4, i % 4]
+        cute.arch.sync_threads()
+        for e in range(tid, T * R, 256):
+            i, ln = e // 32, e % 32
+            value = Float32(0)
+            for w in cutlass.range_constexpr(8):
+                value += red[w, i, ln]
+            value = Float32(BFloat16(value))
+            tile, c = i // 4, i % 4
+            t = 16 * (tile // 3) + ln // 4 + 8 * (c // 2)
+            r = 8 * (tile % 3) + 2 * (ln % 4) + c % 2
+            p = R * part + r
+            if p < 128:
+                qblk[t, r] = BFloat16(value)
+            elif self.live(tok0 + t):
+                output[tok0 + t, head, 512 + p - 128] = fp8_bits(value * qsi)
+        cute.arch.fence_proxy("async.shared", space="cta")
+        cute.arch.sync_threads()
+        cute.arch.cluster_wait()
+        if (part < senders) & (tid < self.cluster):
+            copy_s2cluster(
+                qn.iterator + part * T * R, qblk.iterator, barq, tid, size=T * R * 2
+            )
+        wait(barq)
+        cute.arch.cluster_arrive_relaxed()  # peers may exit once qblk is delivered
+        b_frags = [
+            ldmatrix(uk.iterator + (32 * kp + lane) * (L + 8) + 8 * warp, trans=True)
+            for kp in range(4)
+        ]
+        for m in cutlass.range_constexpr(MT):
+            out = zeros(4, Float32)
+            t = 16 * m + lane % 8 + 8 * ((lane // 8) % 2)
+            for ks in cutlass.range_constexpr(8):
+                p = 16 * ks + 8 * (lane // 16)
+                a = ldmatrix(qn.iterator + ((p // R) * T + t) * R + p % R)
+                b = b_frags[ks // 2][2 * (ks % 2) : 2 * (ks % 2) + 2]
+                mma_into(out, a, b, bf16=True)
+            for h in cutlass.range_constexpr(2):  # rows group, group + 8: 2 bytes
+                row = 16 * m + group + 8 * h
+                if self.live(tok0 + row):
+                    lo, hi = (Float32(BFloat16(out[2 * h + e])) * qsi for e in range(2))
+                    col = L * part + 8 * warp + 2 * quarter
+                    store_fp8_pair(at(output, tok0 + row, head, col), lo, hi)
+        if cutlass.const_expr(self.dcp):
+            self.publish(args, tok0, tid)
+        cute.arch.cluster_wait()
+        if cutlass.const_expr(self.dcp):
+            self.gathered(args, tid)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/value.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MLA value projection (W_UV) and sigmoid gate of the BF16 attention latent."""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import BFloat16, Float32, Uint32
+
+from ..common.primitives import (
+    alloc,
+    at,
+    bf16x2,
+    copy_g2s,
+    launch,
+    load_words,
+    mma_into,
+    sigmoid,
+    store_words,
+    unpack_bf16,
+    zeros,
+)
+
+
+class MlaValue:
+    """W_UV and the sigmoid gate, [B, H, 512] BF16 latent -> [B, H*128]. One
+    128-thread CTA per (32 outputs, head, 16 rows); warp w owns outputs
+    8w..8w+7 over the whole latent (BF16 MMA, FP32 accumulation). W_UV must have a
+    contiguous latent axis (the absorbed kv_b_proj view). Rows are tokens, ``q`` per
+    request (``lengths`` per request)."""
+
+    def __init__(self, heads: int, batch: int, q: int = 1):
+        self.heads, self.batch, self.q = heads, batch, q
+        self.grid = (4, heads, (batch + 15) // 16)
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, self.grid, 128)
+
+    @cute.kernel
+    def kernel(self, args):
+        latent, w_uv, gate, lengths, output = args
+        rows = alloc(cutlass.utils.SmemAllocator(), BFloat16, (16, 544), 128)
+        tid, _, _ = cute.arch.thread_idx()
+        chunk, head, tile = cute.arch.block_idx()
+        lane, warp = tid % 32, tid // 32
+        group, quarter = lane // 4, lane % 4
+        t0, v = tile * 16, chunk * 32 + warp * 8 + group
+        # B fragments: 8 latent values per 32-wide chunk; A uses the same permutation.
+        wb = [
+            load_words(at(w_uv, head, 32 * c + 8 * quarter, v), count=4)
+            for c in range(16)
+        ]
+        cute.arch.griddepcontrol_wait()
+        for c in range(tid, 16 * 64, 128):
+            r, j = c // 64, c % 64
+            if t0 + r < self.batch:
+                src = at(latent, t0 + r, head, j * 8)
+                copy_g2s(rows.iterator + r * 544 + j * 8, src)
+        cute.arch.cp_async_commit_group()
+        col = head * 128 + chunk * 32 + warp * 8 + 2 * quarter
+        gates = []
+        for r in cutlass.range_constexpr(2):
+            word = Uint32(0)
+            if t0 + group + 8 * r < self.batch:
+                word = load_words(
+                    at(gate, t0 + group + 8 * r, col), count=1, mutable=True
+                )[0]
+            gates.append(unpack_bf16(word))
+        cute.arch.griddepcontrol_launch_dependents()
+        cute.arch.cp_async_wait_group(0)
+        cute.arch.sync_threads()
+        parts = zeros((2, 4), Float32)  # two accumulation chains
+        for c in cutlass.range_constexpr(16):
+            x, y = (
+                load_words(
+                    rows.iterator + r * 544 + 32 * c + 8 * quarter, count=4, shared=True
+                )
+                for r in (group, group + 8)
+            )
+            for s in cutlass.range_constexpr(2):
+                a = (x[2 * s], y[2 * s], x[2 * s + 1], y[2 * s + 1])
+                b = (wb[c][2 * s], wb[c][2 * s + 1])
+                mma_into(parts[s, None], a, b, bf16=True)
+        acc = [parts[0, k] + parts[1, k] for k in range(4)]
+        for r in cutlass.range_constexpr(2):
+            t = t0 + group + 8 * r
+            if t < self.batch:
+                # Padded rows may hold a NaN gate; keep inactive rows zero.
+                lo, hi = Float32(0), Float32(0)
+                if lengths[t // self.q] > 0:
+                    lo, hi = (
+                        Float32(BFloat16(acc[2 * r + e])) * sigmoid(gates[r][e], False)
+                        for e in range(2)
+                    )
+                store_words(at(output, t, col), (bf16x2(lo, hi),))

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Prepared first dense MLP, sharing projection and SiTU decode arithmetic."""
+
+from functools import cache
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import architecture, dynamic_fake, dynamic_view, stream
+from ..common.primitives import launch
+from ..gemm.projection import PreparedProjection
+from .shared import situ
+
+
+class DenseSitu:
+    def __init__(self, hidden, padded):
+        self.hidden, self.padded = hidden, padded
+
+    @cute.jit
+    def __call__(self, *args):
+        grid = (cute.ceil_div(self.padded, 256), args[0].shape[0], 1)
+        launch(self.kernel, args, grid, 256)
+
+    @cute.kernel
+    def kernel(self, args):
+        x, out, beta, linear = args[0:4]
+        tile, row, _ = cute.arch.block_idx()
+        tid, _, _ = cute.arch.thread_idx()
+        col = tile * 256 + tid
+        cute.arch.griddepcontrol_wait()
+        if col < self.padded:
+            value = cutlass.Float32(0)
+            if col < self.hidden:
+                gate = cutlass.Float32(x[row, col])
+                up = cutlass.Float32(x[row, self.hidden + col])
+                value = situ(gate, up, beta, linear, native=True)
+            out[row, col] = cutlass.BFloat16(value)
+        cute.arch.griddepcontrol_launch_dependents()
+
+
+@cache
+def _compile_activation(hidden, padded, arch):
+    args = (dynamic_fake(2), dynamic_fake(2), *map(cutlass.Float32, (4, 25)))
+    kernel = DenseSitu(hidden, padded)
+    return cute.compile(kernel, *args, make_fake_stream(), options=f"--gpu-arch={arch}")
+
+
+class PreparedDenseMLP:
+    """Dense MLP; TP16 pads down K 2112 -> 2176 and the activation zero-fills it."""
+
+    def __init__(self, mlp, batches, reduce=None):
+        gate_up, down = mlp.gate_up_proj.weight, mlp.down_proj.weight
+        self.device, self.reduce = gate_up.device, reduce
+        self.hidden = hidden = gate_up.shape[0] // 2
+        self.padded = padded = (hidden + 127) // 128 * 128
+        linear = -1 if mlp.act_fn.linear_beta is None else mlp.act_fn.linear_beta
+        self.scalars = tuple(map(cutlass.Float32, (mlp.act_fn.beta, linear)))
+        if padded != hidden:
+            down = torch.nn.functional.pad(down.detach(), (0, padded - hidden))
+        self.gates = gate_up.new_empty((64, 2 * hidden))
+        self.hidden_rows = gate_up.new_empty((64, padded))
+        tile = {b: next(n for n in (8, 16, 32, 64) if b <= n) for b in batches}
+        up_tiles = {8: (64, 8, 4, 8), 16: (64, 16, 2, 8), 32: (64, 32, 2, 8)}
+        up_tiles[64] = (64, 64, 2, 5)
+        down_tiles = {8: (128, 8, 1, 6), 16: (128, 16, 1, 6), 32: (128, 32, 1, 5)}
+        down_tiles[64] = (64, 64, 1, 5)
+        policy = lambda tiles: {b: tiles[n] for b, n in tile.items()}
+        self.gate = PreparedProjection(gate_up, policy(up_tiles))
+        self.down = PreparedProjection(down, policy(down_tiles))
+        self.gate.prepare(), self.down.prepare()
+        self.activation = _compile_activation(hidden, padded, architecture(self.device))
+
+    def run(self, x, out):
+        gate, hidden = self.gates[: x.shape[0]], self.hidden_rows[: x.shape[0]]
+        self.gate.run(x, gate)
+        views = (dynamic_view(gate), dynamic_view(hidden))
+        self.activation(*views, *self.scalars, stream(self.device))
+        self.down.run(hidden, out)
+        if self.reduce is not None:
+            reduced = self.reduce(out)
+            if reduced.data_ptr() != out.data_ptr():
+                out.copy_(reduced)
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Persistent routed experts for B<=64 (B>8: grouped units of <=8 tokens per expert)
+over one shared core and model-owned workspace. The entry wait completes the
+predecessor (incl. workspace rearm/counters) before shared state is read; output rows
+keep the caller's top16 order."""
+
+import functools
+
+import cuda.bindings.driver as cuda
+import cutlass as c
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import check_residency, fake, stream, tensor_maps, to_cute
+from ..common.policy import batch_table
+from .experts import (
+    PLAN_BYTES,
+    ExpertsGeometry,
+    ExpertsGroupKernel,
+    ExpertsKernel,
+    ExpertsPlanKernel,
+    fc2_pairs,
+)
+
+
+def tensor(value, align=16):
+    return to_cute(value.reshape(-1), align)
+
+
+def streamed(batch):
+    """B<=8 on the incremental producer: ~1-1.5% at B4-8, neutral at B1-2, which keep
+    the previous producer."""
+    return bool(batch_table("4-8:1").get(batch, 0))
+
+
+def pregrouped(tp, batch):
+    """B>8 experts read the unit table a one-CTA kernel builds after routing, beside
+    the fused shared expert (hidden at B9-32; at B64 its ~2.3 us after routing outlast
+    the shared expert)."""
+    table = batch_table(PREGROUP[tp])
+    return batch > 8 and bool(table.get(batch, 0))
+
+
+PREGROUP = {8: "9-32:1", 16: "9-32:1"}
+# With the shared join skipped (moe.prepared.JOIN), TP8 B16 releases at entry.
+EARLY = {8: "9-16:1", 16: ""}
+
+
+def early_group(tp, batch):
+    """Pregrouped batches whose unit-table kernel releases the experts at its entry."""
+    table = batch_table(EARLY[tp])
+    return pregrouped(tp, batch) and bool(table.get(batch, 0))
+
+
+def stealing(batch):
+    """B>8 experts balance the FC2(B) tail by claims."""
+    return batch > 8 and bool(batch_table("9-64:1").get(batch, 0))
+
+
+@functools.cache
+def compile_experts(tp, experts, batch, grid, device, publish=False, *knobs):
+    """Only code and geometry are cached; no owner tensors or descriptors. ``knobs``:
+    (B<=8 stream, pregrouped, steal)."""
+    defaults = (streamed(batch), False, False)
+    stream, grouped, steal = knobs + defaults[len(knobs) :]
+    capacity = 1 << (batch - 1).bit_length()
+    moe = ExpertsKernel(tp, capacity, experts, grid, publish, stream, grouped, steal)
+    flat = lambda n, dtype=c.Uint8, align=16: fake((n,), (1,), dtype, align)
+    box = (c.Int64(0), c.Int32(0), c.Int32(0)) if publish else ()
+    if publish == "twoshot":  # receive slots, owners' bases, parity x TP + rank
+        box = (fake((3 * tp, 64, 3584 // tp)), fake((tp,), None, c.Int64), c.Int32(0))
+    x, ids = fake((batch, 3584), (3584, 1)), flat(batch * 16, c.Int32)
+    x = flat(PLAN_BYTES) if moe.grouped else x  # grouped: the MXFP8 plan
+    scales = (flat(experts * n) for n in (moe.w13_scale_bytes, moe.w2_scale_bytes))
+    out = flat(batch * 16 * (1 if publish else 3584), c.BFloat16)
+    maps = flat(768 if moe.paired else 256, align=64)
+    tensors = (x, ids, maps, *scales, flat(moe.workspace_bytes), out)
+    scalars = (c.Int32(batch), *map(c.Float32, (100, 0.25, 0.04)), *box)
+    with torch.accelerator.device_index(device):
+        return cute.compile(moe, *tensors, *scalars, make_fake_stream())
+
+
+@functools.cache
+def compile_group(tp, experts, batch, grid, device, early=False):
+    """The pregrouped experts' unit-table kernel: (ids, workspace, batch); ``early``:
+    it releases the experts at entry."""
+    capacity = 1 << (batch - 1).bit_length()
+    kernel = ExpertsGroupKernel(tp, capacity, experts, grid, early=early)
+    flat = lambda n, dtype=c.Uint8: fake((n,), (1,), dtype, 16)
+    args = (flat(batch * 16, c.Int32), flat(kernel.workspace_bytes), c.Int32(batch))
+    with torch.accelerator.device_index(device):
+        return cute.compile(kernel, *args, make_fake_stream())
+
+
+@functools.cache
+def compile_plan(batch, device):
+    """Plan writer for callers without the routing kernel's copy role in front."""
+    args = (fake((batch, 3584)), fake((PLAN_BYTES,), (1,), c.Uint8), make_fake_stream())
+    with torch.accelerator.device_index(device):
+        return cute.compile(ExpertsPlanKernel(batch), *args)
+
+
+def scratch(tp, device):
+    """Model-owned expert buffers: the workspace and the grouped plan (MXFP8 x, written
+    by the routing kernel or ``compile_plan``)."""
+    size, u8 = (ExpertsGeometry(tp).workspace_bytes, torch.uint8)
+    workspace = torch.full((size,), 255, device=device, dtype=u8)
+    plan = torch.empty(PLAN_BYTES, device=device, dtype=u8)
+    return {"workspace": workspace, "plan": plan}
+
+
+class PersistentExperts:
+    def __init__(self, w13, s13, w2, s2, *, beta=4.0, linear_beta=25.0):
+        experts, rows, _ = w13.shape
+        self.tp, self.grid = (8, 136) if rows == 768 else (16, 152)
+        self.device, self.experts = w13.device, experts
+        self.weights = (w13, s13, w2, s2)
+        # TP8: B3-4 on 148 CTAs; B5-8 publish on 128-136 CTAs, leaving SMs for the
+        # one-wave tail; grouped B9-64 on 152.
+        grids = "3-4:148,5-6:128,7:132,9-64:152" if self.tp == 8 else ""
+        self.grids = batch_table(grids)
+        scalars = (beta * linear_beta, 1 / beta, 1 / linear_beta)
+        self.scalars = tuple(map(c.Float32, scalars))
+        # The experts kernel uses maps 0-1 (L2 promotion on w13); the TP8 B1 cluster
+        # experts use maps 2-4 (w13, w2 two-chunk and one-chunk boxes, promoted); TP8
+        # FC2 pairs (B4-64) use map 5 (w2 chunk 2 of 256 rows).
+        w2_tile = (192, 3, 64) if self.tp == 8 else (128, 2, 32)
+        specs = [(w13, rows, 1792, 28, 64, 2, True), (w2, 3584, *w2_tile, 2, False)]
+        if self.tp == 8:
+            specs += [specs[0]] + [(w2, 3584, 192, 3, 64, box, True) for box in (2, 1)]
+        dtype = cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B
+        swizzle = cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B
+        l2 = cuda.CUtensorMapL2promotion
+        l2 = (l2.CU_TENSOR_MAP_L2_PROMOTION_NONE, l2.CU_TENSOR_MAP_L2_PROMOTION_L2_256B)
+        maps = [
+            (weight.data_ptr(), dtype, (128, experts * n, chunks), (row_bytes, stride))
+            + ((128, 128, box), swizzle, l2[promote])
+            for weight, n, row_bytes, chunks, stride, box, promote in specs
+        ]
+        if self.tp == 8:
+            shape = (128, experts * 3584, 3), (192, 64), (128, 256, 1)
+            maps.append((w2.data_ptr(), dtype, *shape, swizzle, l2[False]))
+        self.maps = tensor_maps(maps, self.device)
+        scales = (tensor(s.view(torch.uint8)) for s in (s13, s2))
+        self.views = (to_cute(self.maps[:256], 64), *scales)
+        self.all_maps = to_cute(self.maps, 64)  # TP8 FC2 pairs: maps 0-5
+        self.prepared = {}
+
+    def prepare(
+        self, batch, buffers, out, publish=(), planned=False, twoshot=False, early=False
+    ):
+        """``publish``: (multicast, parity, rank) of the WaveMoeTail latent slot to
+        publish (``twoshot``: the ``common.twoshot`` push box); ``planned``: B>8 calls
+        find the plan written (by the routing kernel), else run writes it from x first;
+        ``early``: the caller runs ``group`` (pregrouped batches) before ``run``."""
+        mode = "twoshot" if twoshot else bool(publish)
+        grouped = pregrouped(self.tp, batch)
+        args: tuple = (batch, self.grid_for(batch), self.device.index, mode)
+        args += (streamed(batch), grouped, stealing(batch))
+        code = compile_experts(self.tp, self.experts, *args)
+        # CTAs wait on other clusters' partials: the whole grid must be resident.
+        check_residency(code, self.device.index, args[1], 384, "Kimi K3 MoE experts")
+        ws, maps = tensor(buffers["workspace"]), self.views[0]
+        plan: tuple | None = None
+        capacity = 1 << (batch - 1).bit_length()
+        if fc2_pairs(self.tp, capacity, batch > 8 or streamed(batch)):
+            maps = self.all_maps
+        group = None
+        if batch > 8:
+            plan = (tensor(buffers["plan"]), compile_plan(batch, self.device.index))
+            plan = plan[:1] if planned else plan
+        if grouped:
+            release = early_group(self.tp, batch)
+            grouping = compile_group(self.tp, self.experts, *args[:3], release)
+            group = grouping, ws, early
+        self.prepared[batch] = (out, tensor(out), ws, code, publish, plan, maps, group)
+
+    def group(self, ids):
+        """Pregrouped batches prepared ``early``: the unit-table kernel (run it on the
+        stream after routing; ``run`` then skips it)."""
+        group = self.prepared[ids.shape[0]][7]
+        if group is not None and group[2]:
+            self.launch_group(ids, group)
+
+    def launch_group(self, ids, group):
+        code, ws, _ = group
+        code(tensor(ids), ws, c.Int32(ids.shape[0]), stream(self.device))
+
+    def grid_for(self, batch):
+        return self.grids.get(batch, self.grid)
+
+    def run(self, x, ids, weights=None):
+        batch = x.shape[0]
+        out, output, workspace, code, box, plan, maps, group = self.prepared[batch]
+        output = tensor(weights) if box else output
+        source = to_cute(x, 16)
+        if plan is not None:
+            if len(plan) > 1:
+                plan[1](source, plan[0], stream(self.device))
+            source = plan[0]
+        if group is not None and not group[2]:
+            self.launch_group(ids, group)
+        args = (source, tensor(ids), maps, *self.views[1:], workspace, output)
+        code(*args, c.Int32(batch), *self.scalars, *box, stream(self.device))
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/front.py

@@ -0,0 +1,538 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MoE front: the weight-stationary [router shard; down shard] GEMM publishes logits
+and the latent shard into a multicast mailbox row; the routing kernel polls all 896
+logits, ranks the top16 (sigmoid + bias, renormalized) and Lamport-copies the latent."""
+
+import functools
+
+import cutlass as c
+import cutlass.cute as cute
+import torch
+from cutlass import Int32, Uint32, Uint64
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import fake, stream, to_cute
+from ..common.mailbox import poll, rearm
+from ..common.policy import batch_table
+from ..common.primitives import (
+    alloc,
+    asm,
+    bits_float,
+    div_rn,
+    float_bits,
+    launch,
+    redux,
+    sigmoid,
+    store_words,
+    view,
+)
+from ..gemm.projection import PreparedProjection
+from ..gemm.ws_gemm import WsGemm
+from .experts import PLAN_BYTES, plan_chunk
+from .shared import compile_down, compile_shared, situ, weight_maps
+from .tail_runtime import CAP
+
+
+@asm(Int32)
+def shared_add(address, value):
+    code = "atom.shared.add.u32 $0,[$1],$2;"
+    return code, "=r,r,r,~{memory}", (address, Int32(value))
+
+
+@cute.jit
+def ordered(score):
+    """Monotone uint32 key of a float; NaN maps to 0 (below every number)."""
+    bits = Uint32(float_bits(score))
+    key = ~bits if (bits & Uint32(0x80000000)) != 0 else bits | Uint32(0x80000000)
+    return key if score == score else Uint32(0)
+
+
+class RouteSitu:
+    def __init__(
+        self, batch, shared, renorm, early, roles, polled, plan=False
+    ):  # fmt: skip
+        # B>8: few registers, so routing CTAs co-reside with the front and the fused
+        # shared expert, and polled logits skip the front's completion wait.
+        self.lean = batch > 8
+        self.batch, self.shared, self.experts, self.renorm = batch, shared, 896, renorm
+        # early: release the persistent experts at entry (safe once down drains first);
+        # polled: logits arrive as FP32 words in mailbox columns [3584, 5376) (Lamport);
+        # plan: the copy role also writes the latent as MXFP8 into the experts' plan.
+        self.early, self.polled, self.plan = early, polled, plan
+        self.roles = tuple(r for r, on in zip(("route", "copy", "situ"), roles) if on)
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, len(self.roles) * self.batch, 256)
+
+    @cute.kernel
+    def kernel(self, args):
+        logits, bias, ids, weights, mailbox, latent, gate_up, hidden = args[0:8]
+        scale, beta, linear_beta = args[8:11]
+        plan = args[11] if self.plan else None
+        tid, _, _ = cute.arch.thread_idx()
+        cta, _, _ = cute.arch.block_idx()
+        token, role = cta % self.batch, cta // self.batch
+        if c.const_expr(self.early):
+            cute.arch.griddepcontrol_launch_dependents()
+        for r in c.range_constexpr(len(self.roles)):
+            if role == r:
+                if c.const_expr(self.roles[r] == "route"):
+                    self.route(logits, bias, ids, weights, scale, token, tid, mailbox)
+                elif c.const_expr(self.roles[r] == "copy"):
+                    # Lamport: no dependency wait; the sentinel poll is the barrier.
+                    base = mailbox.iterator + token * mailbox.shape[2] + tid * 8
+                    if tid < 192:
+                        self.copy(base, latent, token, tid, 2, plan)
+                    else:
+                        self.copy(base, latent, token, tid, 1, plan)
+                else:
+                    cute.arch.griddepcontrol_wait()
+                    rounds = (self.shared + 255) // 256  # issue every load first
+                    pair = cute.make_rmem_tensor((rounds, 2), c.Float32)
+                    for i in c.range_constexpr(rounds):
+                        if tid + 256 * i < self.shared:
+                            pair[i, 0] = gate_up[token, tid + 256 * i]
+                            pair[i, 1] = gate_up[token, self.shared + tid + 256 * i]
+                    for i in c.range_constexpr(rounds):
+                        if tid + 256 * i < self.shared:
+                            gate = c.Float32(c.BFloat16(pair[i, 0]))
+                            up = c.Float32(c.BFloat16(pair[i, 1]))
+                            value = situ(gate, up, beta, linear_beta)
+                            hidden[token, tid + 256 * i] = c.BFloat16(value)
+        # Late release: the persistent experts must not take SMs from the down GEMM.
+        if c.const_expr(not self.early):
+            cute.arch.griddepcontrol_launch_dependents()
+
+    @cute.jit
+    def select(self, key, threshold, buffers):
+        """B>8 candidates (keys at or above ``threshold``) append to ``dense`` at a
+        per-warp offset from the shared count, then rank themselves: 8 threads per
+        candidate (4 past 32 candidates) each count their share of the candidates above
+        it (key, then lower expert). Ranks below 16 store the expert and its sigmoid."""
+        dense, sig, chosen, picked, found = buffers
+        tid, _, _ = cute.arch.thread_idx()
+        lane = tid % 32
+        packed = cute.make_rmem_tensor(4, Uint64)
+        ballots = cute.make_rmem_tensor(4, Uint32)
+        total = Int32(0)
+        for i in c.range_constexpr(4):  # packed 0: not a candidate
+            expert = self.expert(tid, i)
+            hit = (expert < self.experts) & (key[i] >= threshold)
+            ballots[i] = cute.arch.vote_ballot_sync(hit)
+            packed[i] = Uint64(0)
+            if hit:
+                packed[i] = (Uint64(key[i]) << 32) | Uint64(~Uint32(expert))
+            total += cute.arch.popc(ballots[i])
+        offset = Int32(0)
+        if (lane == 0) & (total > 0):
+            offset = shared_add(found.iterator, total)
+        offset = cute.arch.shuffle_sync(offset, 0)
+        below = (Uint32(1) << lane) - Uint32(1)
+        for i in c.range_constexpr(4):
+            if packed[i] != Uint64(0):
+                dense[offset + cute.arch.popc(ballots[i] & below)] = packed[i]
+            offset += cute.arch.popc(ballots[i])
+        cute.arch.sync_threads()
+        n = found[0]
+        if n <= 32:  # usually: 8 threads per candidate, 4 slots each
+            part, index = tid % 8, tid // 8
+            mine, rank = Uint64(0), Int32(0)
+            if index < n:
+                mine = Uint64(dense[index])
+            for k in c.range_constexpr(4):  # slots past n do not count
+                j = part + 8 * k
+                rank += Int32((j < n) & (Uint64(dense[j]) > mine))
+            for shift in c.range_constexpr(3):
+                rank += cute.arch.shuffle_sync_bfly(rank, 1 << shift)
+            if (part == 0) & (index < n) & (rank < 16):
+                expert = Int32(~Uint32(mine))
+                chosen[rank], picked[rank] = expert, sig[expert]
+        else:
+            part = tid % 4
+            for first in c.range(0, n, 64, unroll=1):
+                index = first + tid // 4
+                mine, rank = Uint64(0), Int32(0)
+                if index < n:
+                    mine = Uint64(dense[index])
+                if n <= 64:  # every load issues at once; slots past n do not count
+                    for k in c.range_constexpr(16):
+                        j = part + 4 * k
+                        rank += Int32((j < n) & (Uint64(dense[j]) > mine))
+                else:
+                    for j in c.range(part, n, 4, unroll=2):
+                        rank += Int32(Uint64(dense[j]) > mine)
+                rank += cute.arch.shuffle_sync_bfly(rank, 1)
+                rank += cute.arch.shuffle_sync_bfly(rank, 2)
+                if (part == 0) & (index < n) & (rank < 16):
+                    expert = Int32(~Uint32(mine))
+                    chosen[rank], picked[rank] = expert, sig[expert]
+        cute.arch.sync_threads()
+
+    @cute.jit
+    def publish(self, chosen, picked, ids, weights, scale, token):
+        """Warp 0: the 16 ids by rank and their sigmoids renormalized (lane r holds
+        rank r; the butterfly sums in the same order for every lane)."""
+        tid, _, _ = cute.arch.thread_idx()
+        lane = tid % 32
+        if tid < 32:
+            value = c.Float32(0)
+            if lane < 16:
+                ids[token, lane] = chosen[lane]
+                value = picked[lane]
+            total = value
+            for shift in c.range_constexpr(4):
+                total += cute.arch.shuffle_sync_bfly(total, 1 << shift)
+            if lane < 16:
+                value = value * scale
+                if c.const_expr(self.renorm):
+                    value = div_rn(value, total + c.Float32(1e-20))
+                weights[token, lane] = c.BFloat16(value)
+
+    def expert(self, tid, i):
+        return tid * 4 + i if self.polled else tid + i * 256
+
+    @cute.jit
+    def copy(self, base, latent, token, tid, count: c.Constexpr[int], plan=None):
+        got = poll(base, 2048, count)
+        for i in c.range_constexpr(count):
+            dst = latent.iterator + token * latent.stride[0] + (tid + 256 * i) * 8
+            words = [got[i, k] for k in range(4)]
+            store_words(dst, words)
+            rearm(base + i * 2048)
+            if c.const_expr(plan is not None):
+                plan_chunk(plan, token, tid + 256 * i, words)
+
+    @cute.jit
+    def route(self, logits, bias, ids, weights, scale, token, tid, mailbox):
+        """Exact top16 by threads [0, 256): candidates at or above the 16th largest
+        lane-group maximum are compacted and ranked by an unrolled 64-way compare."""
+        warp, lane = tid // 32, tid % 32
+        smem = c.utils.SmemAllocator()
+        km = alloc(smem, Uint32, (8, 32))
+        cand, dense = (alloc(smem, Uint64, shape) for shape in ((8, 128), 64))
+        sig = alloc(smem, c.Float32, 1024)
+        counts = alloc(smem, Int32, 8)
+        cutoff = alloc(smem, Uint32, 1)
+        chosen = alloc(smem, Int32, 16)
+        if c.const_expr(self.lean):  # candidate count; the chosen experts' sigmoids
+            found, picked = alloc(smem, Int32, 1), alloc(smem, c.Float32, 16)
+        # Thread t owns experts t + 256 i, or 4 t + i when polling 16-byte chunks.
+        offsets = cute.make_rmem_tensor(4, c.Float32)
+        for i in c.range_constexpr(4):  # bias is a weight: load before the wait
+            offsets[i] = c.Float32(0)
+            if self.expert(tid, i) < self.experts:
+                offsets[i] = bias[self.expert(tid, i)]
+        raw = cute.make_rmem_tensor(4, c.Float32)
+        key, top = cute.make_rmem_tensor(4, Uint32), Uint32(0)
+        for i in c.range_constexpr(4):
+            key[i] = Uint32(0)
+        if c.const_expr(self.polled):
+            if tid < 224:
+                cell = mailbox.iterator + token * mailbox.shape[2] + 3584 + tid * 8
+                got = poll(cell, 0, 1)
+                rearm(cell)
+                for i in c.range_constexpr(4):
+                    raw[i] = bits_float(got[0, i])
+                if c.const_expr(self.lean):  # all four experts: one block, no branches
+                    scores = [sigmoid(raw[i]) for i in range(4)]
+                    for i in c.range_constexpr(4):
+                        sig[4 * tid + i] = scores[i]
+                        key[i] = ordered(scores[i] + offsets[i])
+            if c.const_expr(not self.lean):
+                cute.arch.griddepcontrol_wait()
+        else:
+            cute.arch.griddepcontrol_wait()
+            for i in c.range_constexpr(4):
+                if self.expert(tid, i) < self.experts:
+                    raw[i] = logits[token, self.expert(tid, i)]
+        for i in c.range_constexpr(4):
+            expert = self.expert(tid, i)
+            if c.const_expr(not (self.lean and self.polled)):  # noqa: SIM102
+                if expert < self.experts:
+                    score = sigmoid(raw[i])
+                    sig[expert] = score
+                    key[i] = ordered(score + offsets[i])
+            top = max(top, key[i])
+        km[warp, lane] = top
+        if c.const_expr(self.lean):  # noqa: SIM102
+            if tid == 0:
+                found[0] = Int32(0)
+        cute.arch.sync_threads()
+        if warp == 0:
+            group = km[0, lane]
+            for w in c.range_constexpr(1, 8):
+                group = max(group, km[w, lane])
+            km[0, lane] = group
+            cute.arch.sync_warp()
+            rank = Int32(0)
+            for j in c.range_constexpr(32):
+                rank += Int32(km[0, j] > group)
+            low = redux(group if rank < 16 else Uint32(0xFFFFFFFF))
+            if lane == 0:
+                cutoff[0] = low
+        cute.arch.sync_threads()
+        if c.const_expr(self.lean):
+            dense = view(cand.iterator, 0, Uint64, (1024,))
+            self.select(key, cutoff[0], (dense, sig, chosen, picked, found))
+            self.publish(chosen, picked, ids, weights, scale, token)
+            return
+        threshold, offset = cutoff[0], Int32(0)
+        for i in c.range_constexpr(4):
+            expert = self.expert(tid, i)
+            hit = (expert < self.experts) & (key[i] >= threshold)
+            ballot = cute.arch.vote_ballot_sync(hit)
+            if hit:
+                below = cute.arch.popc(ballot & ((Uint32(1) << lane) - Uint32(1)))
+                packed = (Uint64(key[i]) << 32) | Uint64(~Uint32(expert))
+                cand[warp, offset + below] = packed
+            offset += cute.arch.popc(ballot)
+        if lane == 0:
+            counts[warp] = offset
+        cute.arch.sync_threads()
+        if warp == 0:
+            n = Int32(0)
+            for v in c.range_constexpr(8):
+                n += counts[v]
+            for half in c.range_constexpr(2):  # compact the regions densely
+                index = lane + 32 * half
+                rest, region = Int32(index), Int32(0)
+                for v in c.range_constexpr(8):
+                    if (region == v) & (rest >= counts[v]):
+                        rest -= counts[v]
+                        region += 1
+                value = Uint64(0)
+                if index < n:
+                    value = cand[region, rest]
+                dense[index] = value
+            cute.arch.sync_warp()
+            if n <= 64:
+                for half in c.range_constexpr(2):
+                    mine, rank = dense[lane + 32 * half], Int32(0)
+                    for j in c.range_constexpr(64):
+                        rank += Int32(dense[j] > mine)
+                    if (lane + 32 * half < n) & (rank < 16):
+                        chosen[rank] = Int32(~Uint32(mine))
+            else:  # rare: rank every candidate against all regions
+                for q in range(lane, n, 32):
+                    rest, region = Int32(q), Int32(0)
+                    for v in c.range_constexpr(8):
+                        if (region == v) & (rest >= counts[v]):
+                            rest -= counts[v]
+                            region += 1
+                    mine, rank = cand[region, rest], Int32(0)
+                    for v in c.range_constexpr(8):
+                        for j in range(counts[v]):
+                            rank += Int32(cand[v, j] > mine)
+                    if rank < 16:
+                        chosen[rank] = Int32(~Uint32(mine))
+            cute.arch.sync_warp()
+            expert, value = chosen[lane % 16], c.Float32(0)
+            if lane < 16:
+                ids[token, lane] = expert
+                value = sig[expert]
+            total = value
+            for shift in c.range_constexpr(4):
+                total += cute.arch.shuffle_sync_bfly(total, 1 << shift)
+            if lane < 16:
+                value = value * scale
+                if c.const_expr(self.renorm):
+                    value = div_rn(value, total + c.Float32(1e-20))
+                weights[token, lane] = c.BFloat16(value)
+
+
+# Front GEMM geometry "ks.kb.tg.rt.ch" per batch, by packed rows (TP8 560, TP16 288),
+# from GB300 sweeps. ks: cluster K-split; kb: 32-wide
+# K blocks per warp, so a row tile takes 224 / (ks * kb) warps (at most 16); tg: tokens
+# per group (0: the single-pass path); rt: 16-row tiles per CTA; ch: accumulator chains
+# per 8-token tile (0: 4 / tiles per group).
+FRONT = {
+    560: "1-4:4.14.0.1.0,5-32:16.7.16.5.2,33-64:16.14.16.5.2",
+    288: "1-4:4.14.0.1.0,5-64:16.7.16.3.2",
+}
+WIDTH = 3584 + 2 * 896  # front mailbox row: BF16 latent, then FP32 logit words
+# Batches whose shared expert is the fused kernel (``SharedUmma``), not the split-K
+# chain; of those, split batches run gate/up + SiTU beside the front and the down rows
+# (``SharedDown``) on SPLIT SMs beside the routed experts (0: fused).
+FUSED, SPLIT = "9-64:1", "9-12:6,13-24:4,25-64:2"
+# Per-TP overrides of SPLIT entries, (q = 1, verify steps): on real prompts (GB300
+# SPEED-Bench 8k mixed) SharedDown on 2 SMs (~123 us at M32) outlasts the experts (TP8
+# q=1 B32: 3 SMs, -5.4% step). Verify steps (q > 1) share experts among a request's q
+# tokens: TP8 DL7 M16/32/64 -5.9/-12.9/-1.1%, TP16 M32 -5.9%.
+SPLITS = {8: ("25-32:3", "13-16:6,25-32:4,33-64:3"), 16: ("", "25-32:3")}
+# Split batches whose front GEMM rows join the gate/up kernel as extra M-tiles (one
+# kernel instead of two that slow each other ~2x side by side).
+MERGED = {8: "41-64:1", 16: "33-64:1"}
+# Shared-expert split-K tactics (M tile, N tile, split, stages) by width (TP8 768, TP16
+# 384): gate_up, then down.
+SHARED_TACTICS = {
+    768: (
+        "1-8:128.8.4.6,9-16:64.16.4.10,17-32:64.32.4.8,33-64:64.64.4.5",
+        "1-8:128.8.1.6,9-16:64.16.1.6,17-32:64.32.1.6,33-64:64.64.1.6",
+    ),
+    384: (
+        "1-16:64.8.4.8,17-32:64.16.4.10,33-64:64.32.4.8",
+        "1-8:128.8.1.3,9-16:64.16.1.3,17-32:64.32.1.3,33-64:64.64.1.3",
+    ),
+}
+
+
+@functools.cache
+def front_table(n):
+    return batch_table(FRONT.get(n, ""))
+
+
+def front_ks(batch, n):
+    return front_table(n).get(batch, (4,))[0]
+
+
+def front_geometry(batch, n, ks, kb=None, tg=None, rt=None, ch=None):
+    """(rt, kq, staged, tg, ch) of the front GEMM over ``n`` rows. Unset arguments come
+    from the table when ``ks`` is its K-split, else from the single pass; token groups
+    whose x slice does not fit in SMEM fall back to the single pass."""
+    entry = front_table(n).get(batch, ())
+    table = entry[1:] if entry[:1] == (ks,) else (14, 0, 1, 0)
+    kb, tg, rt, ch = (t if v is None else v for v, t in zip((kb, tg, rt, ch), table))
+    kq = min(16, 224 // (ks * kb))
+    kb, nt = 224 // (ks * kq), -(-batch // 8)
+    staged = max(kb - 8, 0) if kq == 16 else 0
+    used = (
+        rt * kq * (nt * 544 + max(staged, 1) * 1024) + -(-nt * rt // ks) * ks * kq * 512
+    )
+    if nt * 8 * (7168 // ks * 2 + 64) > 220 * 1024 - used:
+        tg = 0
+    return rt, kq, staged, tg, ch
+
+
+@functools.cache
+def compile_route(batch, shared, renorm, roles, device, polled=False, plan=False):
+    kernel = RouteSitu(batch, shared, renorm, batch <= 2, roles, polled, plan)
+    f32 = c.Float32
+    args: tuple = (fake((batch, 896), None, f32), fake((896,), None, f32))
+    args += (fake((batch, 16), None, c.Int32), fake((batch, 16)), fake((1, 64, WIDTH)))
+    args += (fake((batch, 3584)), fake((batch, 2 * shared), None, f32))
+    args += (fake((batch, shared)), *(f32(v) for v in (1, 4, 25)))
+    args += (fake((PLAN_BYTES,), None, c.Uint8),) * plan + (make_fake_stream(),)
+    with torch.accelerator.device_index(device):
+        return cute.compile(kernel, *args)
+
+
+@functools.cache
+def compile_publish(n, split, valid, batch, device, ks, geometry=None):
+    """Front GEMM that only publishes; routing then runs in the polled kernel."""
+    rt, kq, staged, tg, ch = geometry or front_geometry(batch, n, ks)
+    kernel = WsGemm(n, 7168, batch, rt, kq, staged, split, valid, ks, tg, ch)
+    args: tuple = (fake((batch, 7168)), fake((n, 7168)), fake((1, 64, WIDTH)))
+    args += (c.Int64(0), c.Int32(0), c.Int32(0), fake((batch, 896), None, c.Float32))
+    with torch.accelerator.device_index(device):
+        return cute.compile(kernel, *args, make_fake_stream())
+
+
+class SharedFront:
+    def __init__(
+        self, moe, packed, box, storage, batches, beta, linear_beta, plan=None,
+        counters=None, q=1,
+    ):  # fmt: skip
+        """``plan``: the experts' plan buffer, written by the routing kernel at B>8;
+        ``counters``: the fused shared expert's model-owned barrier counters; ``q``: the
+        deployment's verify length (its SPLITS entries)."""
+        experts = moe.shared_experts
+        gate_up, down = experts.gate_up_proj.weight, experts.down_proj.weight
+        self.shared = shared = gate_up.shape[0] // 2
+        tp = 6144 // shared
+        shard, valid = 3584 // tp, 896 // tp
+        self.device, self.packed, self.storage = packed.device, to_cute(packed), storage
+        gu, dn = map(batch_table, SHARED_TACTICS[shared])
+        self.gate_up = PreparedProjection(
+            gate_up, {b: gu[b] for b in batches}, output_dtype=torch.float32
+        ).prepare()
+        self.down = PreparedProjection(down, {b: dn[b] for b in batches}).prepare()
+        scale, renorm = moe.routed_scaling_factor, moe.moe_renormalize
+        self.scalars = tuple(map(c.Float32, (scale, beta, linear_beta)))
+        self.box: tuple = (to_cute(box.buffer), c.Int64(box.mc))
+        self.box += (c.Int32(box.rank * valid), c.Int32(box.rank * shard))
+        bias = to_cute(moe.gate.e_score_correction_bias)
+        dev, n = packed.device.index, packed.shape[0]
+        fused = batch_table(FUSED)
+        fused = {b for b in batches if b > 8 and counters is not None and fused.get(b)}
+        verify = q > 1
+        split = batch_table(f"{SPLIT},{SPLITS[tp][verify]}")
+        self.split = {b: split[b] for b in fused if split.get(b)}  # down SMs
+        if fused:  # TMA maps and counters of the fused kernel
+            self.maps = weight_maps(gate_up, down)
+            self.fused = to_cute(self.maps, 64), to_cute(counters)
+        merged = batch_table(MERGED[tp])
+        self.merged = {b for b in self.split if merged.get(b)}
+        if self.merged:  # gate/up, down and front maps; unused RS arguments
+            self.merged_maps = weight_maps(gate_up, down, front=packed)
+            # Kept alive here: a CuTe view does not hold its tensor.
+            self.unused = torch.zeros(tp, dtype=torch.int64, device=packed.device)
+            self.merged_args = to_cute(self.merged_maps, 64), to_cute(self.unused)
+            geometry = (n, n - shard, valid, box.buffer.shape[-1])
+        self.plans: dict[int, tuple] = {}
+        for b in batches:
+            views = tuple(t[:b] for t in storage)
+            logits, gate, ids, weights, h, _, latent = map(to_cute, views)
+            front = compile_publish(n, n - shard, valid, b, dev, front_ks(b, n))
+            planned = plan is not None and b > 8
+            roles = (True, True, False)
+            route = compile_route(b, shared, renorm, roles, dev, True, planned)
+            situ = compile_route(b, shared, renorm, (False, False, True), dev)
+            args = (logits, bias, ids, weights, self.box[0], latent, gate, h)
+            extra = (to_cute(plan),) if planned else ()
+            code = None
+            if b in self.merged:
+                gate = compile_shared(tp, b, dev, True, geometry)
+                code = gate, compile_down(tp, b, dev, self.split[b])
+            elif b in self.split:
+                gate = compile_shared(tp, b, dev, gate=True)
+                code = gate, compile_down(tp, b, dev, self.split[b])
+            elif b in fused:
+                code = compile_shared(tp, b, dev)
+            self.plans[b] = views, front, route, situ, args + self.scalars, extra, code
+
+    def route(self, x):
+        """Router logits + published down shard, then top16 and the latent copy (and
+        at B>8 the experts' MXFP8 plan). Merged batches publish them from the gate/up
+        kernel, which also leaves h for ``shared_expert``'s down rows."""
+        views, front, route, _, args, extra, code = self.plans[x.shape[0]]
+        s = stream(self.device)
+        if x.shape[0] in self.merged:
+            maps, unused = self.merged_args
+            work = (to_cute(x), maps, args[7], self.fused[1], unused)
+            work += (Int32(0), Int32(0), Int32(CAP), *self.scalars[1:], *self.box[1:])
+            code[0](*work, s)
+        else:
+            front(to_cute(x), self.packed, *self.box, args[0], s)
+        route(*args, *extra, s)
+        return views[2], views[3], views[6]
+
+    def merged_batch(self, batch):
+        return batch in self.merged
+
+    def planned(self, batch):
+        return bool(self.plans[batch][5])
+
+    def shared_expert(self, x, tail):
+        """The shared expert's rows into the tail's RS mailboxes (``tail.push``); the
+        fused kernel peer-stores them itself and leaves the returned rows unwritten.
+        Returns (rows, the split's down launch to run beside the experts, or None)."""
+        views, _, _, situ, args, _, code = self.plans[x.shape[0]]
+        s = stream(self.device)
+        if code is not None:
+            maps, counters = self.fused
+            mailbox = (tail.peers, *tail.ring, Int32(CAP))
+            front = (to_cute(x), maps, args[7], counters, *mailbox, *self.scalars[1:])
+            if x.shape[0] not in self.split:
+                code(*front, s)
+                return views[5], None
+            gate, down = code
+            if x.shape[0] not in self.merged:  # merged: h comes from ``route``
+                gate(*front, s)
+            return views[5], lambda: down(args[7], maps, *mailbox, stream(self.device))
+        self.gate_up.run(x, views[1])
+        situ(*args, s)
+        shared = self.down.run(views[4], views[5])
+        tail.push(shared)
+        return shared, None

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Prepared MoE provider: one model owner serializes its scratch, no tensor pointers
+are rebound at execution, and the side stream joins before the tail."""
+
+import torch
+
+from ..common.policy import batch_table
+from .experts_runtime import scratch as expert_scratch
+from .shared import scratch as shared_scratch
+
+
+class MoEScratch:
+    def __init__(self, tp, device, batches):
+        self.tp, self.device, self.batches = tp, torch.device(device), tuple(batches)
+        cap, h = max(batches), 6144 // tp
+        with torch.accelerator.device_index(self.device.index):
+            self.stream = torch.cuda.Stream(device=self.device)
+            self.begin, self.shared_done = torch.cuda.Event(), torch.cuda.Event()
+            self.gate_done = torch.cuda.Event()  # the split shared expert's front half
+            f32, i32, bf16 = torch.float32, torch.int32, torch.bfloat16
+            empty = lambda s, t=bf16: torch.empty(s, dtype=t, device=self.device)
+            self.latent = empty((cap, 3584))
+            self.front: tuple = (empty((cap, 896), f32), empty((cap, 2 * h), f32))
+            self.front += (empty((cap, 16), i32), empty((cap, 16)), empty((cap, h)))
+            self.front += (empty((cap, 7168)), self.latent)
+            self.expanded = empty((cap * 16, 3584))
+            self.experts = expert_scratch(tp, self.device)
+            # The experts' plan (MXFP8 x), written by the routing kernel at B>8.
+            self.plan = self.experts["plan"]
+            self.shared = shared_scratch(self.device)  # fused shared-expert counters
+        self.views = {}
+        for b in batches:
+            self.views[b] = (self.latent[:b], self.expanded[: b * 16], self.experts)
+
+
+# TP8 B16: the experts start beside the split gate/up half (+1% E2E).
+JOIN = {8: "9-16:0", 16: ""}
+
+
+class PreparedMoE:
+    def __init__(self, front, experts, tail, state, *, prefetch=None, layer_idx=None):
+        self.front, self.experts, self.tail, self.state = front, experts, tail, state
+        self.prefetch, self.layer_idx = prefetch, layer_idx
+        self.small_batch = None  # optional small-batch chain
+        # Split batches whose routed experts launch beside the shared expert's
+        # gate/up half instead of after it.
+        self.join = batch_table(JOIN[state.tp])
+        self.sms = torch.cuda.get_device_properties(state.device).multi_processor_count
+        for b, sms in front.split.items():  # the split's down CTAs take the SMs left
+            experts.grids[b] = min(experts.grid_for(b), self.sms - sms)
+        for b in state.batches:
+            _, expanded, buffers = state.views[b]
+            publish, planned = tail.publish.get(b, ()), front.planned(b)
+            twoshot = b in tail.twoshots
+            experts.prepare(b, buffers, expanded, publish, planned, twoshot, True)
+
+    def run(self, x, out, *, consume=None):
+        state, b = self.state, x.shape[0]
+        if self.small_batch is not None and b in self.small_batch.batches:
+            return self.small_batch.run(out, consume)
+        main = torch.cuda.current_stream(state.device)
+        # The shared expert runs on the side stream beside the routed experts, or, when
+        # the persistent expert grid spans every SM, beside the front only. Split: its
+        # gate/up half beside the front, its down rows beside the experts.
+        before = self.experts.grid_for(b) >= self.sms
+        if self.front.merged_batch(b):
+            # One kernel publishes the routing logits and leaves h; the down rows then
+            # run beside the experts.
+            ids, weights, latent = self.front.route(x)
+            state.begin.record(main)
+            with torch.cuda.stream(state.stream):
+                state.stream.wait_event(state.begin)
+                shared, beside = self.front.shared_expert(x, self.tail)
+                beside()
+                state.shared_done.record(state.stream)
+            self.experts.group(ids)
+            return self.finish(b, out, latent, ids, weights, shared, consume)
+        state.begin.record(main)
+        with torch.cuda.stream(state.stream):
+            state.stream.wait_event(state.begin)
+            shared, beside = self.front.shared_expert(x, self.tail)
+            if beside is not None:
+                state.gate_done.record(state.stream)
+                beside()
+            state.shared_done.record(state.stream)
+        ids, weights, latent = self.front.route(x)
+        self.experts.group(ids)  # pregrouped unit table, beside the shared expert
+        if before or (beside is not None and self.join.get(b, 1)):
+            main.wait_event(state.shared_done if before else state.gate_done)
+        return self.finish(b, out, latent, ids, weights, shared, consume, before)
+
+    def finish(self, b, out, latent, ids, weights, shared, consume, before=False):
+        """The routed experts, the next layer's prefetch fork, then the tail once the
+        shared expert's side stream is done."""
+        main = torch.cuda.current_stream(self.state.device)
+        expanded = self.experts.run(latent, ids, weights)
+        if self.prefetch is not None:
+            self.prefetch.fork("next_attention", self.layer_idx, b)
+        if not before:
+            main.wait_event(self.state.shared_done)
+        return self.tail.run(expanded, weights, shared, out, consume=consume)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/routing.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Top16 routing (threshold, regional candidates and exact composite ranking) and the
+helpers the TP8 and TP16 small-batch fronts share."""
+
+from typing import Any
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Uint32, Uint64
+
+from ..common.primitives import (
+    bits_float,
+    copy_evict,
+    float_bits,
+    gate_bf16_bits,
+    mailbox_load,
+    mailbox_store,
+    no_neg_zero,
+    published,
+    redux,
+    store_shared_words,
+    view,
+    zeros,
+)
+
+
+@cute.jit
+def bf16_word(lo, hi):
+    return no_neg_zero(Uint32(gate_bf16_bits(lo)) | (Uint32(gate_bf16_bits(hi)) << 16))
+
+
+@cute.jit
+def poll_words(words, pointer, valid):
+    """Copy four mailbox words into ``words``; clear ``valid`` if any is unpublished."""
+    values = mailbox_load(pointer, count=4)
+    for j in cutlass.range_constexpr(4):
+        words[j] = values[j]
+        valid = valid & (values[j] != Uint32(0x80000000))
+    return valid
+
+
+@cute.jit
+def gather_x(xmoe, par, cta, dst, first_chunk, active, tokens=1, h2=None):
+    """Poll this CTA's x_moe rows (and the TP16 H2 slices if ``h2 = (warp, lane,
+    pairs)``) until published, then stage the row chunks at ``dst``."""
+    words = zeros((tokens, 2, 4), Uint32)
+    h2_words: Any = None
+    if cutlass.const_expr(h2 is not None):
+        h2_words = zeros((tokens, 4, 4), Uint32)
+    ready = False
+    while not ready:
+        valid = True
+        for token in cutlass.range_constexpr(tokens):
+            rows = (par * xmoe.shape[1] + cta % xmoe.shape[1]) * xmoe.shape[2]
+            base = xmoe.iterator + rows * 7168 + token * 7168
+            if active:
+                for chunk in cutlass.range_constexpr(2):
+                    pointer = base + (first_chunk + 28 * chunk) * 8
+                    valid = poll_words(words[token, chunk, None], pointer, valid)
+            if cutlass.const_expr(h2 is not None):
+                warp, lane, pairs = h2
+                if (warp < 7) & (pairs > 0):
+                    for chunk in cutlass.range_constexpr(4):
+                        pointer = base + (32 * warp + lane + 224 * chunk) * 8
+                        valid = poll_words(h2_words[token, chunk, None], pointer, valid)
+        ready = cute.arch.vote_all_sync(valid)
+    if active:
+        for token in cutlass.range_constexpr(tokens):
+            for chunk in cutlass.range_constexpr(2):
+                pointer = dst + token * 14336 + (first_chunk + 28 * chunk) * 16
+                values = tuple(words[token, chunk, j] for j in range(4))
+                store_shared_words(pointer, values)
+    return words, h2_words
+
+
+@cute.jit
+def score_record(score, biased):
+    """Published (score, score+bias) bits; NaN or 0x80000000 bits become (0, -inf)."""
+    score_bits, biased_bits = Uint32(float_bits(score)), Uint32(float_bits(biased))
+    sentinel = (score_bits == Uint32(0x80000000)) | (biased_bits == Uint32(0x80000000))
+    if sentinel | (score != score) | (biased != biased):
+        score_bits, biased_bits = Uint32(0), Uint32(0xFF800000)
+    return score_bits, biased_bits
+
+
+@cute.jit
+def stage_rows(slots, pitch, slot, weights, first, rows, barrier, skip=0):
+    """Bulk-copy weight rows ``first + skip + i`` into slots ``slot + i``."""
+    cute.arch.mbarrier_arrive_and_expect_tx(barrier, rows * 14336)
+    for row in range(rows):
+        dst = slots + (slot + row) * pitch
+        src = weights.iterator + (first + skip + row) * 7168
+        copy_evict(dst, src, size=14336, barrier=barrier)
+
+
+class Top16Fast:
+    """Eight-warp top16 of 896 experts, lower ids winning ties. ``scores`` holds FP32
+    (score, score+bias) bits per [parity, token, expert]; ``run`` polls ``token``'s
+    parity-``par`` records until published and writes ``ids[16*token:]``. Runs on the
+    first 256 threads of the CTA with named barrier 1 and 13,200 bytes of SMEM."""
+
+    def __init__(self, batch):
+        self.batch = batch
+
+    @cute.jit
+    def arrays(self, storage):
+        km = view(storage, 0, Uint32, (8, 32), (32, 1))
+        gc = view(storage, 1024, Uint32, 32)
+        wc = view(storage, 1152, Uint32, 8)
+        threshold = view(storage, 1184, Uint32, 1)
+        cc = view(storage, 1200, Uint64, 904)
+        cs = view(storage, 8432, Uint32, 904)
+        cc2 = view(storage, 12048, Uint64, 64)
+        selected_scores = view(storage, 12560, Float32, 80)
+        selected_ids = view(storage, 12880, Int32, 80)
+        return km, gc, wc, threshold, cc, cs, cc2, selected_scores, selected_ids
+
+    @cute.jit
+    def rank_small(self, arrays, counts, n, lane, ids, groups: cutlass.Constexpr[int]):
+        cc, cs, cc2, selected_scores, selected_ids = arrays[4:]
+        mine = cute.make_rmem_tensor(groups, Uint64)
+        score = cute.make_rmem_tensor(groups, Uint32)
+        for half in cutlass.range_constexpr(groups):
+            index = Uint32(lane + 32 * half)
+            total = Uint32(0)
+            slot = index
+            for warp in cutlass.range_constexpr(7):
+                total += counts[warp]
+                past = index >= total
+                gap = (128 if warp < 4 else 96) - counts[warp]
+                slot += gap if past else Uint32(0)
+            slot = min(slot, Uint32(895))
+            mine[half] = cc[slot]
+            score[half] = cs[slot]
+            if index >= n:
+                mine[half] = Uint64(0)
+            cc2[index] = mine[half]
+        cute.arch.sync_warp()
+        for half in cutlass.range_constexpr(groups):
+            rank = zeros(4, Uint32)
+            for q in cutlass.range_constexpr(32 * groups):
+                rank[q % 4] += Uint32(cc2[q] > mine[half])
+            r = (rank[0] + rank[1]) + (rank[2] + rank[3])
+            index = Uint32(lane + 32 * half)
+            selected = (index < n) & (r < 16)
+            slot = r if selected else 16 + index
+            expert = Int32(Uint32(0xFFFFFFFF) - Uint32(mine[half]))
+            selected_ids[slot] = expert
+            selected_scores[slot] = bits_float(score[half])
+            if selected:
+                mailbox_store(ids + r, (Uint32(expert),))
+
+    @cute.jit
+    def run(self, scores, ids, storage, par, token):
+        tid, _, _ = cute.arch.thread_idx()
+        warp, lane = tid // 32, tid % 32
+        records = zeros((4, 2), Uint32)
+        pending = Uint32(0)
+        for i in cutlass.range_constexpr(4):
+            if warp + 8 * i < 28:
+                pending |= Uint32(1 << i)
+        while pending != 0:
+            for i in cutlass.range_constexpr(4):
+                if (pending & Uint32(1 << i)) != 0:
+                    row = (par * self.batch + token) * 896 + lane
+                    pointer = scores.iterator + (row + 32 * (warp + 8 * i)) * 2
+                    words = mailbox_load(pointer, count=2)
+                    records[i, 0], records[i, 1] = words[0], words[1]
+            for i in cutlass.range_constexpr(4):
+                done = (pending & Uint32(1 << i)) != 0
+                if published((records[i, j] for j in range(2)), done):
+                    pending &= Uint32(0xFFFFFFFF ^ (1 << i))
+        cute.arch.barrier(barrier_id=1, number_of_threads=256)
+        ids = ids.iterator + token * 16
+        arrays = self.arrays(storage)
+        km, gc, wc, threshold, cc, cs, _, selected_scores, selected_ids = arrays
+        key = cute.make_rmem_tensor(4, Uint32)
+        maximum = Uint32(0)
+        for i in cutlass.range_constexpr(4):
+            biased = bits_float(records[i, 1])
+            bits = records[i, 1]
+            negative = (bits & Uint32(0x80000000)) != 0
+            ordered = ~bits if negative else bits | Uint32(0x80000000)
+            key[i] = ordered if (warp + 8 * i < 28) & (biased == biased) else Uint32(0)
+            maximum = max(maximum, key[i])
+        km[warp, lane] = maximum
+        cute.arch.barrier(barrier_id=1, number_of_threads=256)
+        if warp == 0:
+            group_max = km[0, lane]
+            for w in cutlass.range_constexpr(1, 8):
+                group_max = max(group_max, km[w, lane])
+            composite = (group_max & Uint32(0xFFFFFFE0)) | Uint32(31 - lane)
+            gc[lane] = composite
+            cute.arch.sync_warp()
+            ranks = zeros(4, Int32)
+            for q in cutlass.range_constexpr(32):
+                ranks[q % 4] += Int32(gc[q] > composite)
+            r = (ranks[0] + ranks[1]) + (ranks[2] + ranks[3])
+            least = redux(composite if r < 16 else Uint32(0xFFFFFFFF))
+            threshold[0] = least & Uint32(0xFFFFFFE0)
+        cute.arch.barrier(barrier_id=1, number_of_threads=256)
+        cutoff = threshold[0]
+        lane_mask = (Uint32(1) << lane) - Uint32(1)
+        region = Uint32(128 * warp if warp < 4 else 512 + 96 * (warp - 4))
+        offset = Uint32(0)
+        for i in cutlass.range_constexpr(4):
+            live = warp + 8 * i < 28
+            ballot = cute.arch.vote_ballot_sync(live & (key[i] >= cutoff))
+            candidate = ((ballot >> lane) & 1) != 0
+            compact = region + offset + cute.arch.popc(ballot & lane_mask)
+            position = compact if candidate else Uint32(896 + warp)
+            expert = Uint32(lane + 32 * (warp + 8 * i))
+            cc[position] = (Uint64(key[i]) << 32) | Uint64(Uint32(0xFFFFFFFF) - expert)
+            cs[position] = records[i, 0]
+            offset += cute.arch.popc(ballot)
+        wc[warp] = offset
+        cute.arch.barrier(barrier_id=1, number_of_threads=256)
+        if warp == 0:
+            counts = cute.make_rmem_tensor(8, Uint32)
+            for w in cutlass.range_constexpr(8):
+                counts[w] = wc[w]
+            n = ((counts[0] + counts[1]) + (counts[2] + counts[3])) + (
+                (counts[4] + counts[5]) + (counts[6] + counts[7])
+            )
+            if n <= 32:
+                self.rank_small(arrays, counts, n, lane, ids, 1)
+            elif n <= 64:
+                self.rank_small(arrays, counts, n, lane, ids, 2)
+            else:
+                prefix = cute.make_rmem_tensor(8, Uint32)
+                total = Uint32(0)
+                for w in cutlass.range_constexpr(8):
+                    prefix[w] = total
+                    total += counts[w]
+                for index in range(lane, n, 32):
+                    slot = Uint32(index)
+                    for w in cutlass.range_constexpr(7):
+                        past = index >= prefix[w + 1]
+                        gap = (128 if w < 4 else 96) - counts[w]
+                        slot += gap if past else Uint32(0)
+                    mine = cc[slot]
+                    rank = Uint32(0)
+                    for w in cutlass.range_constexpr(8):
+                        base = 128 * w if w < 4 else 512 + 96 * (w - 4)
+                        for q in range(counts[w]):
+                            rank += Uint32(cc[base + q] > mine)
+                    if rank < 16:
+                        selected_expert = Int32(Uint32(0xFFFFFFFF) - Uint32(mine))
+                        selected_ids[rank] = selected_expert
+                        selected_scores[rank] = bits_float(cs[slot])
+                        mailbox_store(ids + rank, (Uint32(selected_expert),))
+            cute.arch.sync_warp()
+        return selected_scores

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/shared.py

@@ -0,0 +1,684 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused shared expert for B>8: gate/up, SiTU, down and the routed tail's reduce-scatter
+push in one persistent tcgen05 kernel, synchronized by DSMEM and one global-memory
+barrier instead of kernel bounds.
+
+120 CTAs of 160 threads in clusters of ks = 120 / tiles (TP8 5, TP16 10). Loader warps
+0-3 feed a 128B-swizzled ring of W [64, 64] (TMA) and x [N, 64] (cp.async) K blocks to
+warp 4, which issues each block's four M=64 kind::f16 MMAs warp-uniformly
+(``mma_bf16_block``). Gate/up M-tile m (32 gate rows and the matching 32 up rows)
+belongs to cluster m, whose CTAs split its K. Each CTA pushes its FP32 partial rows
+over DSMEM to the owner of their h columns (``width`` columns per owner); an owner
+thread per (column, token) sums the parts in K (rank) order, rounds gate and up to BF16
+for SiTU, and the owner stores h rows as vectors. After the grid's h barrier, CTAs
+below 112 run 64 down rows (the first ring's worth of weights already staged) and
+peer-store them, -0 made +0, into the tail's RS mailboxes. The last CTA through the
+barrier re-arms the counters."""
+
+import functools
+
+import cuda.bindings.driver as cuda
+import cutlass as c
+import cutlass.cute as cute
+import torch
+from cutlass import Constexpr, Float32, Int32, Int64, Uint32
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import fake, tensor_maps
+from ..common.primitives import (
+    acquire_tensormap,
+    alloc,
+    asm,
+    at,
+    bits_float,
+    commit_asm,
+    copy_g2s,
+    copy_s2cluster,
+    div_rn,
+    fence_tmem,
+    float_bits,
+    launch,
+    load_tma2,
+    load_tmem,
+    mma_bf16_block,
+    multicast_words,
+    named_barrier,
+    pack_pair,
+    push_float,
+    release_tmem,
+    sigmoid,
+    smem_desc,
+    smem_words,
+    split,
+    store_relaxed,
+    store_words,
+    thread,
+    view,
+    wait,
+)
+from ..gemm.ws_gemm import positive_zero, positive_zero32
+
+tanh = functools.partial(cute.math.tanh, fastmath=False)
+GRID, ROWS = 120, 112  # CTAs; CTAs with 64 down rows
+# With the MoE front rows: the K split by TP and a deeper ring that still fits beside
+# three o-proj consumer CTAs (<= 30 KB each; the kernel launches at the consumer's
+# entry and loads its first ring meanwhile).
+FRONT_KS, FRONT_RING = {8: 4, 16: 6}, {8: 112 * 1024, 16: 96 * 1024}
+COUNTERS = 4  # h ready, exit
+
+
+@cute.jit
+def situ(gate, up, beta, linear_beta, *, native: c.Constexpr = False):
+    if c.const_expr(native):
+        inv_beta, half = div_rn(c.Float32(1), beta), c.Float32(0.5)
+        g = half * beta * tanh(gate * inv_beta) * (c.Float32(1) + tanh(gate * half))
+        if linear_beta > 0:
+            up = linear_beta * tanh(up * div_rn(c.Float32(1), linear_beta))
+        return g * up
+    g = beta * tanh(div_rn(gate, beta)) * sigmoid(gate)
+    if linear_beta > 0:
+        up = linear_beta * tanh(div_rn(up, linear_beta))
+    return g * up
+
+
+@asm()
+def prefetch_tma2(descriptor, column, row):
+    """L2 prefetch of the 2D TMA box at (``column``, ``row``)."""
+    code = "cp.async.bulk.prefetch.tensor.2d.L2.global.tile [$0,{$1,$2}];"
+    return code, "l,r,r", (descriptor, Int32(column), Int32(row))
+
+
+@asm()
+def release_count(address):
+    return "red.release.gpu.global.add.u32 [$0],1;", "l,~{memory}", (address,)
+
+
+@asm(Uint32)
+def acquire_count(address):
+    return "ld.acquire.gpu.global.u32 $0,[$1];", "=r,l,~{memory}", (address,)
+
+
+@asm(Uint32)
+def fetch_count(address):
+    return "atom.relaxed.gpu.global.add.u32 $0,[$1],1;", "=r,l,~{memory}", (address,)
+
+
+@cute.jit
+def wait_count(address, target):
+    seen = Int32(acquire_count(address))
+    while seen < target:
+        seen = Int32(acquire_count(address))
+
+
+def weight_maps(gate_up, down, front=None):
+    """TMA maps (W [64, 32] boxes, 128B swizzle) of gate/up, down, then (optional) the
+    MoE front GEMM's packed router / latent rows."""
+    dtype = cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16
+    swizzle = cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B
+    l2 = cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B
+    weights = (gate_up, down) if front is None else (gate_up, down, front)
+    specs = [
+        (w.data_ptr(), dtype, (w.shape[1], w.shape[0]), (2 * w.stride(0),))
+        + ((64, 32), swizzle, l2)
+        for w in weights
+    ]
+    return tensor_maps(specs, gate_up.device)
+
+
+class SharedBase:
+    """What the fused kernel and the split down kernel share: geometry, TMA weight and
+    cp.async row tiles, the M=64 accumulator readout and the RS push of 64 rows."""
+
+    def __init__(self, tp, batch):
+        self.tp, self.batch, self.hidden = tp, batch, 6144 // tp
+        self.down = self.hidden // 64  # down K blocks
+        self.n = (batch + 7) // 8 * 8
+        self.columns = max(32, 1 << (self.n - 1).bit_length())  # one accumulator
+
+    @cute.jit
+    def load_weight(self, tile, descriptor, barrier, row0, row1, column):
+        """W rows [row0, + 32) and [row1, + 32), 64 columns from ``column``, as a
+        K-major 128B-swizzled [64, 64] tile (two TMA boxes)."""
+        cute.arch.mbarrier_arrive_and_expect_tx(barrier, 8192)
+        load_tma2(tile, descriptor, barrier, column, row0)
+        load_tma2(tile + 4096, descriptor, barrier, column, row1)
+
+    @cute.jit
+    def load_rows(self, tile, source, column):
+        """Rows [0, batch) of ``source``, 64 columns from ``column``, as the B tile
+        (threads [0, 128))."""
+        tid, _, _, _ = thread()
+        count: Constexpr = self.batch * 8
+        for i in c.range_constexpr((count + 127) // 128):
+            item = tid + 128 * i
+            row, chunk = item // 8, item % 8
+            dst = tile + row // 8 * 1024 + row % 8 * 128 + (chunk ^ row % 8) * 16
+            if c.const_expr((i + 1) * 128 <= count):
+                copy_g2s(dst, at(source, row, column + chunk * 8))
+            else:
+                if item < count:  # noqa: SIM102
+                    copy_g2s(dst, at(source, row, column + chunk * 8))
+
+    @cute.jit
+    def accumulated(self, tmem):
+        """This lane's accumulator row (M=64: rows 16 warp + lane, lane < 16); later
+        MMAs may then reuse the columns."""
+        _, warp, _, _ = thread()
+        values = load_tmem(tmem + Uint32(32 * warp << 16), self.columns)
+        fence_tmem(before=True)
+        return values
+
+    @cute.jit
+    def elected(self, barrier):
+        with cute.arch.elect_one():
+            commit_asm(barrier)
+
+    @cute.jit
+    def push(self, mailbox, staging, values, row0):
+        """Down rows [row0, + 64) (``values`` from ``accumulated``) -> FP32 [token,
+        row] tile at ``staging`` -> BF16 (-0 made +0) 16-byte peer stores into the RS
+        mailbox slot; ``mailbox`` is (peers, parity, rank, capacity). Warps 0-3."""
+        peers, parity, rank, capacity = mailbox
+        tid, warp, lane, _ = thread()
+        tile = view(staging, 0, Float32, (self.n, 68), (68, 1))
+        if lane < 16:
+            for t in c.range_constexpr(self.batch):
+                tile[t, 16 * warp + lane] = values[t]
+        named_barrier(1, 128)
+        shard: Constexpr = 7168 // self.tp
+        peer, column = row0 // shard, row0 % shard
+        count: Constexpr = self.batch * 8
+        for i in c.range_constexpr((count + 127) // 128):
+            item = tid + 128 * i
+            if item < count:
+                t, part = item // 8, item % 8
+                reduced = []
+                for half in c.range_constexpr(2):
+                    words = smem_words(at(tile, t, part * 8 + half * 4))
+                    reduced += [bits_float(w) for w in words]
+                packed = tuple(
+                    pack_pair(reduced[2 * v], reduced[2 * v + 1], policy="latent")
+                    for v in range(4)
+                )
+                row = (parity * self.tp + rank) * capacity + t
+                offset = Int64(row * shard + column + part * 8)
+                store_relaxed(peers[peer] + offset * 2, packed)
+        named_barrier(1, 128)
+
+
+class SharedUmma(SharedBase):
+    def __init__(self, tp, batch, budget=80 * 1024, gate=False, front=None):
+        super().__init__(tp, batch)
+        # gate: the split variant's front-window half (gate/up and SiTU into h; the
+        # down rows run in ``SharedDown`` beside the routed experts). front: (rows,
+        # split, valid, mailbox row) of the MoE front GEMM, whose packed router /
+        # latent rows join as extra M-tiles (gate mode, on the main stream after the
+        # o-proj consumer: the first ring's weights load before the dependency wait).
+        self.gate, self.front = (gate or front is not None, front)
+        self.tiles = 2 * self.hidden // 64  # gate/up M-tiles, one per cluster
+        self.grid = GRID
+        if front is not None:
+            self.gtiles = self.tiles
+            self.tiles += -(-front[0] // 64)
+            self.grid = FRONT_KS[tp] * self.tiles
+        self.ks = self.grid // self.tiles  # cluster size: the K split
+        self.blocks = -(-112 // self.ks)  # gate/up K blocks per CTA, at most
+        self.stage = 8192 + self.n * 128
+        later = 0 if self.gate else self.down
+        self.stages = min(self.blocks + later, budget // self.stage)
+        # Owners 0 .. 32 / width - 1 take a tile's h columns ``width`` at a time.
+        self.width = 8 if self.ks < 8 else 4
+        self.owners = 32 // self.width
+        self.received = self.ks * 2 * self.width * self.n * 4  # an owner's bytes
+
+    @cute.jit
+    def __call__(self, *args):
+        options = dict(pdl=self.front is not None, min_blocks_per_mp=1)
+        launch(self.kernel, args, self.grid, 160, self.ks, **options)
+
+    @cute.kernel
+    def kernel(self, args):
+        maps = args[1]
+        tid, warp, lane, cta = thread()
+        rank = cta % self.ks
+        stages, stage = self.stages, self.stage
+        smem = c.utils.SmemAllocator()
+        ring = alloc(smem, c.Uint8, stages * stage, 1024).iterator
+        received = alloc(smem, Float32, self.received // 4, 16).iterator
+        staged = alloc(smem, c.BFloat16, self.n * self.width, 16)  # owned h rows
+        if c.const_expr(self.front is not None):  # owned front rows [token, half]
+            staged = (staged, alloc(smem, Float32, self.n * 2 * (self.width + 1), 16))
+        bars = alloc(smem, c.Uint64, 2 * stages + 3, 8).iterator
+        slot = alloc(smem, Uint32, 1)
+        if tid == 0:
+            for s in c.range_constexpr(stages):
+                cute.arch.mbarrier_init(bars + s, 129)  # TMA + 128 cp.async arrivals
+                cute.arch.mbarrier_init(bars + stages + s, 1)
+            for i in c.range_constexpr(3):
+                cute.arch.mbarrier_init(bars + 2 * stages + i, 1)
+            cute.arch.mbarrier_init_fence()
+            # Every rank's gate and up partials of the owned columns.
+            nbytes = self.received if rank < self.owners else 0
+            cute.arch.mbarrier_arrive_and_expect_tx(bars + 2 * stages + 2, nbytes)
+            acquire_tensormap(maps.iterator)
+            acquire_tensormap(maps.iterator + 128)
+            if c.const_expr(self.front is not None):
+                acquire_tensormap(maps.iterator + 256)
+        if warp == 4:
+            cute.arch.alloc_tmem(self.columns, slot.iterator)
+            cute.arch.relinquish_tmem_alloc_permit()
+        fence_tmem(before=True)
+        cute.arch.sync_threads()
+        fence_tmem()
+        cute.arch.cluster_arrive_relaxed()
+        cute.arch.cluster_wait()  # every owner's receive barrier is armed
+        tmem = slot[0]
+        if c.const_expr(self.front is not None):
+            # Weights do not depend on the predecessor: the first ring's worth.
+            if tid == 0:
+                m, k0, q = self.span(cta)
+                for kb in c.range_constexpr(stages):
+                    if kb < q:
+                        self.load_tile(ring + kb * stage, maps, bars + kb, m, k0 + kb)
+                # The rest of the front tiles' boxes (on the routing path) into L2.
+                if m >= self.gtiles:
+                    for kb in c.range(stages, q, unroll=1):
+                        self.prefetch_tile(maps, m, k0 + kb)
+            cute.arch.griddepcontrol_wait()
+            cute.arch.griddepcontrol_launch_dependents()
+        if warp < 4:
+            self.stream(args, ring, (received, staged), bars, tmem, cta)
+        else:
+            self.multiply(ring, bars, tmem, cta)
+        release_tmem(tmem, self.columns, warp, 4)
+        cute.arch.cluster_arrive_relaxed()
+        cute.arch.cluster_wait()
+
+    @cute.jit
+    def load_tile(self, tile, maps, barrier, m, k):
+        """K block ``k`` of M-tile ``m``: 32 gate and the matching 32 up rows, or (front
+        tiles past the gate/up ones) 64 consecutive front rows."""
+        if m < self.gtiles:
+            rows = (32 * m, self.hidden + 32 * m)
+            self.load_weight(tile, maps.iterator, barrier, *rows, 64 * k)
+        else:
+            row = 64 * (m - self.gtiles)
+            self.load_weight(tile, maps.iterator + 256, barrier, row, row + 32, 64 * k)
+
+    @cute.jit
+    def prefetch_tile(self, maps, m, k):
+        """L2 prefetch of ``load_tile``'s two boxes."""
+        if m < self.gtiles:
+            prefetch_tma2(maps.iterator, 64 * k, 32 * m)
+            prefetch_tma2(maps.iterator, 64 * k, self.hidden + 32 * m)
+        else:
+            row = 64 * (m - self.gtiles)
+            prefetch_tma2(maps.iterator + 256, 64 * k, row)
+            prefetch_tma2(maps.iterator + 256, 64 * k, row + 32)
+
+    @cute.jit
+    def span(self, cta):
+        """(tile, first K block, K blocks) of the CTA's gate/up work."""
+        k0, count = split(112, cta % self.ks, self.ks)
+        return cta // self.ks, k0, count
+
+    @cute.jit
+    def multiply(self, ring, bars, tmem, cta):
+        """Warp 4: every K block's four K16 MMAs and its stage release; gate/up
+        completes barrier 0 of three, down (the same columns) barrier 1."""
+        stages, q = self.stages, self.span(cta)[2]
+        total = q
+        if c.const_expr(not self.gate):
+            total = q + (self.down if cta < ROWS else 0)
+        instruction: Constexpr = 1 << 4 | 1 << 7 | 1 << 10 | self.n >> 3 << 17 | 4 << 24
+        done = bars + 2 * stages
+        for kb in c.range(total, unroll=1):
+            s = kb % stages
+            wait(bars + s, kb // stages % 2, scope="cta")
+            cute.arch.fence_proxy("async.shared", space="cta")
+            fence_tmem()
+            tile = ring + s * self.stage
+            a, b = smem_desc(tile, True), smem_desc(tile + 8192, True)
+            index = kb - q  # block index within the product
+            if kb < q:
+                index = kb
+            release = Int32(kb + stages < total)
+            mma_bf16_block(tmem, a, b, instruction, index, bars + stages + s, release)
+            if kb == q - 1:
+                self.elected(done)
+        if c.const_expr(not self.gate):  # noqa: SIM102
+            if cta < ROWS:
+                self.elected(done + 1)
+
+    @cute.jit
+    def stream(self, args, ring, buffers, bars, tmem, cta):
+        """Loader warps: gate/up blocks, the first down weights, the partial push,
+        owned h columns, the h barrier, then the down blocks' h (and remaining
+        weights) and the push."""
+        x, maps, h, counters = args[0:4]
+        stages, stage, hidden = self.stages, self.stage, self.hidden
+        tid, _, _, _ = thread()
+        m, k0, q = self.span(cta)
+        gate_up, down = maps.iterator, maps.iterator + 128
+        for kb in c.range(q, unroll=1):
+            s = kb % stages
+            tile = ring + s * stage
+            if kb >= stages:
+                wait(bars + stages + s, (kb // stages - 1) % 2, scope="cta")
+            k = k0 + kb
+            if c.const_expr(self.front is not None):
+                if (tid == 0) & (kb >= stages):  # the first ring went before the wait
+                    self.load_tile(tile, maps, bars + s, m, k)
+            elif tid == 0:
+                rows = (32 * m, hidden + 32 * m)
+                self.load_weight(tile, gate_up, bars + s, *rows, 64 * k)
+            self.load_rows(tile + 8192, x, 64 * k)
+            cute.arch.cp_async_mbarrier_arrive_noinc(bars + s)
+        received, staged = buffers
+        if c.const_expr(self.gate):
+            self.partial(received, bars, tmem, cta, ring)
+            self.reduce(args, received, staged, bars, cta)
+            return
+        early = min(self.down, stages)
+        row0 = 64 * cta
+        if cta < ROWS:
+            for j in c.range(early, unroll=1):
+                s = (q + j) % stages
+                if q + j >= stages:
+                    wait(bars + stages + s, ((q + j) // stages - 1) % 2, scope="cta")
+                if tid == 0:
+                    tile = ring + s * stage
+                    self.load_weight(tile, down, bars + s, row0, row0 + 32, 64 * j)
+        self.partial(received, bars, tmem, cta)
+        self.reduce(args, received, staged, bars, cta)
+        if tid == 0:
+            wait_count(counters.iterator, GRID)
+            if fetch_count(counters.iterator + 1) == GRID - 1:
+                counters[0], counters[1] = Int32(0), Int32(0)
+        named_barrier(1, 128)
+        if cta < ROWS:
+            for j in c.range(self.down, unroll=1):
+                s = (q + j) % stages
+                tile = ring + s * stage
+                if j >= early:
+                    wait(bars + stages + s, ((q + j) // stages - 1) % 2, scope="cta")
+                    if tid == 0:
+                        self.load_weight(tile, down, bars + s, row0, row0 + 32, 64 * j)
+                self.load_rows(tile + 8192, h, 64 * j)
+                cute.arch.cp_async_mbarrier_arrive_noinc(bars + s)
+            wait(bars + 2 * self.stages + 1, 0, scope="cta")
+            fence_tmem()
+            self.push(args[4:8], ring, self.accumulated(tmem), row0)
+
+    @cute.jit
+    def partial(self, received, bars, tmem, cta, ring=None):
+        """Push the CTA's FP32 gate/up rows to their h columns' owners: row r (gate
+        r, or up r - 32) of column g = r % 32 lands in owner g // width's slot
+        [rank, r // 32, g % width]. With the front rows, the slots are staged in the
+        drained ring and each owner's leaves in one DSMEM bulk copy."""
+        tid, warp, lane, _ = thread()
+        rank = cta % self.ks
+        wait(bars + 2 * self.stages, 0, scope="cta")
+        fence_tmem()
+        values = self.accumulated(tmem)
+        if c.const_expr(self.front is not None):
+            share: Constexpr = 2 * self.width * self.n  # one owner's floats
+            if lane < 16:
+                row = 16 * warp + lane
+                g = row % 32
+                owner = g // self.width
+                cell = ((owner * 2 + row // 32) * self.width + g % self.width) * self.n
+                for v in c.range_constexpr(self.n // 4):
+                    words = tuple(float_bits(values[4 * v + i]) for i in range(4))
+                    at_v = cell + self.slot(4 * v, g % self.width)
+                    store_words(ring + at_v * 4, words, shared=True)
+            cute.arch.fence_proxy("async.shared", space="cta")
+            named_barrier(1, 128)
+            if tid < self.owners:
+                arrived = bars + 2 * self.stages + 2
+                source = ring + tid * share * 4
+                dst = received + rank * share
+                copy_s2cluster(dst, source, arrived, tid, size=share * 4)
+            return
+        if lane < 16:
+            row = 16 * warp + lane
+            g = row % 32
+            cell = ((rank * 2 + row // 32) * self.width + g % self.width) * self.n
+            arrived = bars + 2 * self.stages + 2  # the owner's receive barrier
+            for v in c.range_constexpr(self.n // 4):
+                chunk = tuple(values[4 * v + i] for i in range(4))
+                push_float(received + cell + 4 * v, chunk, arrived, g // self.width)
+
+    @cute.jit
+    def slot(self, t, j):
+        """Token ``t``'s place in a partial row of column ``j``: with the front rows
+        the row's 16-byte chunks rotate by ``j`` (conflict-free staging stores)."""
+        if c.const_expr(self.front is None):
+            return t
+        chunks: Constexpr = self.n // 4
+        return (t // 4 + j) % chunks * 4 + t % 4
+
+    @cute.jit
+    def reduce(self, args, received, staged, bars, cta):
+        """Owner: a thread per (column, token) sums the parts in rank order and
+        applies SiTU; the rows leave in ``width``-column vectors. Then the CTA counts
+        into the h barrier (fused only)."""
+        h, counters = args[2:4]
+        beta, linear_beta = args[8:10]
+        tid, _, _, _ = thread()
+        m, rank = cta // self.ks, cta % self.ks
+        wait(bars + 2 * self.stages + 2, 0)  # peers' DSMEM stores
+        parts = view(received, 0, Float32, self.received // 4)
+        if c.const_expr(self.front is None):
+            if rank < self.owners:
+                self.owned_h(h, parts, staged, m, rank, beta, linear_beta)
+        else:
+            staged, rows = staged
+            if rank < self.owners:
+                if m < self.gtiles:
+                    self.owned_h(h, parts, staged, m, rank, beta, linear_beta)
+                else:
+                    self.publish(args, parts, rows, m - self.gtiles, rank)
+        named_barrier(1, 128)
+        if c.const_expr(not self.gate):  # noqa: SIM102
+            if tid == 0:
+                cute.arch.fence_acq_rel_gpu()
+                release_count(counters.iterator)
+
+    @cute.jit
+    def owned_h(self, h, parts, staged, m, rank, beta, linear_beta):
+        """Owner of gate/up tile ``m``: SiTU of the summed parts into ``width``-column
+        h vectors."""
+        tid, _, _, _ = thread()
+        batch, width = self.batch, self.width
+        for i in c.range_constexpr((width * batch + 127) // 128):
+            item = tid + 128 * i
+            if item < width * batch:
+                j, t = item // batch, item % batch
+                sums = [Float32(0), Float32(0)]
+                for source in c.range_constexpr(self.ks):
+                    for half in c.range_constexpr(2):
+                        cell = (source * 2 + half) * width + j
+                        sums[half] = sums[half] + parts[cell * self.n + self.slot(t, j)]
+                gate = Float32(c.BFloat16(sums[0]))
+                up = Float32(c.BFloat16(sums[1]))
+                value = c.BFloat16(situ(gate, up, beta, linear_beta))
+                staged[t * width + j] = value
+        named_barrier(1, 128)
+        if tid < batch:
+            row = staged.iterator + tid * width
+            words = smem_words(row) if width == 8 else smem_words(row, count=2)
+            store_words(at(h, tid, 32 * m + width * rank), words)
+
+    @cute.jit
+    def publish(self, args, parts, rows, f, rank):
+        """Owner of front tile ``f``: sums each (row, token) in rank order, then
+        publishes ``width``-row vectors like the front GEMM: FP32 router logits (rows
+        < valid) and BF16 latent columns (rows >= split), -0 made +0, as multicast
+        stores into the front mailbox (``mc``; this rank's ``slot`` / ``base``)."""
+        mc, slot, base = args[10:13]
+        _, split, valid, pitch = self.front
+        tid, _, _, _ = thread()
+        batch, width = self.batch, self.width
+        shard: Constexpr = 3584 // self.tp
+        for i in c.range_constexpr((width * batch + 127) // 128):
+            item = tid + 128 * i
+            if item < width * batch:
+                j, t = item // batch, item % batch
+                for half in c.range_constexpr(2):
+                    total = Float32(0)
+                    for source in c.range_constexpr(self.ks):
+                        cell = (source * 2 + half) * width + j
+                        total = total + parts[cell * self.n + self.slot(t, j)]
+                    rows[(t * 2 + half) * (width + 1) + j] = total
+        named_barrier(1, 128)
+        if tid < 2 * batch:
+            t, half = tid // 2, tid % 2
+            row = 64 * f + 32 * half + width * rank
+            values = [rows[tid * (width + 1) + v] for v in range(width)]
+            if row < valid:
+                for v in c.range_constexpr(width // 4):
+                    words = tuple(positive_zero32(values[4 * v + i]) for i in range(4))
+                    offset = t * pitch + 3584 + 2 * (slot + row + 4 * v)
+                    multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2")
+            elif (row >= split) & (row - split < shard):
+                pairs = range(0, width, 2)
+                words = tuple(positive_zero(values[i], values[i + 1]) for i in pairs)
+                offset = t * pitch + base + row - split
+                multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2")
+
+
+class SharedDown(SharedBase):
+    """The split variant's down projection beside the routed experts: ``sms``
+    persistent CTAs (one per SM the experts' grid leaves) stage h [N, hidden] once,
+    take every ``sms``-th 64-row tile, stream its W blocks through a TMA ring (warp 5)
+    into alternating TMEM accumulators (warp 4) and peer-store each tile like the
+    fused kernel (warps 0-3). Exactly the fused kernel's down numerics."""
+
+    def __init__(self, tp, batch, sms, budget=200 * 1024):
+        super().__init__(tp, batch)
+        self.sms, self.image = sms, self.down * self.n * 128  # h as B tiles
+        ring = budget - self.image - self.n * 68 * 4
+        self.stages = min(24, ring // 8192)  # W [64, 64] K blocks (8 KB)
+
+    @cute.jit
+    def __call__(self, *args):
+        options = dict(pdl=False, min_blocks_per_mp=1)
+        launch(self.kernel, args, self.sms, 192, **options)
+
+    @cute.kernel
+    def kernel(self, args):
+        h, maps = args[0:2]
+        tid, warp, lane, cta = thread()
+        stages, columns = self.stages, self.columns
+        smem = c.utils.SmemAllocator()
+        image = alloc(smem, c.Uint8, self.image, 1024).iterator
+        ring = alloc(smem, c.Uint8, stages * 8192, 1024).iterator
+        staging = alloc(smem, c.Uint8, self.n * 68 * 4, 16).iterator
+        # full [stages], empty [stages], accumulator full [2] / empty [2], h ready
+        bars = alloc(smem, c.Uint64, 2 * stages + 5, 8).iterator
+        slot = alloc(smem, Uint32, 1)
+        acc_full, acc_empty, ready = (bars + 2 * stages + i for i in (0, 2, 4))
+        if tid == 0:
+            for s in c.range_constexpr(stages):
+                cute.arch.mbarrier_init(bars + s, 1)
+                cute.arch.mbarrier_init(bars + stages + s, 1)
+            for i in c.range_constexpr(2):
+                cute.arch.mbarrier_init(acc_full + i, 1)
+                cute.arch.mbarrier_init(acc_empty + i, 128)
+            cute.arch.mbarrier_init(ready, 128)
+            cute.arch.mbarrier_init_fence()
+        if warp == 4:
+            cute.arch.alloc_tmem(2 * columns, slot.iterator)
+            cute.arch.relinquish_tmem_alloc_permit()
+        fence_tmem(before=True)
+        cute.arch.sync_threads()
+        fence_tmem()
+        tmem = slot[0]
+        tiles = (ROWS - cta + self.sms - 1) // self.sms
+        if warp < 4:
+            for j in c.range_constexpr(self.down):
+                self.load_rows(image + j * self.n * 128, h, 64 * j)
+            cute.arch.cp_async_mbarrier_arrive_noinc(ready)
+            for i in c.range(tiles, unroll=1):
+                wait(acc_full + i % 2, i // 2 % 2, scope="cta")
+                fence_tmem()
+                values = self.accumulated(tmem + Uint32(i % 2 * columns))
+                cute.arch.mbarrier_arrive(acc_empty + i % 2)
+                self.push(args[2:6], staging, values, 64 * (cta + i * self.sms))
+        elif warp == 4:
+            self.multiply(image, ring, bars, tmem, tiles)
+        elif lane == 0:
+            down = maps.iterator + 128
+            acquire_tensormap(down)  # in the thread that issues its TMAs
+            for i in c.range(tiles, unroll=1):
+                row0 = 64 * (cta + i * self.sms)
+                for j in c.range(self.down, unroll=1):
+                    g = i * self.down + j
+                    s = g % stages
+                    if g >= stages:
+                        wait(bars + stages + s, (g // stages - 1) % 2, scope="cta")
+                    self.load_weight(
+                        ring + s * 8192, down, bars + s, row0, row0 + 32, 64 * j
+                    )
+        release_tmem(tmem, 2 * columns, warp, 4)
+
+    @cute.jit
+    def multiply(self, image, ring, bars, tmem, tiles):
+        """Warp 4: each tile's K blocks into accumulator half i % 2."""
+        stages, columns = self.stages, self.columns
+        instruction: Constexpr = 1 << 4 | 1 << 7 | 1 << 10 | self.n >> 3 << 17 | 4 << 24
+        acc_full, acc_empty, ready = (bars + 2 * stages + i for i in (0, 2, 4))
+        total = tiles * self.down
+        wait(ready, 0, scope="cta")
+        cute.arch.fence_proxy("async.shared", space="cta")
+        for i in c.range(tiles, unroll=1):
+            if i >= 2:
+                wait(acc_empty + i % 2, (i // 2 - 1) % 2, scope="cta")
+            fence_tmem()
+            d = tmem + Uint32(i % 2 * columns)
+            for j in c.range(self.down, unroll=1):
+                g = i * self.down + j
+                s = g % stages
+                wait(bars + s, g // stages % 2, scope="cta")
+                fence_tmem()
+                a = smem_desc(ring + s * 8192, True)
+                b = smem_desc(image + j * self.n * 128, True)
+                release = Int32(g + stages < total)
+                mma_bf16_block(d, a, b, instruction, j, bars + stages + s, release)
+            self.elected(acc_full + i % 2)
+
+
+@functools.cache
+def compile_shared(tp, batch, device, gate=False, front=None):
+    """Arguments: x, ``weight_maps``, h, counters, RS peers, parity, rank, capacity,
+    beta, linear beta. ``gate``: the split variant's gate/up + SiTU half; ``front``
+    (rows, split, valid, mailbox row): with the MoE front GEMM's rows, whose results
+    go to the front mailbox (more arguments: multicast address, logit slot, latent
+    base; ``weight_maps`` with the front's)."""
+    maps = fake((256 if front is None else 384,), None, c.Uint8, 64)
+    args: tuple = (fake((batch, 7168)), maps, fake((batch, 6144 // tp)))
+    args += (fake((COUNTERS,), None, Int32), fake((tp,), None, Int64))
+    args += (Int32(0), Int32(0), Int32(0), Float32(4), Float32(25))
+    if front is not None:
+        args += (Int64(0), Int32(0), Int32(0))
+    ring = 80 * 1024 if front is None else FRONT_RING[tp]
+    kernel = SharedUmma(tp, batch, ring, gate, front)
+    with torch.accelerator.device_index(device):
+        return cute.compile(kernel, *args, make_fake_stream())
+
+
+@functools.cache
+def compile_down(tp, batch, device, sms):
+    """Arguments: h, ``weight_maps``, RS peers, parity, rank, capacity."""
+    args: tuple = (fake((batch, 6144 // tp)), fake((256,), None, c.Uint8, 64))
+    args += (fake((tp,), None, Int64), Int32(0), Int32(0), Int32(0))
+    kernel = SharedDown(tp, batch, sms)
+    with torch.accelerator.device_index(device):
+        return cute.compile(kernel, *args, make_fake_stream())
+
+
+def scratch(device):
+    """Model-owned fused shared-expert counters (zeroed; re-armed by each launch)."""
+    return torch.zeros(COUNTERS, dtype=torch.int32, device=device)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/chain.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Small-batch MoE chains: front -> routed experts -> (shared expert) -> polling tail.
+
+TP8 B1: ``MoeFrontTP8`` -> ``ClusterExpertsTP8`` -> ``MoeSharedTP8`` -> ``MoeTail``
+(front_poll, trigger 3). TP16 B1/B2: ``MoeFrontTP16`` (shared expert inside, on
+designated SMs) -> ``ClusterExpertsTP16`` -> ``MoeTail``. The o-projection consumer
+publishes the post-attention row into the x_moe mailbox (parity = layer index % 2) and
+bumps an epoch hint; the co-resident front has staged its weights by then. MoE tail
+polls the G/W/S slots (three parities, 0xffff-armed, re-armed after use) and publishes
+the next AttnRes input."""
+
+import functools
+
+import cutlass.cute as cute
+import torch
+from cutlass import Float32, Float64, Int32, Uint64
+from cutlass.cute.runtime import make_fake_stream
+
+from ...common.host import fake, layout, stream, to_cute
+from ...oproj.prepared import OprojMailbox
+from ..tail_runtime import POLL_CAP, compile_tail
+from .experts import ClusterExpertsTP8, ClusterExpertsTP16
+from .front_tp8 import MoeFrontTP8, MoeSharedTP8, SharedSituTables
+from .front_tp16 import MoeFrontTP16, initialize_designated
+
+# Kernel class and trailing scalar arguments of each small-batch launch.
+_KERNELS = {
+    "front": (MoeFrontTP8, (Int32,) * 3 + (Uint64, Float32)),
+    "front16": (MoeFrontTP16, (Int32,) * 3 + (Uint64,) + (Float32,) * 4 + (Float64,)),
+    "experts_tp16": (ClusterExpertsTP16, (Float32,) * 3),
+    "experts_tp8": (ClusterExpertsTP8, (Float32,) * 3),
+    "shared": (MoeSharedTP8, (Int32,) * 2),
+    "tables": (SharedSituTables, (Float32,) * 2),
+}
+
+
+@functools.cache
+def _compile(kind, specs, device, *args, **kwargs):
+    kernel, scalars = _KERNELS[kind]
+    operands = [fake(*spec) for spec in specs] + [t(0) for t in scalars]
+    with torch.accelerator.device_index(device):
+        return cute.compile(kernel(*args, **kwargs), *operands, make_fake_stream())
+
+
+def _launch(kind, tensors, *args, **kwargs):
+    specs = tuple(map(layout, tensors))
+    code = _compile(kind, specs, tensors[0].device.index, *args, **kwargs)
+    return code, tuple(map(to_cute, tensors))
+
+
+def u32(t):
+    return t.view(torch.uint32)
+
+
+class SmallBatchBox:
+    """Model-owned small-batch storage, shared serially by every MoE layer."""
+
+    def __init__(self, group, device, tp, rank):
+        def armed(shape, fill, dtype=torch.bfloat16):
+            value = torch.empty(shape, dtype=dtype, device=device)
+            value.view(torch.int32).fill_(fill)
+            return value
+
+        zeros = functools.partial(torch.zeros, dtype=torch.int32, device=device)
+        self.tp, self.rank, self.device = tp, rank, device
+        self.sms = torch.cuda.get_device_properties(device).multi_processor_count
+        self.xmoe = armed((2, 1, 4, 7168), -0x80000000)
+        self.epoch = zeros(256)
+        self.scores = armed((7168,), -0x80000000, torch.int32)
+        self.ids = armed((16,), -1, torch.int32)
+        self.w, self.s = armed((3, POLL_CAP, 16), -1), armed((3, POLL_CAP, 7168), -1)
+        self.g = armed((3, POLL_CAP * 16, 3584), -1)
+        self.h2 = OprojMailbox(group, device, tp, rank, (2, POLL_CAP, 3584))
+        self.latent = zeros(3584, dtype=torch.bfloat16)
+        self.hg, self.counter = zeros(6144), zeros(1)
+        self.tables = {}
+        self.batches = (1,) if tp == 8 else (1, 2)
+        if tp == 16:  # resident-SM placement mask plus per-batch route/gate scratch
+            self.designated = zeros(self.sms + 16)
+            initialize_designated(self.designated)
+            self.scratch, i32 = {}, torch.int32
+            for b in self.batches:
+                scores = armed((2 * b * 896 * 2,), -0x80000000, i32)
+                self.scratch[b] = scores, armed((b * 16,), -1, i32), zeros(2 * b * 768)
+        self.front_views = tuple(self.xmoe[p] for p in range(2))
+        self.hint_views = tuple(self.epoch[p * 32 : p * 32 + 224] for p in range(2))
+
+    def situ_tables(self, beta, linear_beta):
+        key = float(beta), float(linear_beta)
+        if key not in self.tables:
+            table = torch.empty(131072, dtype=torch.float32, device=self.device)
+            code, args = _launch("tables", (table,))
+            code(*args, Float32(beta), Float32(linear_beta), stream(self.device))
+            self.tables[key] = table
+        return self.tables[key]
+
+
+class SmallBatchChain:
+    """One MoE layer's small-batch chain; MoE tail weights come from the layer's
+    prepared tail."""
+
+    def __init__(self, box, moe, experts, tail, layer_idx, beta, linear_beta, prefetch):
+        self.box, self.tail_weights, self.tail_box = (
+            box,
+            tail.tail_weights,
+            tail.tail_box,
+        )
+        self.prefetch, self.layer_idx = prefetch, layer_idx
+        self.parity, self.slot = layer_idx % 2, int(tail.tail_weights.parity)
+        self.batches, shared = box.batches, moe.shared_experts
+        router, bias = moe.gate.weight, moe.gate.e_score_correction_bias
+        gate_up, shared_down = shared.gate_up_proj.weight, shared.down_proj.weight
+        rows = 3584 // box.tp
+        down = moe.routed_expert_down_proj.weight.narrow(0, box.rank * rows, rows)
+        scales = tuple(experts.weights[i].view(torch.uint8).reshape(-1) for i in (1, 3))
+        scale, renorm = moe.routed_scaling_factor, moe.moe_renormalize
+        self.scalars: tuple = (Int32(box.rank), Uint64(box.h2.mc), Float32(scale))
+        head: tuple = (router, bias, box.xmoe, u32(box.epoch))
+        if box.tp == 8:
+            tensors = head + (u32(box.scores), box.ids, box.w, down, box.h2.buffer)
+            tensors += (box.latent, gate_up, u32(box.hg))
+            self.front = _launch("front", tensors, grid=box.sms - 8, renorm=renorm)
+            tables = box.situ_tables(beta, linear_beta)
+            tensors = (box.xmoe, shared_down, tables, u32(box.hg), box.s)
+            self.shared = _launch("shared", tensors, grid=box.sms - 96)
+            tensors = (box.latent, box.ids, experts.maps[256:640].view(3, 128), *scales)
+            routed, args = _launch(
+                "experts_tp8", tensors + (box.g[self.slot, :16], box.counter)
+            )
+            self.routed = routed, args + experts.scalars
+            self.tail = compile_tail(8, 1, box.device.index, "front_poll", 3)
+            return
+        maps, self.codes = experts.maps[:256].view(2, 128), {}
+        for b in self.batches:
+            scores, ids, hg = box.scratch[b]
+            tensors = head + (u32(scores), ids, box.w, down, box.h2.buffer, gate_up)
+            tensors += (u32(hg), shared_down, box.s, box.designated)
+            front = _launch("front16", tensors, b, grid=box.sms - 7 * b, renorm=renorm)
+            latent = box.h2.buffer[self.parity, :b].reshape(-1)
+            g = box.g[self.slot, : b * 16]
+            routed = _launch("experts_tp16", (latent, ids, maps, *scales, g), b)
+            tail = compile_tail(16, b, box.device.index, "front_poll", 2)
+            self.codes[b] = front, routed, tail
+        lb = linear_beta
+        self.scalars += (Float32(beta), Float32(1 / beta), Float32(lb), Float64(1 / lb))
+        self.experts = experts.scalars
+
+    def run(self, out, consume):
+        b, s = out.shape[0], stream(self.box.device)
+        par, slot = Int32(self.parity), Int32(self.slot)
+        if self.box.tp == 8:
+            self.front[0](*self.front[1], par, slot, *self.scalars, s)
+            self.routed[0](*self.routed[1], s)
+            self.shared[0](*self.shared[1], par, slot, s)
+            tail = self.tail
+        else:
+            (front, front_args), (routed, routed_args), tail = self.codes[b]
+            front(*front_args, par, slot, *self.scalars, s)
+            routed(*routed_args, *self.experts, s)
+        if self.prefetch is not None:
+            self.prefetch.fork("next_attention", self.layer_idx, b)
+        self.tail_weights.run(
+            self.tail_box, self.box.g, self.box.w, self.box.s, tail, b
+        )
+        if consume(out) is not out:
+            raise RuntimeError("MoE tail output must be consumed from its mailbox ring")
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/experts.py

@@ -0,0 +1,541 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Small-batch routed experts: ``ClusterExpertsTP16`` (TP16 B1/B2, six-CTA split-K
+clusters exchanging quantized hidden rows) and ``ClusterExpertsTP8`` (TP8 B1, one
+six-CTA cluster per routing slot)."""
+
+from typing import Any
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Int64, Uint8, Uint32, Uint64
+from cutlass.cute.arch import store_async_dsmem
+
+from ...common.primitives import (
+    acquire_tensormap_asm,
+    alloc,
+    chunk_mma,
+    commit_asm,
+    copy_evict,
+    copy_scale,
+    fence_tmem,
+    float_bits,
+    fp8_bits,
+    launch,
+    load_tma3,
+    load_tmem,
+    load_volatile,
+    load_words,
+    mma_steps,
+    mx_exponent,
+    mx_rescale,
+    pack_fp8,
+    prefetch_bytes,
+    prefetch_map,
+    published,
+    release_tmem,
+    require_zero_tmem,
+    shared_address,
+    situ,
+    smem_desc,
+    split,
+    store_bf16_pair,
+    store_pair,
+    store_shared_words,
+    thread,
+    unpack_bf16,
+    view,
+    wait,
+    wait_tmem,
+    warp_max_abs,
+)
+
+
+@cute.jit
+def quantize_mx8(packed, row):
+    """(E8M0 scale, FP8 words) of eight BF16 values; four lanes share a 32-block."""
+    values = cute.make_rmem_tensor(8, Float32)
+    for k in cutlass.range_constexpr(4):
+        values[2 * k], values[2 * k + 1] = unpack_bf16(packed[row, k])
+    amax = Float32(0)
+    for k in cutlass.range_constexpr(8):
+        amax = max(amax, abs(values[k]))
+    for k in (1, 2):
+        amax = max(amax, cute.arch.shuffle_sync_bfly(amax, k))
+    sb = mx_exponent(amax)
+    rs = mx_rescale(sb)
+    lo, hi = (pack_fp8([values[4 * h + k] * rs for k in range(4)]) for h in (0, 1))
+    return sb, (lo, hi)
+
+
+class _ClusterExperts:
+    """Shared launch of the six-CTA-cluster expert kernels."""
+
+    kernel: Any
+    threads: int
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, 96, self.threads, cluster=6, min_blocks_per_mp=1)
+
+
+class ClusterExpertsTP16(_ClusterExperts):
+    """Paired CTAs split FC1 and exchange quantized hidden rows; the TP16 intermediate
+    has 192 live values padded to 256. IDs [B, 16], x BF16 [B, 3584], output BF16
+    [B * 16, 3584]. Polled IDs are re-armed per cluster; the caller re-arms latent rows
+    after all consumers complete."""
+
+    threads = 288
+
+    def __init__(self, batch: int):
+        """Polled IDs and latent rows; successors are released before the final wait."""
+        self.batch = batch
+        self.off_sfa = 6 * 32768
+        self.off_x = self.off_sfa + 6 * 1024
+        self.off_h = self.off_x + 7 * 2048
+        self.off_xsf = self.off_h + batch * 2048
+        self.off_hsf = self.off_xsf + 14 * 128 + 384
+        self.off_partial = self.off_hsf + batch * 1152
+        self.off_misc = self.off_partial + batch * 256
+        self.off_acc1 = self.off_misc + 104
+        self.off_acc2 = self.off_acc1 + batch * 8
+        self.off_precv = self.off_acc2 + batch * 40
+        self.off_hrecv = self.off_precv + batch * 8
+        self.off_tmem = self.off_hrecv + batch * 8
+        self.off_amax = self.off_tmem + 4
+        self.off_stage = (self.off_amax + batch * 16 + 15) // 16 * 16
+        # 512 B unused reserve keeps the tuned SMEM footprint.
+        self.smem_bytes = (self.off_stage + batch * 64 + 512 + 15) // 16 * 16
+        cols = (8, 64, 128, 184) if batch == 1 else (16, 96, 144, 200)
+        self.col_acc2, self.col_sfa, self.col_xsf, self.col_hsf = cols
+
+    @cute.jit
+    def get_id(self, ids, j, token):
+        expert = Int32(load_volatile(ids.iterator + token * 16 + j)[0])
+        while expert < 0:
+            expert = Int32(load_volatile(ids.iterator + token * 16 + j)[0])
+        return expert
+
+    @cute.jit
+    def issue(self, sm, maps, s13, s2, full, n, e0, e1, rank, tile0, tiles):
+        slot = n % 6
+        cute.arch.mbarrier_arrive_and_expect_tx(full.iterator + slot, 16384 + 1024)
+        if n < 7 * self.batch:
+            expert = e0 if n < 7 else e1
+            ks = 7 * (rank % 2) + n % 7
+            row = expert * 512 + rank // 2 * 128
+            dst = sm.iterator + slot * 32768
+            load_tma3(dst, maps.iterator, full.iterator + slot, row, ks * 2)
+            dst = sm.iterator + self.off_sfa + slot * 1024
+            src = s13.iterator + Int64(expert) * 57344 + rank // 2 * 14336 + ks * 1024
+            copy_evict(dst, src, size=1024, barrier=full.iterator + slot)
+        else:
+            u = n - 7 * self.batch
+            expert = e0 if u < tiles else e1
+            tile = tile0 + u % tiles
+            dst, desc = sm.iterator + slot * 32768, maps.iterator + 128
+            load_tma3(dst, desc, full.iterator + slot, expert * 3584 + tile * 128, 0)
+            dst = sm.iterator + self.off_sfa + slot * 1024
+            src = s2.iterator + Int64(expert) * 28672 + tile * 1024
+            copy_evict(dst, src, size=1024, barrier=full.iterator + slot)
+
+    @cute.kernel
+    def kernel(self, args):
+        x, ids, maps, s13, s2, output, *situ_args = args
+        tid, warp, lane, cta = thread()
+        rank, j = cta % 6, cta // 6
+        half, row_tile = rank % 2, rank // 2
+        tile0, tiles = split(28, rank, 6)
+        sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem_bytes, 1024)
+        full = view(sm.iterator, self.off_misc, Uint64, 6)
+        empty = view(sm.iterator, self.off_misc + 48, Uint64, 6)
+        xall = view(sm.iterator, self.off_misc + 96, Uint64, 1)
+        acc1 = view(sm.iterator, self.off_acc1, Uint64, self.batch)
+        acc2 = view(sm.iterator, self.off_acc2, Uint64, self.batch * 5)
+        precv = view(sm.iterator, self.off_precv, Uint64, self.batch)
+        hrecv = view(sm.iterator, self.off_hrecv, Uint64, self.batch)
+        tmem_base = view(sm.iterator, self.off_tmem, Uint32, 1)
+        if tid == 0:
+            for s in cutlass.range_constexpr(6):
+                cute.arch.mbarrier_init(full.iterator + s, 1)
+                cute.arch.mbarrier_init(empty.iterator + s, 1)
+            cute.arch.mbarrier_init(xall.iterator, 7)
+            for t in cutlass.range_constexpr(self.batch):
+                cute.arch.mbarrier_init(acc1.iterator + t, 1)
+                for i in cutlass.range_constexpr(5):
+                    cute.arch.mbarrier_init(acc2.iterator + 5 * t + i, 1)
+                cute.arch.mbarrier_init(precv.iterator + t, 1)
+                cute.arch.mbarrier_init(hrecv.iterator + t, 1)
+                cute.arch.mbarrier_arrive_and_expect_tx(precv.iterator + t, 256)
+                cute.arch.mbarrier_arrive_and_expect_tx(hrecv.iterator + t, 288)
+            cute.arch.mbarrier_init_fence()
+            acquire_tensormap_asm(maps.iterator)
+            acquire_tensormap_asm(maps.iterator + 128)
+        if warp == 2:
+            cute.arch.alloc_tmem(256, tmem_base.iterator)
+            cute.arch.relinquish_tmem_alloc_permit()
+        cute.arch.cluster_arrive_relaxed()
+        fence_tmem(before=True)
+        cute.arch.sync_threads()
+        fence_tmem()
+        tmem = tmem_base[0]
+        cute.arch.griddepcontrol_launch_dependents()
+
+        if warp == 0:
+            if lane == 0:
+                e0 = e1 = self.get_id(ids, j, 0)
+                for n in cutlass.range_constexpr(6):
+                    self.issue(sm, maps, s13, s2, full, n, e0, e1, rank, tile0, tiles)
+                ks = 7 * half + 6
+                prefetch_map(maps.iterator, e0 * 512 + row_tile * 128, 2 * ks)
+                src = s13.iterator + Int64(e0) * 57344 + row_tile * 14336 + ks * 1024
+                prefetch_bytes(src, 1024)
+                if cutlass.const_expr(self.batch == 2):
+                    e1 = self.get_id(ids, j, 1)
+                for n in cutlass.range(6, self.batch * (7 + tiles), unroll=1):
+                    slot = n % 6
+                    wait(empty.iterator + slot, ((n // 6) - 1) % 2)
+                    self.issue(sm, maps, s13, s2, full, n, e0, e1, rank, tile0, tiles)
+        elif warp == 1:
+            require_zero_tmem(tmem)
+            adesc = smem_desc(sm.iterator, True)
+            sfa = smem_desc(sm.iterator + self.off_sfa, False)
+            xdesc = smem_desc(sm.iterator + self.off_x, True)
+            xsf = smem_desc(sm.iterator + self.off_xsf, False)
+            wait_tmem(xall.iterator)
+            with cute.arch.elect_one():
+                for n in cutlass.range_constexpr(14):
+                    copy_scale(self.col_xsf + n * 4, xsf + n * 8)
+            cute.arch.sync_warp()
+            for t in cutlass.range(self.batch, unroll=1):
+                for n in cutlass.range_constexpr(7):
+                    g = 7 * t + n
+                    slot = (t + n) % 6
+                    wait_tmem(full.iterator + slot, (g // 6) % 2)
+                    with cute.arch.elect_one():
+                        col = self.col_sfa + 8 * slot
+                        copy_scale(col, sfa + 64 * slot)
+                        copy_scale(col + 4, sfa + 64 * slot + 32)
+                        d, a, b = 8 * t, adesc + 2048 * slot, xdesc + 128 * n
+                        mma_steps(d, a, b, col, self.col_xsf + 8 * n, n != 0, False)
+                        commit_asm(empty.iterator + slot)
+                        if cutlass.const_expr(n == 6):
+                            commit_asm(acc1.iterator + t)
+                    cute.arch.sync_warp()
+            for t in cutlass.range(self.batch, unroll=1):
+                wait(hrecv.iterator + t, test=True)
+                cute.arch.fence_proxy("async.shared", space="cta")
+                fence_tmem()
+                hdesc = smem_desc(sm.iterator + self.off_h + t * 2048, True)
+                hsf = smem_desc(sm.iterator + self.off_hsf + t * 1152, False)
+                with cute.arch.elect_one():
+                    for k in cutlass.range_constexpr(6):
+                        copy_scale(self.col_hsf + 24 * t + 4 * k, hsf + k * 8)
+                cute.arch.sync_warp()
+                for i in cutlass.range_constexpr(5):
+                    if i < tiles:
+                        g = self.batch * 7 + t * tiles + i
+                        slot = g % 6
+                        wait_tmem(full.iterator + slot, (g // 6) % 2)
+                        with cute.arch.elect_one():
+                            col = self.col_sfa + 8 * slot
+                            copy_scale(col, sfa + 64 * slot)
+                            copy_scale(col + 4, sfa + 64 * slot + 32)
+                            d, a = self.col_acc2 + 8 * (5 * t + i), adesc + 2048 * slot
+                            mma_steps(d, a, hdesc, col, self.col_hsf + 24 * t, 0, True)
+                            commit_asm(empty.iterator + slot)
+                            commit_asm(acc2.iterator + 5 * t + i)
+                        cute.arch.sync_warp()
+        else:
+            n = warp - 2
+            packed = cute.make_rmem_tensor((self.batch, 4), Uint32)
+            for t in cutlass.range_constexpr(self.batch):
+                pointer = x.iterator + t * 3584 + 1792 * half + 256 * n + lane * 8
+                ready = False
+                while not ready:
+                    words = load_volatile(pointer, count=4)
+                    ready = published(words)
+                    for k in cutlass.range_constexpr(4):
+                        packed[t, k] = words[k]
+            block, k = lane // 4, lane % 4
+            kc, kb = 16 * n + 2 * block + k // 2, 8 * n + block
+            for t in cutlass.range_constexpr(self.batch):
+                sb, words = quantize_mx8(packed, t)
+                offset = self.off_x + kc // 8 * 1024 + t * 128 + ((kc % 8) ^ t) * 16
+                offset += (k % 2) * 8
+                store_shared_words(sm.iterator + offset, words)
+                if k == 0:
+                    sm[self.off_xsf + kb // 4 * 128 + 16 * t + kb % 4] = Uint8(sb)
+            cute.arch.fence_proxy("async.shared", space="cta")
+            cute.arch.sync_warp()
+            if lane == 0:
+                cute.arch.mbarrier_arrive(xall.iterator)
+            if (warp >= 4) & (warp < 8):
+                ew = warp % 4
+                own = ew // 2 == half
+                tl = tmem + (Uint32(32 * ew) << 16)
+                cute.arch.cluster_wait()
+                partial = view(sm.iterator, self.off_partial, Float32, self.batch * 64)
+                maxima = view(sm.iterator, self.off_amax, Float32, self.batch * 4)
+                for t in cutlass.range(self.batch, unroll=1):
+                    wait_tmem(acc1.iterator + t)
+                    value = load_tmem(tl + 8 * t + t)
+                    index = 64 * t + (ew % 2) * 32 + lane
+                    if not own:
+                        dst = cute.recast_ptr(partial.iterator + index, dtype=Uint32)
+                        bits = Uint32(float_bits(value))
+                        store_async_dsmem(dst, bits, precv.iterator + t, rank ^ 1)
+                    else:
+                        wait(precv.iterator + t, test=True)
+                        value += partial[index]
+                        other = cute.arch.shuffle_sync_bfly(value, 8)
+                        up = (lane & 8) == 0
+                        hidden = situ(other, value, *situ_args) if up else Float32(0)
+                        amax = warp_max_abs(hidden)
+                        if lane == 0:
+                            maxima[4 * t + ew] = amax
+                        cute.arch.barrier(barrier_id=1, number_of_threads=64)
+                        sb = mx_exponent(max(amax, maxima[4 * t + (ew ^ 1)]))
+                        if up:
+                            byte = fp8_bits(hidden * mx_rescale(sb))
+                            stage = self.off_stage + 64 * t + 16 * ew
+                            sm[stage + 2 * (lane % 8) + lane // 16] = byte
+                        cute.arch.sync_warp()
+                        if lane < 6:
+                            src = sm.iterator + self.off_stage + 64 * t + 16 * ew
+                            h_words = load_words(src, count=4, shared=True)
+                            kc = 4 * row_tile + ew
+                            h_offset = self.off_h + 2048 * t + kc // 8 * 1024
+                            h_offset += kc % 8 * 16
+                            dst = sm.iterator + cute.assume(h_offset, divby=16)
+                            store_async_dsmem(dst, h_words, hrecv.iterator + t, lane)
+                            if ew % 2 == 1:
+                                sw = (Uint32(sb) * Uint32(0x01010101),) * 4
+                                scale_offset = self.off_hsf + 1152 * t + 128 * rank
+                                dst = sm.iterator + cute.assume(scale_offset, divby=16)
+                                store_async_dsmem(dst, sw, hrecv.iterator + t, lane)
+                for t in cutlass.range(self.batch, unroll=1):
+                    for i in range(tiles):
+                        wait_tmem(acc2.iterator + 5 * t + i)
+                        value = load_tmem(tl + self.col_acc2 + 8 * (5 * t + i))
+                        other = cute.arch.shuffle_sync_down(value, 8)
+                        tile = tile0 + i
+                        column = tile * 128 + 32 * ew + 4 * (lane % 8) + lane // 8
+                        if (lane & 8) == 0:
+                            dst = output.iterator + (t * 16 + j) * 3584 + column
+                            store_pair(dst, value, other)
+        release_tmem(tmem, 256, warp, 2)
+        if (tid == 0) & (rank == 0):
+            for t in cutlass.range(self.batch, unroll=1):
+                ids[t * 16 + j] = Int32(-1)
+        if (cta == 0) & (tid == 0):  # one-thread exit wait, as in the front
+            cute.arch.griddepcontrol_wait()
+
+
+def _align_1024(value):
+    return (value + 1023) // 1024 * 1024
+
+
+class ClusterExpertsTP8(_ClusterExperts):
+    """FP4 x MXFP8 FC1 -> SiTU -> MXFP8 -> FC2 of x BF16 [3584] and IDs int32 [16] into
+    BF16 [16, 3584]. Weights stay in the loaded MoE weights' shuffled layout (three TMA
+    maps, two scale arrays); mailbox mode consumes/re-arms IDs and the latent row via a
+    grid ticket."""
+
+    slots = 6
+    slot_bytes = 32768
+    off_sfa = slots * slot_bytes
+    off_x = _align_1024(off_sfa + slots * 1024)
+    off_xsf = _align_1024(off_x + 27 * 128 + 1024)
+    off_h = _align_1024(off_xsf + 28 * 128 + 384)
+    off_hsf = off_h + 3 * 1024
+    off_hsc = off_hsf + 3 * 128 + 384
+    off_partial = off_hsc + 32
+    off_stage = off_partial + 1024
+    off_amax = off_stage + 64
+    off_bar = off_amax + 16
+    off_tmem = off_bar + (2 * slots + 8 + 3) * 8
+    smem_bytes = off_tmem + 16
+    threads = 256
+
+    @cute.jit
+    def issue(self, sm, maps, scales, index, row, chunk, offset, n, full, empty, count):
+        slot = n % 6
+        phase = (n // 6) % 2
+        if n >= 6:
+            wait(empty.iterator + slot, phase ^ 1)
+        cute.arch.mbarrier_arrive_and_expect_tx(full.iterator + slot, count * 8704)
+        dst, desc = sm.iterator + slot * 32768, maps.iterator + index * 128
+        load_tma3(dst, desc, full.iterator + slot, row, chunk)
+        dst, src = sm.iterator + self.off_sfa + slot * 1024, scales.iterator + offset
+        copy_evict(dst, src, size=count * 512, barrier=full.iterator + slot)
+
+    @cute.kernel
+    def kernel(self, args):
+        x, ids, maps, s13, s2, output, counter, *situ_args = args
+        tid, warp, lane, cta = thread()
+        rank, j = cta % 6, cta // 6
+        sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem_bytes, 1024)
+        full = view(sm.iterator, self.off_bar, Uint64, 6)
+        empty = view(sm.iterator, self.off_bar + 48, Uint64, 6)
+        accb = view(sm.iterator, self.off_bar + 96, Uint64, 8)
+        xbar, hbar = accb.iterator + 8, accb.iterator + 9
+        hsc = view(sm.iterator, self.off_hsc, Uint32, 6)
+        tmem_base = view(sm.iterator, self.off_tmem, Uint32, 1)
+        t2lo, t2hi = 28 * rank // 6, 28 * (rank + 1) // 6
+        if tid == 0:
+            for s in cutlass.range_constexpr(6):
+                cute.arch.mbarrier_init(full.iterator + s, 1)
+                cute.arch.mbarrier_init(empty.iterator + s, 1)
+            for a in cutlass.range_constexpr(8):
+                cute.arch.mbarrier_init(accb.iterator + a, 1)
+            cute.arch.mbarrier_init(xbar, 6)
+            cute.arch.mbarrier_init(hbar, 1)
+            cute.arch.mbarrier_arrive_and_expect_tx(hbar, 6 * 68)
+            cute.arch.mbarrier_init_fence()
+            for i in cutlass.range_constexpr(3):
+                acquire_tensormap_asm(maps.iterator + i * 128)
+        if warp == 1:
+            cute.arch.alloc_tmem(256, tmem_base.iterator)
+            cute.arch.relinquish_tmem_alloc_permit()
+            fence_tmem(before=True)
+        cute.arch.sync_threads()
+        cute.arch.cluster_arrive_relaxed()
+        fence_tmem()
+        tmem = tmem_base[0]
+        cute.arch.griddepcontrol_launch_dependents()
+        cute.arch.griddepcontrol_wait()
+        if warp == 0:
+            if lane == 0:
+                expert = ids[j]
+                for n in range(14):
+                    row, chunk = expert * 768 + rank * 128, 2 * n
+                    offset = Int64(expert) * 86016 + rank * 14336 + n * 1024
+                    self.issue(sm, maps, s13, 0, row, chunk, offset, n, full, empty, 2)
+                n = Int32(14)
+                for tile in range(t2lo, t2hi):
+                    offset = Int64(expert) * 43008 + tile * 1536
+                    row = expert * 3584 + tile * 128
+                    self.issue(sm, maps, s2, 1, row, 0, offset, n, full, empty, 2)
+                    n += 1
+                    offset += 1024
+                    self.issue(sm, maps, s2, 2, row, 2, offset, n, full, empty, 1)
+                    n += 1
+        elif warp == 1:
+            adesc = smem_desc(sm.iterator, True)
+            sfa = smem_desc(sm.iterator + self.off_sfa, False)
+            xdesc = smem_desc(sm.iterator + self.off_x, True)
+            xsf = smem_desc(sm.iterator + self.off_xsf, False)
+            hdesc = smem_desc(sm.iterator + self.off_h, True)
+            hsf = smem_desc(sm.iterator + self.off_hsf, False)
+            wait_tmem(xbar)
+            for n in range(14):
+                slot = n % 6
+                wait_tmem(full.iterator + slot, (n // 6) % 2)
+                for half in cutlass.range_constexpr(2):
+                    kc = 2 * n + half
+                    empty_address = acc_address = Uint32(0)
+                    if cutlass.const_expr(half == 1):
+                        empty_address = shared_address(empty.iterator + slot)
+                        if n == 13:
+                            acc_address = shared_address(accb.iterator)
+                    a = (adesc + slot * 2048 + half * 1024, xdesc + kc * 8)
+                    sf = (tmem + 64 + 8 * slot + 4 * half, sfa + slot * 64 + half * 32)
+                    sfx = (tmem + 128 + 4 * kc, xsf + kc * 8, 1, Int32(kc != 0))
+                    chunk_mma(tmem, *a, *sf, *sfx, empty_address, acc_address)
+                cute.arch.sync_warp()
+            wait(hbar)
+            if lane < 3:
+                word = (hsc[2 * lane] & 0xFFFF) | (hsc[2 * lane + 1] << 16)
+                dst = sm.iterator + self.off_hsf + lane * 128
+                store_shared_words(dst, (word,) * 4)
+            cute.arch.fence_proxy("async.shared", space="cta")
+            cute.arch.sync_warp()
+            n = Int32(14)
+            for tile in range(t2lo, t2hi):
+                a = 2 + tile - t2lo
+                for stage in cutlass.range_constexpr(2):
+                    slot = n % 6
+                    wait_tmem(full.iterator + slot, (n // 6) % 2)
+                    for half in cutlass.range_constexpr(2 if stage == 0 else 1):
+                        kh: cutlass.Constexpr = 2 * stage + half
+                        empty_address = acc_address = Uint32(0)
+                        if cutlass.const_expr(half == (1 if stage == 0 else 0)):
+                            empty_address = shared_address(empty.iterator + slot)
+                        if cutlass.const_expr(kh == 2):
+                            acc_address = shared_address(accb.iterator + a)
+                        ops: tuple = (tmem + 8 * a, adesc + slot * 2048 + half * 1024)
+                        ops += (hdesc + kh * 64, tmem + 64 + 8 * slot + 4 * half)
+                        ops += (sfa + slot * 64 + half * 32, tmem + 240 + 4 * kh)
+                        ops += (hsf + kh * 8, Int32(tile == t2lo), int(kh != 0))
+                        chunk_mma(*ops, empty_address, acc_address)
+                    cute.arch.sync_warp()
+                    n += 1
+        else:
+            xtid = tid - 64
+            packed = cute.make_rmem_tensor((3, 4), Uint32)
+            for i in cutlass.range_constexpr(3):
+                if 32 * (warp - 2) + 192 * i < 448:
+                    words = load_words(
+                        x.iterator + 8 * (xtid + 192 * i), count=4, mutable=True
+                    )
+                    for k in cutlass.range_constexpr(4):
+                        packed[i, k] = words[k]
+            for i in cutlass.range_constexpr(3):
+                f = xtid + 192 * i
+                if 32 * (warp - 2) + 192 * i < 448:
+                    sb, words = quantize_mx8(packed, i)
+                    kc, chunk = f // 16, (f % 16) // 2
+                    offset = self.off_x + kc * 128 + (chunk ^ kc % 8) * 16 + f % 2 * 8
+                    store_shared_words(sm.iterator + offset, words)
+                    kb = f // 4
+                    if f % 4 == 0:
+                        sm[self.off_xsf + (kb // 4) * 128 + kb % 4] = Uint8(sb)
+            cute.arch.fence_proxy("async.shared", space="cta")
+            cute.arch.sync_warp()
+            if lane == 0:
+                cute.arch.mbarrier_arrive(xbar)
+            if warp >= 4:
+                ew = warp % 4
+                address = tmem + (Uint32(32 * ew) << 16)
+                up = (lane & 8) == 0
+                local = 2 * (lane % 8) + lane // 16
+                cute.arch.cluster_wait()
+                wait_tmem(accb.iterator)
+                value = load_tmem(address)
+                other = cute.arch.shuffle_sync_bfly(value, 8)
+                hidden = situ(other, value, *situ_args) if up else Float32(0)
+                amax = abs(hidden)
+                for off in cutlass.range_constexpr(5):
+                    amax = max(amax, cute.arch.shuffle_sync_bfly(amax, 1 << off))
+                maxima = view(sm.iterator, self.off_amax, Float32, 4)
+                if lane == 0:
+                    maxima[ew] = amax
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+                sb = mx_exponent(max(amax, maxima[ew ^ 1]), True)
+                if up:
+                    byte = fp8_bits(hidden * mx_rescale(sb))
+                    sm[self.off_stage + ew * 16 + local] = byte
+                cute.arch.sync_warp()
+                if lane == 0:
+                    src = sm.iterator + self.off_stage + ew * 16
+                    words = load_words(src, count=4, shared=True)
+                    kc = 4 * rank + ew
+                    offset = self.off_h + (kc // 8) * 1024 + (kc % 8) * 16
+                    sw = Uint32(mx_exponent(max(maxima[0], maxima[1]), True))
+                    sw |= Uint32(mx_exponent(max(maxima[2], maxima[3]), True)) << 8
+                    for peer in range(6):
+                        dst = sm.iterator + cute.assume(offset, divby=16)
+                        store_async_dsmem(dst, words, hbar, peer)
+                        if ew == 0:
+                            store_async_dsmem(hsc.iterator + rank, sw, hbar, peer)
+                for tile in range(t2lo, t2hi):
+                    a = 2 + tile - t2lo
+                    wait_tmem(accb.iterator + a)
+                    value = load_tmem(address + 8 * a)
+                    other = cute.arch.shuffle_sync_down(value, 8)
+                    if up:
+                        column = tile * 128 + 32 * ew + 4 * (lane % 8) + lane // 8
+                        dst = output.iterator + j * 3584 + column
+                        store_bf16_pair(dst, value, other)
+        release_tmem(tmem, 256, warp, 1)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py

@@ -0,0 +1,580 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TP16 B1/B2 MoE front: SM placement, early router/latent-down staging, MMA router
+and shared gate/up, FP32 routing, latent multicast/gather and shared SiTU/down rows."""
+
+import functools
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import BFloat16, Float32, Int32, Uint8, Uint32, Uint64
+from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
+
+from .....ops.cute_dsl.latent_moe_tail.primitives import to_cute
+from ...common.primitives import (
+    alloc,
+    asm,
+    bf16,
+    copy_bulk,
+    div_rn,
+    early_weights,
+    fma,
+    gate_bf16_bits,
+    init_barriers,
+    launch,
+    load_f32x4,
+    mailbox_load,
+    mailbox_store,
+    mma_into,
+    multicast_words,
+    named_barrier,
+    prefetch_bytes,
+    shared_situ,
+    sigmoid,
+    sleep_ns,
+    smem_words,
+    split,
+    thread,
+    unpack_bf16,
+    view,
+    wait,
+    warp_sum,
+    zeros,
+)
+from ..routing import Top16Fast, bf16_word, gather_x, score_record, stage_rows
+
+
+@asm()
+def sreg(name, dtype=Uint32, *, pure=False):
+    bits, kind = (64, "l") if dtype is Uint64 else (32, "r")
+    return f"mov.u{bits} $0,%{name};", f"={kind}", (), dtype, pure
+
+
+@asm(Uint32)
+def claim_add(pointer):
+    return "atom.relaxed.gpu.global.add.u32 $0,[$1],1;", "=r,l,~{memory}", (pointer,)
+
+
+@asm()
+def counter_reduce(pointer, value, *, op="add", release=False):
+    ordering = "release" if release else "relaxed"
+    code = f"red.{ordering}.gpu.global.{op}.{'b32' if op == 'or' else 'u32'} [$0],$1;"
+    return code, "l,r,~{memory}", (pointer, Uint32(value))
+
+
+@asm(Uint32)
+def load_arrivals(pointer):
+    return "ld.acquire.gpu.global.u32 $0,[$1];", "=r,l,~{memory}", (pointer,)
+
+
+class FrontPlacement:
+    """One claiming thread per CTA; every CTA publishes its arrival."""
+
+    def __init__(self, batch):
+        self.batch = batch
+        self.offset = 8 if batch == 2 else 0
+
+    @cute.jit
+    def claim(self, epoch, designated, par, cta):
+        pointer = epoch.iterator + (2 + par) * 32 + self.offset
+        sm = sreg("smid", pure=True)
+        claimed = Int32(-1)
+        if sm < designated.shape[0] - 16:  # noqa: SIM102 - preserve guarded load
+            if designated[sm] != 0:
+                claimed = Int32(claim_add(pointer))
+                counter_reduce(pointer + 2, Uint32(0xFFFFFFFF) - Uint32(cta), op="max")
+                if cutlass.const_expr(self.batch == 2):
+                    mask = pointer + 3 + (cta // 32)
+                    counter_reduce(mask, Uint32(1) << (cta % 32), op="or")
+        counter_reduce(pointer + 1, Uint32(1), release=True)
+        return claimed
+
+    @cute.jit
+    def rearm(self, epoch, par):
+        pointer = epoch.iterator + (2 + (par ^ 1)) * 32 + self.offset
+        for i in cutlass.range_constexpr(8 if self.batch == 2 else 3):
+            mailbox_store(pointer + i, (Uint32(0),))
+
+    @cute.jit
+    def role(self, epoch, par, cta, claimed, grid):
+        pointer = epoch.iterator + (2 + par) * 32 + self.offset
+        count = Int32(0)
+        role = Int32(-1)
+        if claimed >= 0:
+            arrived = load_arrivals(pointer + 1)
+            while arrived < grid:
+                sleep_ns(64)
+                arrived = load_arrivals(pointer + 1)
+            count = Int32(mailbox_load(pointer)[0])
+            first = Uint32(0xFFFFFFFF) - mailbox_load(pointer + 2)[0]
+            if cutlass.const_expr(self.batch == 1):
+                role = claimed | (Int32(first == cta) << 16)
+            else:
+                order = Int32(0)
+                for word in cutlass.range_constexpr(5):
+                    bits = mailbox_load(pointer + 3 + word)[0]
+                    if word < cta // 32:
+                        order += Int32(cute.arch.popc(bits))
+                    elif word == cta // 32:
+                        below = (Uint32(1) << (cta % 32)) - Uint32(1)
+                        order += Int32(cute.arch.popc(bits & below))
+                role = order | (Int32(order < 2) << 16)
+        return role, count
+
+
+class FrontPlacementProbe:
+    def __init__(self, cluster):
+        self.cluster = cluster
+
+    @cute.jit
+    def __call__(self, *args):
+        grid, smem = args[0].shape[0] // 2, 220 * 1024
+        launch(self.kernel, args, grid, 128, self.cluster, pdl=False, smem=smem)
+
+    @cute.kernel
+    def kernel(self, args):
+        output, spin = args
+        tid, _, _, cta = thread()
+        sm, cluster = sreg("smid", pure=True), sreg("clusterid.x")
+        start = sreg("globaltimer", Uint64)
+        if tid == 0:
+            output[2 * cta] = Int32(sm)
+            output[2 * cta + 1] = Int32(cluster)
+        while sreg("globaltimer", Uint64) - start < spin:
+            pass
+
+
+@functools.cache
+def _compile_probe(device, cluster, length):
+    output = make_fake_tensor(Int32, (length,), (1,), assumed_align=16)
+    with torch.accelerator.device_index(device):
+        probe = FrontPlacementProbe(cluster)
+        return cute.compile(probe, output, Uint64(20000), make_fake_stream())
+
+
+def plan_designated(groups, *, engine_clusters=16, cluster_size=6):
+    """Largest-capacity-first plan of the routed-expert clusters over the SM groups."""
+    capacities = [len(group) // cluster_size for group in groups]
+    if sum(capacities) < engine_clusters:
+        raise ValueError("SM probe found insufficient routed-engine cluster capacity")
+    order = sorted(range(len(groups)), key=lambda i: (-capacities[i], -len(groups[i])))
+    counts = [0] * len(groups)
+    remaining = engine_clusters
+    for i in order:
+        counts[i] = min(capacities[i], remaining)
+        remaining -= counts[i]
+    keep = set()
+    for group, count in zip(groups, counts):
+        quota = len(group) - cluster_size * count
+        chosen: list[int] = []
+        for sm in sorted(group):
+            if sm % 2 == 0 and sm + 1 in group and len(chosen) + 2 <= quota:
+                chosen.extend((sm, sm + 1))
+        for sm in sorted(group):
+            if len(chosen) < quota and sm not in chosen:
+                chosen.append(sm)
+        keep.update(chosen)
+    return keep, counts
+
+
+def initialize_designated(table: torch.Tensor):
+    """Probe SM groups and fill the model-owned int32[nsm+16] table; returns CPU audit
+    data. Probe storage is freed before capture. Probes: 1- and 3-wave grids of
+    2..16-CTA clusters, 220 KiB SMEM (one CTA per SM), a 20 us global-timer spin."""
+    device = table.device
+    nsm = torch.cuda.get_device_properties(device).multi_processor_count
+    parent = list(range(4096))
+    seen = set()
+    launches = []
+
+    def find(sm):
+        while parent[sm] != sm:
+            parent[sm] = parent[parent[sm]]
+            sm = parent[sm]
+        return sm
+
+    with torch.accelerator.device_index(device.index):
+        stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream)
+        for cluster in range(2, 17):
+            for waves in (1, 3):
+                clusters = max(1, (nsm // cluster) * waves)
+                output = table.new_full((2 * cluster * clusters,), -1)
+                try:
+                    compiled = _compile_probe(device.index, cluster, output.numel())
+                    compiled(to_cute(output), Uint64(20000), stream)
+                except RuntimeError as e:
+                    launches.append(dict(cluster=cluster, waves=waves, error=str(e)))
+                    break
+                torch.accelerator.synchronize(device)
+                members: dict[int, list[int]] = {}
+                for sm, group in output.view(-1, 2).cpu().tolist():
+                    if sm >= 0:
+                        members.setdefault(group, []).append(sm)
+                        seen.add(sm)
+                for group in members.values():
+                    for sm in group[1:]:
+                        parent[find(sm)] = find(group[0])
+                launches.append(dict(cluster=cluster, waves=waves, clusters=clusters))
+    if len(seen) != nsm:
+        raise RuntimeError(f"Placement probe saw {len(seen)} of {nsm} SMs")
+    group_map: dict[int, list[int]] = {}
+    for sm in sorted(seen):
+        group_map.setdefault(find(sm), []).append(sm)
+    groups = sorted(group_map.values(), key=lambda group: group[0])
+    keep, counts = plan_designated(groups)
+    values = [int(sm in keep) for sm in range(nsm + 16)]
+    table.copy_(torch.tensor(values, dtype=torch.int32, device=device))
+    return dict(
+        groups=groups,
+        engine_clusters=counts,
+        designated=sorted(keep),
+        launches=launches,
+    )
+
+
+class MoeFrontTP16:
+    def __init__(self, batch: int, *, grid=None, renorm=True):
+        self.batch = batch
+        self.grid = grid if grid is not None else 141 if batch == 1 else 134
+        self.renorm = renorm
+        self.shared_bytes = 15 * 14336 + (512 * batch + 128) * 4 + 16 + 8
+        self.placement = FrontPlacement(batch)
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, self.grid, 512, min_blocks_per_mp=1)
+
+    @cute.jit
+    def gemv(self, rows, nrows, activation, partials, warp, lane):
+        group, quarter = lane // 4, lane % 4
+        acc = zeros((4, 4), Float32)
+        for iteration in cutlass.range_constexpr(14):
+            offset = (warp * 448 + iteration * 32 + 8 * quarter) * 2
+            weight = x = (Uint32(0),) * 4
+            if group < nrows:
+                weight = smem_words(rows + group * 14336 + offset)
+            if group < 2 * self.batch:
+                x = smem_words(activation + (group // 2) * 14336 + offset)
+            b0 = x[2] if (group & 1) != 0 else x[0]
+            b1 = x[3] if (group & 1) != 0 else x[1]
+            chain: cutlass.Constexpr = iteration % 4
+            a = (weight[0], weight[2], weight[1], weight[3])
+            mma_into(acc[chain, None], a, (b0, b1), bf16=True)
+        if (quarter < self.batch) & (group < nrows):
+            chain_sum = lambda c: acc[c, 0] + acc[c, 3]
+            total = (chain_sum(0) + chain_sum(1)) + (chain_sum(2) + chain_sum(3))
+            partials[(warp * 16 + group) * self.batch + quarter] = total
+
+    @cute.jit
+    def gate_pass(self, slots, rows, start, bars, partials, hg, par, tid):
+        wait(bars + 1)
+        warp, lane = tid // 32, tid % 32
+        self.gemv(slots + 9 * 14336, rows, slots + 7 * 14336, partials, warp, lane)
+        cute.arch.sync_threads()
+        if tid < self.batch * rows:
+            token, row = tid // rows, tid % rows
+            total = Float32(0)
+            for w in cutlass.range_constexpr(16):
+                total += partials[(w * 16 + row) * self.batch + token]
+            bits = Uint32(gate_bf16_bits(total))
+            pointer = hg.iterator + (par * self.batch + token) * 768 + start + row
+            mailbox_store(pointer, ((bits << 16) | Uint32(1),))
+
+    @cute.jit
+    def poll_flag(self, pointer):
+        word = mailbox_load(pointer)[0]
+        while (word & 1) == 0:
+            if cutlass.const_expr(self.batch == 2):
+                sleep_ns(128)
+            word = mailbox_load(pointer)[0]
+        return word
+
+    @cute.jit
+    def shared_dot(self, packed, hq, token):
+        value = Float32(0)
+        for chunk in cutlass.range_constexpr(3):
+            w = (*unpack_bf16(packed[chunk, 0]), *unpack_bf16(packed[chunk, 1]))
+            for j in cutlass.range_constexpr(3, -1, -1):
+                value = fma(w[j], hq[token][chunk][j], value)
+        return warp_sum(value, 32)
+
+    @cute.jit
+    def load_down(self, down, packed, pair, end_pair, lane, fits=True):
+        for row in cutlass.range_constexpr(4):
+            for chunk in cutlass.range_constexpr(3):
+                words = (Uint32(0), Uint32(0))
+                if fits & (pair + row // 2 < end_pair):
+                    words = early_weights(
+                        down.iterator + (2 * pair + row) * 384 + 4 * (lane + 32 * chunk)
+                    )
+                packed[row, chunk, 0], packed[row, chunk, 1] = words[0], words[1]
+
+    @cute.jit
+    def down_dots(self, packed, values, hq, token):
+        for row in cutlass.range_constexpr(4):
+            values[token, row] = self.shared_dot(packed[row, None, None], hq, token)
+
+    @cute.jit
+    def lane_word(self, values, token, lane):
+        lo = values[token, 0] if lane == 0 else values[token, 2]
+        hi = values[token, 1] if lane == 0 else values[token, 3]
+        return bf16_word(lo, hi)
+
+    @cute.jit
+    def store_word(self, output, par3, token, column, word):
+        row = par3 * output.shape[1] + token
+        mailbox_store(output.iterator + row * 7168 + 2 * column, (word,))
+
+    @cute.jit
+    def stream_pairs(self, down, hq, output, par3, begin_pair, end_pair, warp, lane):
+        packed = cute.make_rmem_tensor((4, 3, 2), Uint32)
+        values = cute.make_rmem_tensor((self.batch, 4), Float32)
+        for pair in range(begin_pair + 2 * warp, end_pair, 32):
+            self.load_down(down, packed, pair, end_pair, lane)
+            for token in cutlass.range_constexpr(self.batch):
+                self.down_dots(packed, values, hq, token)
+                if (lane < 2) & (pair + lane < end_pair):
+                    word = self.lane_word(values, token, lane)
+                    self.store_word(output, par3, token, pair + lane, word)
+
+    @cute.jit
+    def shared_outputs(self, down, hsm, output, par3, first, end, warp, lane):
+        begin_pair, end_pair = (first + 1) // 2, end // 2
+        hq = [
+            [
+                load_f32x4(hsm.iterator + t * 384 + 4 * (lane + 32 * c), shared=True)
+                for c in range(3)
+            ]
+            for t in range(self.batch)
+        ]
+        packed = cute.make_rmem_tensor((4, 3, 2), Uint32)
+        values = cute.make_rmem_tensor((self.batch, 4), Float32)
+        saved = cute.make_rmem_tensor((3, self.batch), Uint32)
+        fits = (end_pair - begin_pair) <= 96
+        if cutlass.const_expr(self.batch == 2):
+            for iteration in cutlass.range_constexpr(3):
+                pair = begin_pair + 2 * warp + 32 * iteration
+                self.load_down(down, packed, pair, end_pair, lane, fits)
+                for token in cutlass.range_constexpr(2):
+                    self.down_dots(packed, values, hq, token)
+                    saved[iteration, token] = self.lane_word(values, token, lane)
+        if cutlass.const_expr(self.batch == 1) or not fits:
+            self.stream_pairs(down, hq, output, par3, begin_pair, end_pair, warp, lane)
+        boundary = cute.make_rmem_tensor(self.batch, Uint32)
+        if ((end & 1) != 0) & (warp == 0):
+            edge = end // 2
+            for row in cutlass.range_constexpr(2):
+                for chunk in cutlass.range_constexpr(3):
+                    words = early_weights(
+                        down.iterator + (2 * edge + row) * 384 + 4 * (lane + 32 * chunk)
+                    )
+                    packed[row, chunk, 0], packed[row, chunk, 1] = words[0], words[1]
+            for token in cutlass.range_constexpr(self.batch):
+                lo = self.shared_dot(packed[0, None, None], hq, token)
+                hi = self.shared_dot(packed[1, None, None], hq, token)
+                boundary[token] = bf16_word(lo, hi)
+        if cutlass.const_expr(self.batch == 2):
+            for iteration in cutlass.range_constexpr(3):
+                pair = begin_pair + 2 * warp + 32 * iteration
+                if fits & (lane < 2) & (pair + lane < end_pair):
+                    for token in cutlass.range_constexpr(2):
+                        word = saved[iteration, token]
+                        self.store_word(output, par3, token, pair + lane, word)
+        if ((end & 1) != 0) & (warp == 0) & (lane == 0):
+            for token in cutlass.range_constexpr(self.batch):
+                self.store_word(output, par3, token, end // 2, boundary[token])
+
+    @cute.kernel
+    def kernel(self, args):
+        router, bias, xmoe, epoch, scores, ids, weights, down, latent = args[0:9]
+        gate_up, hg, shared_down, output, designated, par, par3, rank = args[9:17]
+        multicast_base, scale, situ = args[17], args[18], args[19:23]
+        tid, warp, lane, cta = thread()
+        router_first, router_rows = split(896, cta, self.grid)
+        pair_first, pairs = split(112, cta, self.grid)
+        gate_first, gate_rows = Int32(0), Int32(0)
+        if cta >= self.batch:
+            part, parts = cta - self.batch, self.grid - self.batch
+            gate_first, gate_rows = split(768, part, parts)
+        slots = alloc(cutlass.utils.SmemAllocator(), Uint8, self.shared_bytes).iterator
+        top, size = 15 * 14336, 256 * self.batch
+        partials = view(slots, top, Float32, size)
+        latent_partial = view(slots, top + 4 * size, Float32, (16, 8), (8, 1))
+        gate_partial = view(slots, top + 4 * size + 512, Float32, size)
+        bars = cute.recast_ptr(slots + top + 8 * size + 512, dtype=Uint64)
+        roles = view(slots, top + 8 * size + 528, Int32, 2)
+        hsm = view(slots, 0, Float32, self.batch * 384)
+        gate_args = (slots, gate_rows, gate_first, bars, gate_partial, hg, par, tid)
+        cute.arch.griddepcontrol_launch_dependents()
+        active_lane = lane < 28
+        first_chunk = 56 * warp + (lane if active_lane else 0)
+        h2_weights = zeros((2, 4 if self.batch == 2 else 2, 4), Uint32)
+        for row in cutlass.range_constexpr(2):
+            for chunk in cutlass.range_constexpr(4 if self.batch == 2 else 2):
+                valid = active_lane & (pairs > 0)
+                index = first_chunk + 28 * chunk
+                if cutlass.const_expr(self.batch == 2):
+                    valid = (warp < 7) & (pairs > 0)
+                    index = 32 * warp + lane + 224 * chunk
+                if valid:
+                    pointer = down.iterator + (2 * pair_first + row) * 7168 + index * 8
+                    words = early_weights(pointer, count=4)
+                    for j in cutlass.range_constexpr(4):
+                        h2_weights[row, chunk, j] = words[j]
+        if tid == 0:
+            init_barriers(bars, 2, async_proxy=True)
+            claimed = self.placement.claim(epoch, designated, par, cta)
+            stage_rows(slots, 14336, 0, router, router_first, router_rows, bars)
+            if cta == 0:
+                mailbox_store(epoch.iterator + (par ^ 1) * 32, (Uint32(0),))
+                self.placement.rearm(epoch, par)
+            observed = mailbox_load(epoch.iterator + par * 32)[0]
+            while observed < 28 * self.batch:
+                sleep_ns(32)
+                observed = mailbox_load(epoch.iterator + par * 32)[0]
+            role, count = self.placement.role(epoch, par, cta, claimed, self.grid)
+            roles[0], roles[1] = role, count
+        else:
+            if (tid >= 32) & (tid < 32 + self.batch * router_rows):
+                offset = tid - 32
+                token, row = offset // router_rows, offset % router_rows
+                index = ((par ^ 1) * self.batch + token) * 896 + router_first + row
+                mailbox_store(scores.iterator + index * 2, (Uint32(0x80000000),) * 2)
+            if (tid >= 64) & (tid < 64 + self.batch * gate_rows):
+                offset = tid - 64
+                token, row = offset // max(gate_rows, 1), offset % max(gate_rows, 1)
+                index = ((par ^ 1) * self.batch + token) * 768 + gate_first + row
+                hg[index] = Uint32(0)
+            vectors: cutlass.Constexpr = xmoe.shape[1] * xmoe.shape[2] * 896
+            rearm = ((xmoe, vectors, vectors), (latent, 3584, self.batch * 448))
+            for tensor, stride, count in rearm:
+                for vector in range(cta * 511 + tid - 1, count, self.grid * 511):
+                    pointer = tensor.iterator + (par ^ 1) * stride * 8 + vector * 8
+                    mailbox_store(pointer, (Uint32(0x80000000),) * 4)
+        my_bias = Float32(0)
+        if tid < self.batch * router_rows:
+            my_bias = bias[router_first + tid % router_rows]
+        cute.arch.sync_threads()
+        role, count = roles[0], roles[1]
+        shared_resident = role >= 0
+        consumer = shared_resident & ((role & 0x10000) != 0)
+        dense = role & 0xFFFF
+        consumer_token = dense if self.batch == 2 else Int32(0)
+        shared_first, shared_end = Int32(0), Int32(0)
+        if shared_resident:
+            shared_first = 7168 * dense // count
+            shared_end = 7168 * (dense + 1) // count
+        h2 = (warp, lane, pairs) if self.batch == 2 else None
+        words_x, words_h2 = gather_x(
+            xmoe, par, cta, slots + 7 * 14336, first_chunk, active_lane, self.batch, h2
+        )
+        if cutlass.const_expr(self.batch == 1):
+            for row in cutlass.range_constexpr(2):
+                value = Float32(0)
+                if pairs > 0:
+                    for chunk in cutlass.range_constexpr(2):
+                        for j in cutlass.range_constexpr(4):
+                            w0, w1 = unpack_bf16(h2_weights[row, chunk, j])
+                            x0, x1 = unpack_bf16(words_x[0, chunk, j])
+                            value = fma(w0, x0, fma(w1, x1, value))
+                value = warp_sum(value, 32)
+                if lane == 0:
+                    latent_partial[warp, row] = value
+        elif (warp < 7) & (pairs > 0):
+            for token in cutlass.range_constexpr(2):
+                for row in cutlass.range_constexpr(2):
+                    value = Float32(0)
+                    for chunk in cutlass.range_constexpr(4):
+                        for j in cutlass.range_constexpr(4):
+                            w0, w1 = unpack_bf16(h2_weights[row, chunk, j])
+                            x0, x1 = unpack_bf16(words_h2[token, chunk, j])
+                            value = fma(w1, x1, fma(w0, x0, value))
+                    value = warp_sum(value, 32)
+                    if lane == 0:
+                        latent_partial[warp, 2 * token + row] = value
+        cute.arch.sync_warp()
+        if cutlass.const_expr(self.batch == 1):
+            if warp == 1:
+                cute.arch.barrier(barrier_id=3, number_of_threads=512)
+            else:
+                named_barrier(3, 512, arrive=True)
+        elif (warp == 1) & (pairs > 0):
+            cute.arch.barrier(barrier_id=3, number_of_threads=224)
+        elif (warp < 7) & (pairs > 0):
+            named_barrier(3, 224, arrive=True)
+        if (warp == 1) & (pairs > 0) & (lane < self.batch):
+            lo, hi = Float32(0), Float32(0)
+            for w in cutlass.range_constexpr(16 if self.batch == 1 else 7):
+                lo += latent_partial[w, 2 * lane]
+                hi += latent_partial[w, 2 * lane + 1]
+            word = bf16_word(lo, hi)
+            offset = par * 8 * 3584 + lane * 3584 + rank * 224 + 2 * pair_first
+            multicast_words(multicast_base + Uint64(offset * 2), (word,), kind="b32")
+        if warp == 1:
+            cute.arch.sync_warp()
+        if tid == 0:
+            if gate_rows > 0:
+                cute.arch.mbarrier_arrive_and_expect_tx(bars + 1, gate_rows * 14336)
+                dst, src = slots + 9 * 14336, gate_up.iterator + gate_first * 7168
+                copy_bulk(dst, src, bars + 1, gate_rows * 14336)
+            if shared_resident & (shared_end > shared_first):
+                src = shared_down.iterator + shared_first * 384
+                prefetch_bytes(src, (shared_end - shared_first) * 768)
+        wait(bars)
+        self.gemv(slots, router_rows, slots + 7 * 14336, partials, warp, lane)
+        cute.arch.sync_threads()
+        if tid < self.batch * router_rows:
+            token, row = tid // router_rows, tid % router_rows
+            logit = Float32(0)
+            for w in cutlass.range_constexpr(16):
+                logit += partials[(w * 16 + row) * self.batch + token]
+            score = sigmoid(logit)
+            record = score_record(score, score + my_bias)
+            index = (par * self.batch + token) * 896 + router_first + row
+            mailbox_store(scores.iterator + index * 2, record)
+        if (gate_rows > 0) & (~consumer):
+            self.gate_pass(*gate_args)
+        cute.arch.sync_threads()
+        if consumer & (warp < 8):
+            top16 = Top16Fast(self.batch)
+            selected = top16.run(scores, ids, slots, par, consumer_token)
+            if warp == 0:
+                score = selected[lane] if lane < 16 else Float32(0)
+                total = warp_sum(score, 16)
+                if lane < 16:
+                    normalized = score
+                    if cutlass.const_expr(self.renorm):
+                        # Like the TP8 front and vLLM's grouped_topk kernels.
+                        normalized = div_rn(score, total + Float32(1e-20))
+                    bits = Uint32(gate_bf16_bits(normalized * scale))
+                    if bits == 0x8000:
+                        bits = Uint32(0)
+                    high = cute.arch.shuffle_sync_down(bits, 1, mask=0xFFFF)
+                    if (lane & 1) == 0:
+                        slot = par3 * weights.shape[1] + consumer_token
+                        pointer = weights.iterator + slot * 16 + lane
+                        mailbox_store(pointer, (bits | (high << 16),))
+                cute.arch.sync_warp()
+        if consumer & (gate_rows > 0):
+            cute.arch.sync_threads()
+            self.gate_pass(*gate_args)
+        if shared_resident:
+            if consumer:
+                cute.arch.sync_threads()
+            for element in range(tid, self.batch * 384, 512):
+                token, column = element // 384, element % 384
+                gp = hg.iterator + (par * self.batch + token) * 768 + column
+                gate, up = self.poll_flag(gp), self.poll_flag(gp + 384)
+                result = shared_situ(bf16(gate, True), bf16(up, True), *situ)
+                hsm[element] = Float32(BFloat16(result))
+            cute.arch.sync_threads()
+            self.shared_outputs(
+                shared_down, hsm, output, par3, shared_first, shared_end, warp, lane
+            )
+            cute.arch.sync_threads()
+        # One-thread exit wait: completion then covers the o-proj consumer (and its
+        # mailbox re-arm) for later waits.
+        if (cta == 0) & (tid == 0):
+            cute.arch.griddepcontrol_wait()

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp8.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TP8 B1 small-batch front (full router staged in SMEM, x_moe polling, threshold top16,
+latent multicast and gather, shared gate/up) and the shared-expert finisher."""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import BFloat16, Float32, Int32, Uint8, Uint32, Uint64
+
+from ...common.primitives import (
+    alloc,
+    bf16,
+    copy_bulk,
+    copy_g2s,
+    div_rn,
+    early_weights,
+    fma,
+    gate_bf16_bits,
+    init_barriers,
+    launch,
+    load_f32x4,
+    mailbox_load,
+    mailbox_store,
+    mma_into,
+    multicast_words,
+    named_barrier,
+    published,
+    routing_add,
+    routing_sigmoid,
+    routing_weight,
+    sigmoid,
+    smem_words,
+    split,
+    store_gate_mailbox,
+    store_relaxed,
+    table_load,
+    thread,
+    unpack_bf16,
+    view,
+    wait,
+    warp_sum,
+    zeros,
+)
+from ..routing import Top16Fast, bf16_word, gather_x, score_record, stage_rows
+
+
+class MoeSharedTP8:
+    ring_bytes = 57344
+    row_bytes = 14400
+    lsu_rows = 92
+    shared_bytes = ring_bytes + row_bytes + lsu_rows * 1536 + 4096 + 24
+
+    def __init__(self, grid=52):
+        self.grid = grid
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, self.grid, 512, min_blocks_per_mp=1)
+
+    @cute.jit
+    def stage(self, dst, down, low, high, barrier):
+        cute.arch.mbarrier_arrive_and_expect_tx(barrier, (high - low) * 1536)
+        if high > low:
+            copy_bulk(dst, down.iterator + low * 768, barrier, (high - low) * 1536)
+
+    @cute.jit
+    def rows_out(self, base, low, start, high, hsm, output, par3, lane):
+        packed = cute.make_rmem_tensor((6, 2), Uint32)
+        for row in range(start, high, 16):
+            for chunk in cutlass.range_constexpr(6):
+                pointer = base + (row - low) * 1536 + (lane + 32 * chunk) * 8
+                values = smem_words(pointer, count=2)
+                packed[chunk, 0], packed[chunk, 1] = values[0], values[1]
+            self.row_out(row, packed, hsm, output, par3, lane)
+
+    @cute.jit
+    def row_out(self, row, packed, hsm, output, par3, lane):
+        value = Float32(0)
+        for chunk in cutlass.range_constexpr(6):
+            h = load_f32x4(hsm.iterator + 4 * (lane + 32 * chunk), shared=True)
+            w0, w1 = unpack_bf16(packed[chunk, 0])
+            w2, w3 = unpack_bf16(packed[chunk, 1])
+            value = fma(w0, h[0], fma(w1, h[1], fma(w2, h[2], fma(w3, h[3], value))))
+        value = warp_sum(value, 32)
+        if lane == 0:
+            pointer = output.iterator + par3 * output.shape[1] * 7168 + row
+            store_gate_mailbox(pointer, gate_bf16_bits(value))
+
+    @cute.kernel
+    def kernel(self, args):
+        xmoe, down, tables, hg, output, par, par3 = args[0:7]
+        tid, warp, lane, cta = thread()
+        first = 7168 * cta // self.grid
+        end = 7168 * (cta + 1) // self.grid
+        first_end = min(end, first + 37)
+        second_end = min(end, first_end + 9)
+        lsu_count = min(92, end - second_end)
+        storage = alloc(cutlass.utils.SmemAllocator(), Uint8, self.shared_bytes, 1024)
+        ring = storage.iterator
+        xslot = ring + self.ring_bytes
+        lsu = xslot + self.row_bytes
+        partial = lsu + self.lsu_rows * 1536
+        hsm = view(partial, 1024, Float32, 768)
+        bars = cute.recast_ptr(partial + 4096, dtype=Uint64)
+        cute.arch.griddepcontrol_launch_dependents()
+        if tid == 0:
+            init_barriers(bars, 3, async_proxy=True)
+            self.stage(ring, down, first, first_end, bars + 1)
+        active = lane < 28
+        gather_x(xmoe, par, cta, xslot, 56 * warp + (lane if active else 0), active)
+        cute.arch.sync_warp()
+        if lsu_count > 0:
+            for vector in range(tid, lsu_count * 96, 512):
+                dst = lsu + vector * 16
+                copy_g2s(dst, down.iterator + second_end * 768 + vector * 8)
+            cute.arch.cp_async_commit_group()
+        cute.arch.fence_proxy("async.shared", space="cta")
+        cute.arch.sync_threads()
+        if tid == 0:
+            cute.arch.fence_proxy("async.shared", space="cta")
+            self.stage(xslot, down, first_end, second_end, bars + 2)
+        if tid < 384:
+            column = 2 * tid
+            gate = up = (Uint32(0), Uint32(0))
+            valid = False
+            while not valid:
+                gate = mailbox_load(hg.iterator + par * 2 * 1536 + column, count=2)
+                up = mailbox_load(hg.iterator + par * 2 * 1536 + 768 + column, count=2)
+                valid = ((gate[0] & gate[1] & up[0] & up[1]) & Uint32(1)) != 0
+            ga0 = table_load(tables.iterator + (gate[0] >> 16))
+            ua0 = table_load(tables.iterator + 65536 + (up[0] >> 16))
+            ga1 = table_load(tables.iterator + (gate[1] >> 16))
+            ua1 = table_load(tables.iterator + 65536 + (up[1] >> 16))
+            hsm[column] = Float32(BFloat16(ga0 * ua0))
+            hsm[column + 1] = Float32(BFloat16(ga1 * ua1))
+        cute.arch.sync_threads()
+        out = (hsm, output, par3, lane)
+        if lsu_count > 0:
+            cute.arch.cp_async_wait_group(0)
+            cute.arch.sync_threads()
+            self.rows_out(
+                lsu, second_end, second_end + warp, second_end + lsu_count, *out
+            )
+        parts = ((ring, first, first_end), (xslot, first_end, second_end))
+        for part, (base, low, high) in enumerate(parts):
+            if high > low:
+                if warp == 0:
+                    wait(bars + 1 + part)
+                cute.arch.sync_threads()
+                self.rows_out(base, low, low + warp, high, *out)
+        if cta == 0:
+            cute.arch.griddepcontrol_wait()
+
+
+class SharedSituTables:
+    """Separate FP32 gate/up lookup tables over all BF16 patterns."""
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, 256, 256, pdl=False)
+
+    @cute.kernel
+    def kernel(self, args):
+        tables, beta, linear_beta = args[0:3]
+        tid, _, _, cta = thread()
+        index = 256 * cta + tid
+        value = bf16(Uint32(index))
+        gate = div_rn(value, beta)
+        probability = sigmoid(value)
+        tables[index] = (beta * cute.math.tanh(gate, fastmath=False)) * probability
+        up = value
+        if linear_beta > 0:
+            scaled = div_rn(value, linear_beta)
+            up = linear_beta * cute.math.tanh(scaled, fastmath=False)
+        tables[65536 + index] = up
+
+
+class MoeFrontTP8:
+    row_bytes = 14400
+    shared_bytes = 15 * row_bytes + (16 * 16 + 16 * 4) * 4 + 24
+
+    def __init__(self, grid=140, *, renorm=True):
+        self.grid = grid
+        self.renorm = renorm
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, self.grid, 512, min_blocks_per_mp=1)
+
+    @cute.jit
+    def gemv(self, rows, nrows, activation, partials, warp, lane):
+        group, quarter = lane // 4, lane % 4
+        acc = zeros((8, 4), Float32)
+        for iteration in cutlass.range_constexpr(14):
+            offset = (warp * 448 + iteration * 32 + 8 * quarter) * 2
+            lo = hi = x = (Uint32(0),) * 4
+            if group < nrows:
+                lo = smem_words(rows + group * self.row_bytes + offset)
+            if group + 8 < nrows:
+                hi = smem_words(rows + (group + 8) * self.row_bytes + offset)
+            if group < 1:
+                x = smem_words(activation + offset)
+            for half in cutlass.range_constexpr(2):
+                chain: cutlass.Constexpr = (2 * iteration + half) % 8
+                a = (lo[2 * half], hi[2 * half], lo[2 * half + 1], hi[2 * half + 1])
+                mma_into(acc[chain, None], a, (x[2 * half], x[2 * half + 1]), bf16=True)
+        if quarter == 0:
+            for half in cutlass.range_constexpr(2):
+                comp: cutlass.Constexpr = 2 * half
+                partials[warp * 16 + group + 8 * half] = (
+                    (acc[0, comp] + acc[1, comp]) + (acc[2, comp] + acc[3, comp])
+                ) + ((acc[4, comp] + acc[5, comp]) + (acc[6, comp] + acc[7, comp]))
+
+    @cute.jit
+    def gate_pass(self, rows, count, start, barrier, ctx):
+        activation, partials, hg, par, tid, warp, lane = ctx
+        wait(barrier)
+        self.gemv(rows, count, activation, partials, warp, lane)
+        cute.arch.sync_threads()
+        if tid < count:
+            value = Float32(0)
+            for w in cutlass.range_constexpr(16):
+                value += partials[w * 16 + tid]
+            bits = Uint32(gate_bf16_bits(value))
+            pointer = hg.iterator + (par * 2) * 1536 + start + tid
+            mailbox_store(pointer, ((bits << 16) | Uint32(1),))
+        cute.arch.sync_threads()
+
+    @cute.kernel
+    def kernel(self, args):
+        router, bias, xmoe, epoch, scores, ids, weights, down = args[:8]
+        latent, latent_out, gate_up, hg, par, par3, rank, multicast_base = args[8:16]
+        scale = args[16]
+        tid, warp, lane, cta = thread()
+        first_router, router_rows = split(896, cta, self.grid)
+        first_pair, pairs = split(224, cta, self.grid)
+        gu_start, gu_rows = Int32(0), Int32(0)
+        if cta > 0:
+            gu_start, gu_rows = split(1536, cta - 1, self.grid - 1)
+        gu_first = min(gu_rows, 7)
+        gu_second = gu_rows - gu_first
+        storage = alloc(cutlass.utils.SmemAllocator(), Uint8, self.shared_bytes, 128)
+        slots = storage.iterator
+        activation = slots + 7 * self.row_bytes
+        dst = (slots, self.row_bytes)
+        top = 15 * self.row_bytes
+        partials = view(slots, top, Float32, 256)
+        latent_partial = view(slots, top + 1024, Float32, (16, 4), (4, 1))
+        bars = cute.recast_ptr(slots + top + 1280, dtype=Uint64)
+        cute.arch.griddepcontrol_launch_dependents()
+        active_lane = lane < 28
+        first_chunk = 56 * warp + (lane if active_lane else 0)
+        second_chunk = first_chunk + 28
+        latent_weights = zeros((4, 2, 4), Uint32)
+        for row in cutlass.range_constexpr(4):
+            if active_lane & (row < 2 * pairs):
+                for chunk in cutlass.range_constexpr(2):
+                    col = first_chunk if chunk == 0 else second_chunk
+                    pointer = down.iterator + (2 * first_pair + row) * 7168 + 8 * col
+                    words = early_weights(pointer, count=4)
+                    for j in cutlass.range_constexpr(4):
+                        latent_weights[row, chunk, j] = words[j]
+        cute.arch.cp_async_commit_group()
+        if tid == 0:
+            init_barriers(bars, 3, async_proxy=True)
+            stage_rows(*dst, 0, router, first_router, router_rows, bars)
+            if cta == 0:
+                mailbox_store(epoch.iterator + (par ^ 1) * 32, (Uint32(0),))
+            observed = mailbox_load(epoch.iterator + par * 32)[0]
+            while observed < 28:
+                observed = mailbox_load(epoch.iterator + par * 32)[0]
+        else:
+            for i in range(tid - 1, 2 * router_rows, 511):
+                token, row = i // router_rows, i % router_rows
+                index = ((par ^ 1) * 2 + token) * 896 + first_router + row
+                mailbox_store(scores.iterator + index * 2, (Uint32(0x80000000),) * 2)
+            if cta > 0:
+                for i in range(tid - 1, 2 * gu_rows, 511):
+                    token, row = i // gu_rows, i % gu_rows
+                    hg[(par ^ 1) * 2 * 1536 + token * 1536 + gu_start + row] = Uint32(0)
+            slot_vectors: cutlass.Constexpr = xmoe.shape[1] * xmoe.shape[2] * 896
+            for i in range(cta * 511 + tid - 1, slot_vectors, self.grid * 511):
+                pointer = xmoe.iterator + (par ^ 1) * slot_vectors * 8 + i * 8
+                mailbox_store(pointer, (Uint32(0x80000000),) * 4)
+        my_bias = Float32(0)
+        if tid < router_rows:
+            my_bias = bias[first_router + tid]
+        cute.arch.sync_threads()
+        input_words, _ = gather_x(xmoe, par, cta, activation, first_chunk, active_lane)
+        for row in cutlass.range_constexpr(4):
+            if row < 2 * pairs:
+                total = Float32(0)
+                for chunk in cutlass.range_constexpr(2):
+                    for j in cutlass.range_constexpr(4):
+                        w0, w1 = unpack_bf16(latent_weights[row, chunk, j])
+                        x0, x1 = unpack_bf16(input_words[0, chunk, j])
+                        total = fma(w0, x0, fma(w1, x1, total))
+                total = warp_sum(total, 32)
+                if lane == 0:
+                    latent_partial[warp, row] = total
+        cute.arch.sync_warp()
+        if warp == 1:
+            cute.arch.barrier(barrier_id=3, number_of_threads=512)
+            if lane < pairs:
+                lo, hi = Float32(0), Float32(0)
+                for w in cutlass.range_constexpr(16):
+                    lo += latent_partial[w, 2 * lane]
+                    hi += latent_partial[w, 2 * lane + 1]
+                packed = bf16_word(lo, hi)
+                offset = par * 8 * 3584 * 2 + rank * 896 + (first_pair + lane) * 4
+                multicast_words(multicast_base + Uint64(offset), (packed,), kind="b32")
+        else:
+            named_barrier(3, 512, arrive=True)
+        if tid == 0:  # noqa: SIM102
+            if gu_first > 0:
+                stage_rows(*dst, 8, gate_up, gu_start, gu_first, bars + 1)
+        wait(bars)
+        self.gemv(slots, router_rows, activation, partials, warp, lane)
+        cute.arch.sync_threads()
+        if tid == 0:  # noqa: SIM102
+            if gu_second > 0:
+                cute.arch.fence_proxy("async.shared", space="cta")
+                stage_rows(*dst, 0, gate_up, gu_start, gu_second, bars + 2, gu_first)
+        if tid < router_rows:
+            logit = Float32(0)
+            for w in cutlass.range_constexpr(16):
+                logit += partials[w * 16 + tid]
+            score = routing_sigmoid(logit)
+            record = score_record(score, routing_add(score, my_bias))
+            index = par * 2 * 896 + first_router + tid
+            mailbox_store(scores.iterator + index * 2, record)
+        if cta == 0:
+            if warp < 8:
+                selected = Top16Fast(2).run(scores, ids, slots, par, 0)
+                if warp == 0:
+                    score = selected[lane] if lane < 16 else Float32(0)
+                    total = warp_sum(score, 16)
+                    if lane < 16:
+                        value = routing_weight(score, total, scale, renorm=self.renorm)
+                        pointer = weights.iterator + par3 * weights.shape[1] * 16 + lane
+                        store_gate_mailbox(pointer, gate_bf16_bits(value))
+                    cute.arch.sync_warp()
+                for vector in range(tid, 448, 256):
+                    pointer = latent.iterator + par * 8 * 3584 + vector * 8
+                    words = mailbox_load(pointer, count=4, scope="sys")
+                    ready = published(words)
+                    while not ready:
+                        words = mailbox_load(pointer, count=4, scope="sys")
+                        ready = published(words)
+                    mailbox_store(latent_out.iterator + vector * 8, words)
+                    store_relaxed(pointer, (Uint32(0x80000000),) * 4)
+            cute.arch.griddepcontrol_wait()
+        else:
+            cute.arch.sync_threads()
+            ctx = (activation, partials, hg, par, tid, warp, lane)
+            if gu_first > 0:
+                self.gate_pass(
+                    slots + 8 * self.row_bytes, gu_first, gu_start, bars + 1, ctx
+                )
+            if gu_second > 0:
+                self.gate_pass(slots, gu_second, gu_start + gu_first, bars + 2, ctx)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail_runtime.py

@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Routed MoE tail: one fused kernel finalizes this rank's top-16 rows, multicasts the
+latent partial and peer-stores the shared shard (three rotating parities), all-reduces
+and RMS-normalizes per K-slice, runs the K-split up projection and multicasts the output
+row into the next AttnRes mailbox. At TP8 B5+ and TP16 B9+ the experts publish the
+latent partial, WaveMoeTail.push the shared shard and WaveMoeTail does the rest. At
+``twoshot_batches`` the latent all-reduce is ``common.twoshot``: the experts push each
+partial chunk to its column owner, whose WaveMoeTail CTAs sum and multicast it."""
+
+from functools import cache
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import BFloat16, Float32, Int32, Int64, Uint8
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import fake, prepare_shared_carveout, stream, tensor_maps, to_cute
+from ..common.policy import batch_table
+from ..common.twoshot import TwoShot
+from ..oproj.prepared import OprojMailbox
+from .tail import MoeTail, WaveMoeTail
+
+# MoE tail mailbox tokens (padded decode batches); small-batch front_poll G/W/S rows per
+# parity.
+CAP, POLL_CAP = 64, 8
+# Two-shot latent all-reduce batches per TP (B9+): 2-/4-node GB300, us/layer of RS push
+# + experts + WaveMoeTail, max over ranks, multicast vs two-shot: TP8 B16 41.74/42.80,
+# B32 52.20/53.18, B64 79.27/81.50; TP16 B16 33.98/34.49, B32 42.19/42.52, B64
+# 65.05/60.51.
+TWOSHOT = {8: "", 16: "33-64:1"}
+# WaveMoeTail batches that release the next AttnRes at entry.
+EARLY = {8: "9-64:1", 16: "9-64:1"}
+# WaveMoeTail batches whose up projection runs on tcgen05 (B9+; 1-GPU WaveMoeTail us
+# mma.sync vs tcgen05, B16/32/64: TP8 4.53/5.80/9.05 vs 4.45/5.26/7.80, TP16 two-shot
+# 4.50/5.39/7.75 vs 4.52/5.46/7.21).
+UMMA = {8: "9-64:1", 16: "33-64:1"}
+
+
+def flat(t):
+    return to_cute(t.reshape(-1), 16)
+
+
+def published(tp, batch):
+    """The experts publish the latent and WaveMoeTail runs the tail (else MoeTail
+    front)."""
+    return batch > (4 if tp == 8 else 8)
+
+
+def twoshot_batches(tp):
+    table = batch_table(TWOSHOT.get(tp, ""))
+    return {batch for batch, on in table.items() if on and batch > 8}
+
+
+def early_batches(tp):
+    table = batch_table(EARLY.get(tp, ""))
+    return {batch for batch, on in table.items() if on and published(tp, batch)}
+
+
+def umma_batches(tp):
+    table = batch_table(UMMA.get(tp, ""))
+    return {batch for batch, on in table.items() if on and batch > 8}
+
+
+def shares(tp, batch, twoshot=False):
+    """WaveMoeTail clusters share reduced latent rows (all 112 CTAs must be
+    co-resident)."""
+    return published(tp, batch) and WaveMoeTail(batch, tp, twoshot).reduce
+
+
+@cache
+def resident_clusters(device):
+    """Co-resident 16-CTA clusters at one CTA per SM."""
+    with torch.accelerator.device_index(device):
+        return cutlass.utils.HardwareInfo(device).get_max_active_clusters(16)
+
+
+class MoeTailBox:
+    """Model-owned MoE tail mailboxes: latent [3, TP, CAP, 3584] (multicast), shared
+    reduce-scatter [3, TP, CAP, 7168/TP] (peer stores), output [1, CAP, 7168], the
+    local all-reduced latent rows [3, CAP, 3584] that WaveMoeTail clusters share and the
+    two-shot latent buffers (``twoshot_args``/``twoshot_tail`` replace args/tail)."""
+
+    def __init__(self, group, device, tp, rank):
+        self.rank = rank
+        self.latent = OprojMailbox(group, device, tp, rank, (3, tp, CAP, 3584))
+        self.rs = OprojMailbox(group, device, tp, rank, (3, tp, CAP, 7168 // tp))
+        self.out = OprojMailbox(group, device, tp, rank, (1, CAP, 7168))
+        ptrs = self.rs.handle.buffer_ptrs
+        self.peers = torch.tensor(ptrs, dtype=torch.int64, device=device)
+        self.args = (flat(self.latent.buffer), flat(self.rs.buffer))
+        self.tail: tuple = (to_cute(self.peers, 16), flat(self.out.buffer))
+        self.tail += (Int64(self.latent.mc), Int64(self.out.mc))
+        self.reduced = torch.empty((3, CAP, 3584), dtype=torch.bfloat16, device=device)
+        self.reduced.view(torch.int32).fill_(-0x80000000)
+        self.twoshot = box = TwoShot(group, device, tp, rank, 3584, CAP, parities=3)
+        self.twoshot_args = (flat(box.recv), self.args[1])
+        self.twoshot_tail = (*self.tail[:2], Int64(box.gather_mc), self.tail[3])
+        # Experts' push: every parity's receive slots [3 x TP, CAP, 3584/TP] and the
+        # owners' buffer bases (slot parity x TP + rank).
+        self.push_box = (
+            to_cute(box.recv.view(3 * tp, CAP, -1)),
+            to_cute(box.peers[0], 16),
+        )
+
+
+@cache
+def compile_tail(tp, batch, device, mode="front", trigger=2, early=False, umma=False):
+    """front: rows complete at the PDL wait; front_poll: three 0xffff-armed G/W/S parity
+    slots polled per row and re-armed after use; publish (``published``): WaveMoeTail
+    over the published mailboxes, reading the weight shard directly, not by TMA;
+    twoshot: publish over the two-shot receive slots and gather."""
+    m8, twoshot = mode in ("publish", "twoshot"), mode == "twoshot"
+    vec = lambda n, dtype=BFloat16, align=16: fake((n,), None, dtype, align)
+    slots = 3 * POLL_CAP if mode == "front_poll" else batch
+    latent = 3 * tp * CAP * (3584 // tp if twoshot else 3584)
+    tensors: tuple = (vec(latent), vec(3 * tp * CAP * (7168 // tp)))
+    weights = vec(7168 // tp * 3584) if m8 and not umma else vec(128, Uint8, 64)
+    tensors += (weights, vec(3584))
+    tensors += (vec(slots * 16 * 3584), vec(slots * 16), vec(slots * 7168))
+    tensors += (vec(tp, Int64), vec(CAP * 7168))
+    scalars: tuple = (Int64(0), Int64(0), Int32(0), Int32(0), Int32(CAP), Float32(1e-6))
+    scalars += (Int32(2), Int32(CAP))
+    kernel: MoeTail
+    if m8:
+        kernel = wave = WaveMoeTail(batch, tp, twoshot, early, umma)
+        scalars += (vec(3 * CAP * 3584),) if wave.reduce else ()
+    else:
+        kernel = MoeTail(tp, batch, mode=mode, trigger=trigger)
+    options = "" if m8 else "--ptxas-options=--maxrregcount=120"
+    with torch.accelerator.device_index(device):
+        args = (*tensors, *scalars, make_fake_stream())
+        code = cute.compile(kernel, *args, options=options)
+        # MoE tail requests 100% shared carveout so the next large PDL kernel can
+        # coexist.
+        return prepare_shared_carveout(code, device, 100)
+
+
+@cache
+def compile_push(batch, device, tp=8):
+    args = (fake((batch, 7168)), fake((tp,), None, Int64), Int32(0), Int32(0))
+    with torch.accelerator.device_index(device):
+        return cute.compile(
+            WaveMoeTail(batch, tp).push, *args, Int32(CAP), make_fake_stream()
+        )
+
+
+class MoeTailWeights:
+    """One layer's up-projection shard [7168/TP, 3584], RMS weight and TMA map."""
+
+    def __init__(self, tp, weight, gamma, eps, ordinal, count, batches, twoshots=()):
+        self.weight, self.gamma = weight, gamma
+        # Three parities rotate over the model's MoE layers; each launch re-arms the
+        # previous MoE layer's slots of its batch (wrapping across decode steps, where
+        # the first layer re-arms the largest batch the last layer may have published).
+        parity, rearm = ordinal % 3, (ordinal - 1) % count % 3
+        if parity == rearm:  # a launch must never re-arm the slots it publishes to
+            raise ValueError(f"MoE tail parities collide for {count} MoE layers")
+        self.parity, rearm = Int32(parity), Int32(rearm)
+        self.scalars = (Int32(CAP), Float32(eps), rearm)
+        self.wrap = 0 if ordinal else max(batches)
+        cols = 512 if tp == 8 else 448
+        u16 = cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16
+        dims = (cols // 2, 3584 // (cols // 2), 7168 // tp)
+        spec: tuple = (weight.data_ptr(), u16, dims, (cols, 7168), (cols // 2, 2, 32))
+        spec += (cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE,)
+        spec += (cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_NONE,)
+        self.descriptor = tensor_maps([spec], weight.device)
+        self.views = (to_cute(self.descriptor, 64), to_cute(gamma, 16))
+        self.codes, device, early = {}, weight.device.index, early_batches(tp)
+        umma = umma_batches(tp)
+        if any(published(tp, b) and b in umma for b in batches):
+            # WaveMoeTail tcgen05: [64 x lines, 64] boxes of the shard, 128B swizzle
+            spec = (weight.data_ptr(), u16, (3584, 7168 // tp), (7168,))
+            spec += ((64, 128 if tp == 8 else 64),)
+            spec += (cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,)
+            spec += (cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B,)
+            self.umma_map = tensor_maps([spec], weight.device)
+            self.umma_views = (to_cute(self.umma_map, 64), self.views[1])
+        for b in batches:  # WaveMoeTail reads the shard directly, not by TMA
+            twoshot = b in twoshots
+            mode = "twoshot" if twoshot else "publish" if published(tp, b) else "front"
+            views = self.views if mode == "front" else (flat(weight), self.views[1])
+            on = mode != "front" and b in umma
+            views = self.umma_views if on else views
+            code = compile_tail(tp, b, device, mode, early=b in early, umma=on)
+            self.codes[b] = code, views, shares(tp, b, twoshot), twoshot
+        shared = any(shares(tp, b, b in twoshots) for b in batches)
+        if shared and resident_clusters(device) < 7:
+            raise RuntimeError("WaveMoeTail needs seven co-resident 16-CTA clusters")
+
+    def run(self, box, expanded, weights, shared, code=None, batch=None):
+        batch = batch or weights.shape[0]
+        if code is None:
+            code, views, share, twoshot = self.codes[batch]
+        else:
+            views, share, twoshot = self.views, 0, False
+        args, tail = box.args, box.tail
+        if twoshot:
+            args, tail = box.twoshot_args, box.twoshot_tail
+        rows = (flat(expanded), flat(weights), flat(shared), *tail, self.parity)
+        rows += (Int32(box.rank), *self.scalars, Int32(self.wrap or batch))
+        if share:
+            rows += (flat(box.twoshot.gather if twoshot else box.reduced),)
+        code(*args, *views, *rows, stream(expanded.device))
+
+
+class PreparedTail:
+    """The layer's MoE tail; placement = (model MoeTailBox, MoE ordinal, MoE layer
+    count). At WaveMoeTail batches the experts publish into latent slot ``publish[b]``
+    (multicast, parity, rank; at ``twoshots`` the two-shot push box) and the side
+    stream pushes the shared-expert output (``push``)."""
+
+    def __init__(self, op, norm, up, batches, placement):
+        tp, rank = op.contract.tp_size, op.rank
+        local_up = up.narrow(0, rank * (7168 // tp), 7168 // tp)
+        self.tail_box, ordinal, count = placement
+        eps = op.contract.rms_eps
+        m8 = [b for b in batches if published(tp, b)]
+        self.twoshots = twoshot_batches(tp) & set(m8)
+        args = (tp, local_up, norm, eps, ordinal, count, batches, self.twoshots)
+        self.tail_weights = MoeTailWeights(*args)
+        self.peers, _, mc, _ = self.tail_box.tail
+        self.ring = (self.tail_weights.parity, Int32(rank))
+        self.pushers = {b: compile_push(b, norm.device.index, tp) for b in m8}
+        self.publish = {b: (mc, *self.ring) for b in m8}
+        slot = Int32(ordinal % 3 * tp + rank)
+        self.publish.update({b: (*self.tail_box.push_box, slot) for b in self.twoshots})
+
+    def push(self, shared):
+        if shared.shape[0] in self.pushers:
+            args = (to_cute(shared), self.peers, *self.ring, Int32(CAP))
+            self.pushers[shared.shape[0]](*args, stream(shared.device))
+
+    def run(self, expanded, weights, shared, out, *, consume):
+        self.tail_weights.run(self.tail_box, expanded, weights, shared)
+        if consume(out) is not out:
+            raise RuntimeError("The tail output must be consumed from its mailbox ring")
+        return out

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/consumer.py

@@ -0,0 +1,423 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""O-projection all-reduce consumer (``groups`` 128-thread groups poll disjoint
+rank slots, group 0 sums in rank order into the AttnRes) and TP16 fused producer +
+consumer. With ``twoshot`` the mailbox is the one-slot gather row of
+``common.twoshot``: the threads owning this rank's columns first reduce their
+receive slots into it (args 11-13: receive view, gather multicast address, rank)."""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Uint32
+
+from ..attn_res.kernel import arm_exchanges, rstd
+from ..common.mailbox import rearm
+from ..common.primitives import (
+    add_rn,
+    alloc,
+    asm,
+    at,
+    bits_float,
+    copy_g2s,
+    float_bits,
+    fma,
+    launch,
+    load_vector,
+    named_barrier,
+    pack_pair,
+    publish_hint,
+    published,
+    push_float,
+    smem_words,
+    store_relaxed,
+    store_shared_words,
+    store_words,
+    unpack_bf16,
+    zeros,
+)
+from ..common.primitives import wait_phase_zero_cluster as wait_phase_zero
+from ..common.twoshot import reduce
+from .producer import OprojProducer, pack8, tree_reduce
+
+
+@asm((Float32,) * 2, pure=True)
+def pair_op(a, b, c=(), *, op="add"):
+    values = (*a, *b, *c)
+    code = "{ .reg .b64 a,b,c,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; "
+    if op == "fma":
+        code += "mov.b64 c,{$6,$7}; fma.rn.f32x2 d,a,b,c; "
+    else:
+        code += f"{op}.rn.f32x2 d,a,b; "
+    constraints = ",".join(["=f", "=f"] + ["f"] * len(values))
+    return code + "mov.b64 {$0,$1},d; }", constraints, tuple(map(Float32, values))
+
+
+@asm(Float32, pure=True)
+def mul_rn(a, b):
+    return "mul.rn.f32 $0,$1,$2;", "=f,f,f", (*(Float32(v) for v in (a, b)),)
+
+
+def two(values, j):
+    return values[2 * j], values[2 * j + 1]
+
+
+@cute.jit
+def dot8(a, b):
+    s0 = pair_op(two(a, 0), two(b, 0), op="mul")
+    s1 = pair_op(two(a, 1), two(b, 1), op="mul")
+    s0 = pair_op(two(a, 2), two(b, 2), s0, op="fma")
+    s1 = pair_op(two(a, 3), two(b, 3), s1, op="fma")
+    s = pair_op(s0, s1)
+    return add_rn(s[0], s[1])
+
+
+@cute.jit
+def reduce_scatter(values, size: cutlass.Constexpr[int], lane):
+    """Warp-reduce ``size`` values; returns the index of the one this lane holds."""
+    for shift in cutlass.range_constexpr(4, -1, -1):
+        live: cutlass.Constexpr = max(1, size >> (4 - shift))
+        half: cutlass.Constexpr = max(1, live // 2)
+        if cutlass.const_expr(live == 1):
+            values[0] = add_rn(
+                values[0], cute.arch.shuffle_sync_bfly(values[0], 1 << shift)
+            )
+        else:
+            upper = (lane & (1 << shift)) != 0
+            for j in cutlass.range_constexpr(half):
+                send, keep = values[j + half], values[j]
+                if upper:
+                    send, keep = values[j], values[j + half]
+                values[j] = add_rn(keep, cute.arch.shuffle_sync_bfly(send, 1 << shift))
+    idx = cutlass.Int32(0)
+    for shift in cutlass.range_constexpr(4, -1, -1):
+        width: cutlass.Constexpr = size >> (4 - shift)
+        if cutlass.const_expr(width > 1):  # noqa: SIM102
+            if (lane & (1 << shift)) != 0:
+                idx += width // 2
+    return idx
+
+
+@cute.jit
+def fetch(raw, rank, pointer):
+    words = load_vector(pointer, cache="volatile")
+    for j in cutlass.range_constexpr(4):
+        raw[rank, j] = words[j]
+
+
+class OprojConsumer:
+    def __init__(
+        self,
+        nb,
+        *,
+        has_delta=True,
+        groups=4,
+        twoshot=False,
+    ):
+        if not 0 <= nb <= 8:
+            raise ValueError("Output-projection AttnRes supports zero to eight blocks")
+        self.nb, self.has_delta = nb, has_delta
+        self.groups, self.twoshot = groups, twoshot
+        self.kpre = 1 << max(0, (2 * nb - 1).bit_length())
+        self.npost = 4 if nb > 0 else 1
+
+    @cute.jit
+    def wait_exchange(self, barrier):
+        wait_phase_zero(barrier)  # peers complete it by st.async
+
+    @cute.jit
+    def allocate(self):
+        smem = cutlass.utils.SmemAllocator()
+        wpre = alloc(smem, Float32, (4, 16))
+        pre = alloc(smem, Float32, (7, 16))
+        blk = alloc(smem, Float32, 16)
+        post = alloc(smem, Float32, (7, 4, 4))
+        total = alloc(smem, Float32, 4)
+        bars = alloc(smem, cutlass.Uint64, 2, 8)
+        parts = alloc(smem, Float32, (max(1, self.groups - 1), 2, 128, 4))
+        lean = self.groups == 1  # above B8: rows for this nb only (co-residency)
+        pf = alloc(smem, Uint32, (4 + self.nb if lean else 12, 128, 4))
+        return wpre, pre, blk, post, total, bars, parts, pf
+
+    @cute.jit
+    def setup(self, bars):
+        arm_exchanges(bars, (2 * self.nb, 4 * self.npost))
+        cute.arch.cluster_arrive_relaxed()
+
+    @cute.jit
+    def prefetch(self, args, pf, rank, token):
+        _, prefix, blocks, norm, query, out_norm = args[:6]
+        tid, _, _ = cute.arch.thread_idx()
+        h = (rank * 128 + tid) * 8
+        if cutlass.const_expr(self.nb > 0):
+            copy_g2s(pf.iterator + tid * 4, norm.iterator + h)
+            copy_g2s(pf.iterator + 512 + tid * 4, query.iterator + h)
+        copy_g2s(pf.iterator + 1024 + tid * 4, out_norm.iterator + h)
+        if cutlass.const_expr(self.has_delta):
+            copy_g2s(pf.iterator + 1536 + tid * 4, at(prefix, token, h))
+        for block in cutlass.range_constexpr(self.nb):
+            dst = pf.iterator + (4 + block) * 512 + tid * 4
+            copy_g2s(dst, at(blocks, token, block, h))
+        cute.arch.cp_async_commit_group()
+
+    @cute.jit
+    def fragment(self, pf, row, tid):
+        words = smem_words(pf.iterator + (row * 128 + tid) * 4)
+        values = cute.make_rmem_tensor(8, Float32)
+        for j in cutlass.range_constexpr(4):
+            values[2 * j], values[2 * j + 1] = unpack_bf16(words[j])
+        return values
+
+    @cute.jit
+    def combined_weight(self, pf, tid):
+        a, b = self.fragment(pf, 0, tid), self.fragment(pf, 1, tid)
+        result = cute.make_rmem_tensor(8, Float32)
+        for j in cutlass.range_constexpr(4):
+            result[2 * j], result[2 * j + 1] = pair_op(two(a, j), two(b, j), op="mul")
+        return result
+
+    @cute.jit
+    def pre(self, state, rank, eps, mix, wait_groups: cutlass.Constexpr[int] = 0):
+        wpre, pre, blk, _, _, bars, _, pf = state
+        tid, _, _ = cute.arch.thread_idx()
+        warp, lane = tid // 32, tid % 32
+        maximum, exponent_sum, qpart = Float32(0), Float32(0), Float32(0)
+        cute.arch.cp_async_wait_group(wait_groups)
+        if cutlass.const_expr(self.nb > 0):
+            weight = self.combined_weight(pf, tid)
+            partial = zeros(self.kpre, Float32)
+            for block in cutlass.range_constexpr(self.nb):
+                value = self.fragment(pf, 4 + block, tid)
+                partial[block] = dot8(value, value)
+                partial[self.nb + block] = dot8(value, weight)
+            idx = reduce_scatter(partial, self.kpre, lane)
+            if ((lane & (32 // self.kpre - 1)) == 0) & (idx < 2 * self.nb):
+                wpre[warp, idx] = partial[0]
+            named_barrier()
+            cute.arch.cluster_wait()
+            if tid < 2 * self.nb:
+                pre_sum = Float32(0)
+                for w in cutlass.range_constexpr(4):
+                    pre_sum = add_rn(pre_sum, wpre[w, tid])
+                for peer in cutlass.range_constexpr(7):
+                    push_float(at(pre, rank, tid), pre_sum, bars.iterator, peer)
+            self.wait_exchange(bars.iterator)
+            if tid < 2 * self.nb:
+                sums = [Float32(0), Float32(0)]
+                for peer in cutlass.range_constexpr(7):
+                    sums[peer % 2] = add_rn(sums[peer % 2], pre[peer, tid])
+                blk[tid] = add_rn(*sums)
+            named_barrier()
+            logits = [
+                mul_rn(blk[self.nb + b], rstd(blk[b], eps)) for b in range(self.nb)
+            ]
+            maximum = tree_reduce(logits, max)
+            exp = [cute.math.exp(x - maximum, fastmath=True) for x in logits]
+            exponent_sum = tree_reduce(exp, add_rn)
+            for block in cutlass.range_constexpr(self.nb):
+                value = self.fragment(pf, 4 + block, tid)
+                for j in cutlass.range_constexpr(4):
+                    mix[2 * j], mix[2 * j + 1] = pair_op(
+                        (exp[block], exp[block]), two(value, j), two(mix, j), op="fma"
+                    )
+            qpart = dot8(mix, mix)
+        else:
+            cute.arch.cluster_wait()
+        return maximum, exponent_sum, qpart
+
+    def slots(self, tp):
+        """Rank slots each group polls (at least four)."""
+        return max(4, (tp + self.groups - 1) // self.groups)
+
+    @cute.jit
+    def poll_sum(self, mailbox, token, h, first, count):
+        n: cutlass.Constexpr = self.slots(mailbox.shape[0])
+        raw = zeros((n, 4), Uint32)
+        pending = Int32(0)
+        for rank in cutlass.range_constexpr(n):
+            if rank < count:
+                fetch(raw, rank, at(mailbox, first + rank, token, h))
+        for rank in cutlass.range_constexpr(n):
+            if rank < count:  # noqa: SIM102
+                if not published(raw[rank, j] for j in range(4)):
+                    pending |= 1 << rank
+        while pending != 0:
+            for rank in cutlass.range_constexpr(n):
+                if (pending & (1 << rank)) != 0:
+                    fetch(raw, rank, at(mailbox, first + rank, token, h))
+            for rank in cutlass.range_constexpr(n):
+                if (pending & (1 << rank)) != 0:  # noqa: SIM102
+                    if published(raw[rank, j] for j in range(4)):
+                        pending &= ~(1 << rank)
+        values = cute.make_rmem_tensor(8, Float32)
+        for j in cutlass.range_constexpr(4):
+            pair = unpack_bf16(raw[0, j])
+            for rank in cutlass.range_constexpr(1, n):
+                if rank < count:
+                    pair = pair_op(pair, unpack_bf16(raw[rank, j]))
+            values[2 * j], values[2 * j + 1] = pair[0], pair[1]
+        return values
+
+    @cute.jit
+    def publish(self, output, xmoe, hints, token, h, values):
+        tid, _, _ = cute.arch.thread_idx()
+        words = pack8(values)
+        store_words(at(output, token, h), words)
+        if cutlass.const_expr(xmoe is not None):
+            words = pack8(values, "latent")
+            for rep in cutlass.range_constexpr(xmoe.shape[0]):
+                store_relaxed(at(xmoe, rep, token, h), words, scope="gpu")
+            cute.arch.sync_warp()
+            if tid % 32 == 0:
+                publish_hint(hints.iterator)
+
+    @cute.jit
+    def post(self, args, state, block_mix, rank, token):
+        mailbox, prefix, _, _, _, _, output, xmoe, hints, eps, out_eps = args[:11]
+        _, _, _, post, total, bars, parts, pf = state
+        mix, maximum, exponent_sum, qpart = block_mix
+        thread, _, _ = cute.arch.thread_idx()
+        group, tid = thread // 128, thread % 128
+        warp, lane = tid // 32, tid % 32
+        h = (rank * 128 + tid) * 8
+        S: cutlass.Constexpr = (mailbox.shape[0] + self.groups - 1) // self.groups
+        everyone: cutlass.Constexpr = 128 * self.groups
+        if group > 0:
+            first = group * S
+            count = min(S, mailbox.shape[0] - first)
+            values = self.poll_sum(mailbox, token, h, first, count)
+            for half in cutlass.range_constexpr(2):
+                words = tuple(float_bits(values[half * 4 + i]) for i in range(4))
+                store_shared_words(at(parts, group - 1, half, tid, 0), words)
+            named_barrier(2, everyone, arrive=True)
+            for r in cutlass.range_constexpr(self.slots(mailbox.shape[0])):
+                if r < count:
+                    rearm(at(mailbox, first + r, token, h))
+        else:
+            if cutlass.const_expr(self.twoshot):
+                recv, gather_mc, me = args[11:14]
+                if h // recv.shape[2] == me:
+                    reduce(recv, mailbox, gather_mc, me, token, h)
+            if cutlass.const_expr(self.has_delta):
+                old = self.fragment(pf, 3, tid)
+            attn = self.poll_sum(mailbox, token, h, 0, min(S, mailbox.shape[0]))
+            if cutlass.const_expr(self.groups > 1):
+                named_barrier(2, everyone)
+            for other in cutlass.range_constexpr(self.groups - 1):
+                for half in cutlass.range_constexpr(2):
+                    words = smem_words(at(parts, other, half, tid, 0))
+                    for j in cutlass.range_constexpr(2):
+                        k: cutlass.Constexpr = half * 2 + j
+                        part = tuple(map(bits_float, two(words, j)))
+                        attn[2 * k], attn[2 * k + 1] = pair_op(two(attn, k), part)
+            prefix_value = cute.make_rmem_tensor(8, Float32)
+            for j in cutlass.range_constexpr(4):
+                rounded = unpack_bf16(pack_pair(*two(attn, j)))
+                if cutlass.const_expr(self.has_delta):
+                    rounded = pair_op(two(old, j), rounded)
+                    rounded = unpack_bf16(pack_pair(rounded[0], rounded[1]))
+                prefix_value[2 * j], prefix_value[2 * j + 1] = rounded[0], rounded[1]
+            partial = cute.make_rmem_tensor(self.npost, Float32)
+            if cutlass.const_expr(self.nb > 0):
+                weight = self.combined_weight(pf, tid)
+                partial[0], partial[1] = dot8(mix, prefix_value), qpart
+                partial[2] = dot8(prefix_value, prefix_value)
+                partial[3] = dot8(prefix_value, weight)
+            else:
+                partial[0] = dot8(prefix_value, prefix_value)
+            idx = reduce_scatter(partial, self.npost, lane)
+            if idx < self.npost:
+                for peer in range(lane & (32 // self.npost - 1), 7, 32 // self.npost):
+                    dst = at(post, rank, warp, idx)
+                    push_float(dst, partial[0], bars.iterator + 1, peer)
+            words = pack8(prefix_value)
+            store_words(at(prefix, token, h), words)
+            for r in cutlass.range_constexpr(min(S, mailbox.shape[0])):
+                rearm(at(mailbox, r, token, h))
+            if cutlass.const_expr(xmoe is not None):
+                cute.arch.fence_acq_rel_gpu()
+            self.wait_exchange(bars.iterator + 1)
+            if tid < self.npost:
+                sums = [Float32(0)] * 4
+                for peer in cutlass.range_constexpr(7):
+                    for w in cutlass.range_constexpr(4):
+                        sums[w] = add_rn(sums[w], post[peer, w, tid])
+                total[tid] = tree_reduce(sums, add_rn)
+            named_barrier()
+            norm = self.fragment(pf, 2, tid)
+            final = cute.make_rmem_tensor(8, Float32)
+            if cutlass.const_expr(self.nb > 0):
+                lp = mul_rn(total[3], rstd(total[2], eps))
+                mx = max(maximum, lp)
+                alpha = cute.math.exp(maximum - mx, fastmath=True)
+                ep = cute.math.exp(lp - mx, fastmath=True)
+                den = fma(alpha, exponent_sum, ep)
+                a2, t1 = mul_rn(alpha, alpha), total[1]
+                ae2 = mul_rn(mul_rn(2.0, alpha), ep)
+                quad = fma(a2, t1, fma(ae2, total[0], mul_rn(mul_rn(ep, ep), total[2])))
+                scale = rstd(quad, mul_rn(mul_rn(out_eps, den), den))
+                ca, cp = mul_rn(alpha, scale), mul_rn(ep, scale)
+            else:
+                cp = rstd(total[0], out_eps)
+            for j in cutlass.range_constexpr(4):
+                value = pair_op((cp, cp), two(prefix_value, j), op="mul")
+                if cutlass.const_expr(self.nb > 0):
+                    value = pair_op((ca, ca), two(mix, j), value, op="fma")
+                value = pair_op(value, two(norm, j), op="mul")
+                final[2 * j], final[2 * j + 1] = value[0], value[1]
+            self.publish(output, xmoe, hints, token, h, final)
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, (7, args[1].shape[0], 1), 128 * self.groups, 7)
+
+    @cute.kernel
+    def kernel(self, args):
+        rank, token, _ = cute.arch.block_idx()
+        tid, _, _ = cute.arch.thread_idx()
+        state = self.allocate()
+        self.setup(state[5])
+        cute.arch.griddepcontrol_launch_dependents()
+        mix = zeros(8, Float32)
+        maximum, exponent_sum, qpart = Float32(0), Float32(0), Float32(0)
+        if tid < 128:
+            self.prefetch(args, state[7], rank, token)
+            maximum, exponent_sum, qpart = self.pre(state, rank, args[9], mix)
+        self.post(args, state, (mix, maximum, exponent_sum, qpart), rank, token)
+        cute.arch.griddepcontrol_wait()
+
+
+class OprojFused:
+    def __init__(self, nb, *, has_delta=True):
+        self.producer = OprojProducer(768, staged_blocks=0)
+        # TP16 requires prefix/blocks to be ready at launch.
+        self.consumer = OprojConsumer(nb, has_delta=has_delta)
+
+    @cute.jit
+    def __call__(self, *args):
+        if cutlass.const_expr(args[0].shape[0] > 8):
+            raise ValueError("The selected fused schedule supports B1..8")
+        launch(self.kernel, args, (8, 14, 1), 512, 8, min_blocks_per_mp=1)
+
+    @cute.kernel
+    def kernel(self, args):
+        x, weight, rows = args[0], args[1], (args[2], *args[5:])
+        rank, cluster, _ = cute.arch.block_idx()
+        tid, _, _ = cute.arch.thread_idx()
+        cta = cluster * 8 + rank
+        active = (cluster < x.shape[0]) & (rank < 7)
+        state = self.consumer.allocate()
+        smem = self.producer.allocate(x.shape[0])
+        self.consumer.setup(state[5])
+        cute.arch.griddepcontrol_launch_dependents()
+        group0 = active & (tid < 128)
+        if group0:
+            self.consumer.prefetch(rows, state[7], rank, cluster)
+        weights = self.producer.load_weight(weight, smem[2], cta)
+        mix = zeros(8, Float32)
+        maximum, exponent_sum, qpart = Float32(0), Float32(0), Float32(0)
+        if group0:
+            maximum, exponent_sum, qpart = self.consumer.pre(state, rank, rows[9], mix)
+        self.producer.compute(args, smem, weights, cta, False)
+        if active:
+            block_mix = (mix, maximum, exponent_sum, qpart)
+            self.consumer.post(rows, state, block_mix, rank, cluster)

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Output projection fused with its TP all-reduce and post-attention AttnRes: rank
+partials go into Lamport slots of a multicast [TP, 64, 7168] mailbox (parity alternates
+by layer) and the consumer sums and re-arms them. TP16 B<=8 fuses producer and consumer
+(MLA layers below B5 keep them split); above B8 the producer runs 16-token groups.
+At ``twoshot_batches``, given ``TwoShot`` buffers, the all-reduce is ``common.twoshot``
+(chunks pushed to their column owner, owner sums + multicasts, all Lamport-polled), so
+each GPU receives ~2 x 14 KB per token instead of TP x 14 KB."""
+
+from functools import cache
+from typing import Any
+
+import cutlass
+import cutlass.cute as cute
+import torch
+import torch.distributed._symmetric_memory as symm_mem
+from cutlass.cute.runtime import make_fake_stream
+
+from ..common.host import dynamic_fake, dynamic_view, fake, stream, to_cute
+from ..common.policy import batch_table
+from .consumer import OprojConsumer, OprojFused
+from .producer import OprojProducer, OprojUmma
+
+
+class OprojMailbox:
+    """Sentinel-armed multicast mailbox; defaults to [parity, TP, 64, 7168] BF16."""
+
+    def __init__(self, group, device, tp, rank, shape=None):
+        self.tp, self.rank = tp, rank
+        shape = (2, tp, 64, 7168) if shape is None else shape
+        self.buffer = symm_mem.empty(shape, dtype=torch.bfloat16, device=device)
+        self.handle = handle = symm_mem.rendezvous(self.buffer, group)
+        self.mc = int(handle.multicast_ptr or 0)
+        if not self.mc:  # slot() offsets this address
+            raise RuntimeError(
+                "Kimi K3 specialized o-projection needs multicast memory"
+            )
+        self.buffer.view(torch.int32).fill_(-0x80000000)
+        torch.accelerator.synchronize(device)
+        torch.distributed.barrier(group=group)
+
+    def slot(self, parity):
+        return self.buffer[parity], self.mc + parity * self.buffer[0].numel() * 2
+
+
+# Two-shot all-reduce batches per TP (2-/4-node GB300, 32 MLA layers, us/layer, max over
+# ranks, multicast vs two-shot: TP8 B16 8.85/11.40, B32 10.63/11.32, B64 18.06/13.18;
+# TP16 B16 9.78/12.28, B32 14.96/13.05, B64 28.06/14.07).
+TWOSHOT = {8: "33-64:1", 16: "17-64:1"}
+
+
+def twoshot_batches(tp):
+    table = batch_table(TWOSHOT.get(tp, ""))
+    return {batch for batch, on in table.items() if on}
+
+
+def engine(batch, tp):
+    """Producer engine above B8 (GB300): grouped mma.sync for TP8 up to B16, else
+    tcgen05."""
+    default = "mma" if tp == 8 and batch <= 16 else "umma"
+    return default
+
+
+def producer(k, batch, tp, mla, twoshot=False, kind=None):
+    """B<=8: one mma.sync pass, TP8 stages 8 (MLA 4) of 12 K blocks per warp. Above
+    B8 (``kind`` overrides ``engine``): the tcgen05 engine, or 2-tile warps over 8 K
+    parts in 16-token groups, TP8 staging 2 of 6 K blocks per warp."""
+    if batch <= 8:
+        staged = 0 if tp == 16 else 4 if mla else 8
+        return OprojProducer(k, staged_blocks=staged, twoshot=twoshot)
+    if (kind or engine(batch, tp)) == "umma":
+        return OprojUmma(k, twoshot=twoshot)
+    grouped = dict(rows=2, group_tiles=2, twoshot=twoshot)
+    return OprojProducer(k, staged_blocks=0 if tp == 16 else 2, **grouped)
+
+
+@cache
+def compile_oproj(
+    kind, k, batch, tp, nb, has_delta, mla, front=False, cap=64, twoshot=False
+):
+    """Producer, consumer or fused kernel over a [TP, cap, 7168] parity view;
+    ``front``: the consumer also publishes x_moe [1, 4, 7168] and the epoch hint.
+    ``twoshot``: the two-shot pair (arguments as ``TwoShot.views``)."""
+    box = fake((1 if twoshot else tp, cap, 7168), (cap * 7168, 7168, 1))
+    x, w = fake((batch, k), (k, 1)), fake((7168, k), (k, 1))
+    mc, slot, eps = cutlass.Int64(0), cutlass.Int32(0), cutlass.Float32(1e-6)
+    recv, peers = fake((tp, cap, 7168 // tp)), fake((tp,), None, torch.int64)
+    if kind.startswith("producer"):  # producer_mma / producer_umma: engine A/B
+        kernel = producer(k, batch, tp, mla, twoshot, kind[9:] or None)
+        tail = (recv, peers, slot) if twoshot else (box, mc, slot)
+        return cute.compile(kernel, x, w, *tail, make_fake_stream())
+    row, weight = dynamic_fake(2), fake((7168,), (1,))
+    xmoe = fake((1, 4, 7168), (4 * 7168, 7168, 1)) if front else None
+    hints = fake((224,), (1,), torch.int32) if front else None
+    rows = (row, dynamic_fake(3), weight, weight, weight, row, xmoe, hints, eps, eps)
+    if kind == "consumer":  # B64: one 128-thread group per CTA polls every rank
+        options = dict(has_delta=has_delta)
+        groups = 4 if batch <= 32 and not twoshot else 1
+        kernel = OprojConsumer(nb, groups=groups, twoshot=twoshot, **options)
+        extra = (recv, mc, slot) if twoshot else ()
+        return cute.compile(kernel, box, *rows, *extra, make_fake_stream())
+    kernel = OprojFused(nb, has_delta=has_delta)
+    return cute.compile(kernel, x, w, box, mc, slot, *rows, make_fake_stream())
+
+
+class PreparedOproj:
+    """One decoder layer's o-projection + all-reduce + post-attention AttnRes."""
+
+    def __init__(self, decoder, mailbox, batches, mla, front=None, twoshot=None):
+        weight = decoder.self_attn.o_proj.weight
+        norm, out_norm = decoder.mlp_res_norm, decoder.post_attention_layernorm
+        weights = (norm.weight, decoder.mlp_res_proj.weight.squeeze(0), out_norm.weight)
+        self.views = tuple(to_cute(w.reshape(-1)) for w in weights)
+        self.eps = tuple(cutlass.Float32(n.variance_epsilon) for n in (norm, out_norm))
+        parity, rank = decoder.layer_idx % 2, cutlass.Int32(mailbox.rank)
+        local_box, mc = mailbox.slot(parity)
+        self.device, self.box = weight.device, to_cute(local_box)
+        self.head = (self.box, cutlass.Int64(mc), rank)
+        self.twoshot: Any = None
+        if twoshot is not None:
+            recv, peers, gather, gather_mc = twoshot.views(parity)
+            recv = to_cute(recv)
+            self.twoshot = (
+                (recv, to_cute(peers), rank),
+                to_cute(gather),
+                (recv, cutlass.Int64(gather_mc), rank),
+            )
+        self.wview, tp, k = to_cute(weight), mailbox.tp, weight.shape[1]
+        nb = decoder.prev_valid_blocks + decoder.is_block_write_layer
+        delta = not decoder.is_block_write_layer
+        self.front = None if front is None else (*map(to_cute, front[:2]), front[2])
+        self.codes, self.twoshots = {}, twoshot_batches(tp) if self.twoshot else set()
+        for b in batches:
+            publish = front is not None and b in front[2]
+            fused = tp == 16 and (not mla or b >= 5) and b <= 8
+            kinds = ("fused",) if fused else ("producer", "consumer")
+            twoshot = b in self.twoshots and not fused
+            shape = (k, b, tp, nb, delta, mla)
+            self.codes[b] = tuple(
+                compile_oproj(kind, *shape, publish and kind != "producer", 64, twoshot)
+                for kind in kinds
+            )
+
+    def run(self, local, prefix, blocks, out):
+        batch, s = local.shape[0], stream(self.device)
+        codes, x = self.codes[batch], to_cute(local)
+        front = (None, None)
+        if self.front is not None and batch in self.front[2]:
+            front = self.front[:2]
+        rows = (dynamic_view(prefix), dynamic_view(blocks), *self.views)
+        rows += (dynamic_view(out), *front, *self.eps, s)
+        if len(codes) == 1:
+            codes[0](x, self.wview, *self.head, *rows)
+        elif batch in self.twoshots:
+            head, gather, tail = self.twoshot
+            codes[0](x, self.wview, *head, s)
+            codes[1](gather, *rows[:-1], *tail, s)
+        else:
+            codes[0](x, self.wview, *self.head, s)
+            codes[1](self.box, *rows)
+        return out, prefix, blocks

vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/producer.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""O-projection producer for the split and fused AttnRes paths: 112 CTAs of 64 output
+rows. Each warp keeps ``rows`` 16-row weight tiles of one of ``4 * rows`` K parts in
+registers (TP8 stages some K blocks in shared memory), runs BF16 MMA chains and the K
+parts are reduced in order. Weights load before the PDL wait; the split producer
+triggers its consumer right after it. With ``group_tiles`` the batch runs in groups of
+that many 8-token tiles: each group's x arrives through a double-buffered cp.async ring
+during the previous group's MMAs, and is reduced and published during the next one's.
+With ``twoshot`` the partial chunks go to their owners' receive slots instead (the
+first shot of ``common.twoshot``; args 2-4: receive view, owners' bases, rank).
+``OprojUmma`` is the tcgen05 engine of
+the same producer for B>8."""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Constexpr, Float32, Int64, Uint8, Uint32, Uint64
+
+from ..common.primitives import (
+    add_rn,
+    alloc,
+    at,
+    bits_float,
+    copy_g2s,
+    fence_tmem,
+    launch,
+    load_tmem,
+    load_vector,
+    mma_bf16_block,
+    mma_into,
+    multicast_words,
+    named_barrier,
+    pack_pair,
+    prefetch_bytes,
+    release_tmem,
+    smem_desc,
+    smem_words,
+    stage_weight,
+    store_shared_words,
+    thread,
+    view,
+    wait,
+)
+from ..common.twoshot import push
+
+
+def pack8(values, policy="plain"):
+    """Four BF16x2 words of ``values[0:8]``."""
+    return tuple(
+        pack_pair(values[2 * j], values[2 * j + 1], policy=policy) for j in range(4)
+    )
+
+
+def tree_reduce(values, op):
+    tree = list(values)
+    for stride in (1, 2, 4):
+        for block in range(0, len(tree) - stride, 2 * stride):
+            tree[block] = op(tree[block], tree[block + stride])
+    return tree[0]
+
+
+class OprojProducer:
+    def __init__(self, k, *, staged_blocks, rows=1, group_tiles=0, twoshot=False):
+        self.k, self.staged, self.rows = k, staged_blocks, rows
+        self.twoshot = twoshot
+        self.parts = 4 * rows
+        self.kblocks = k // (32 * self.parts)
+        self.register_blocks = self.kblocks - self.staged
+        self.group_tiles = group_tiles
+        # Single pass: TP8 copies the batch to shared memory, TP16 loads registers.
+        self.shared_x = self.staged > 0 and not group_tiles
+        self.register_x = k == 768 and not group_tiles
+        self.x_stride = k // 8 + 4
+        self.chains = max(1, 4 // (rows * max(1, group_tiles)))
+
+    def schedule(self, batches):
+        """(8-token tiles, tiles per group, groups) of ``batches`` tokens."""
+        tiles = (batches + 7) // 8
+        per_group = self.group_tiles or tiles
+        return tiles, per_group, (tiles + per_group - 1) // per_group
+
+    @cute.jit
+    def allocate(self, batches: cutlass.Constexpr[int]):
+        nt, gt, ng = self.schedule(batches)
+        smem = cutlass.utils.SmemAllocator()
+        red = alloc(smem, Float32, (gt, 4, self.parts, 8, 20))
+        rows = nt * 8 if self.shared_x else 0
+        if cutlass.const_expr(self.group_tiles):
+            rows = min(2, ng) * gt * 8
+        xs = alloc(smem, Uint32, rows * self.x_stride * 4 if rows else 1)
+        staged = alloc(smem, Uint32, max(1, 16 * self.staged * self.rows * 2 * 32 * 4))
+        return red, xs, staged
+
+    @cute.jit
+    def load_weight(self, weight, staged, cta):
+        tid, _, _ = cute.arch.thread_idx()
+        lane, warp = tid % 32, tid // 32
+        group, part = lane // 4, lane % 4
+        tile = cta * 4 + warp // self.parts * self.rows
+        kq = warp % self.parts
+        base = (tile * 16 + group) * self.k + kq * self.kblocks * 32 + part * 8
+        if cutlass.const_expr(self.staged > 0):
+            for j in cutlass.range_constexpr(self.staged):
+                for t in cutlass.range_constexpr(self.rows):
+                    unit = (warp * self.staged + j) * self.rows + t
+                    dst = staged.iterator + (unit * 2 * 32 + lane) * 4
+                    src = weight.iterator + base + t * 16 * self.k
+                    src += (self.register_blocks + j) * 32
+                    stage_weight(dst, src)
+                    stage_weight(dst + 128, src + 8 * self.k)
+            cute.arch.cp_async_commit_group()
+        wr = cute.make_rmem_tensor((self.register_blocks, self.rows, 2, 4), Uint32)
+        for j in cutlass.range_constexpr(self.register_blocks):
+            for t in cutlass.range_constexpr(self.rows):
+                for r in cutlass.range_constexpr(2):
+                    src = weight.iterator + base + (t * 2 + r) * 8 * self.k + j * 32
+                    words = load_vector(src, cache="weight")
+                    for v in cutlass.range_constexpr(4):
+                        wr[j, t, r, v] = words[v]
+        return wr
+
+    @cute.jit
+    def stage_x(self, x, xs, first: Constexpr[int], buffer: Constexpr[int]):
+        """One cp.async group: x tokens [first, first + group) into ring ``buffer``."""
+        tid, _, _ = cute.arch.thread_idx()
+        chunks: Constexpr = self.k // 8
+        count: Constexpr = min(x.shape[0] - first, self.group_tiles * 8)
+        rows: Constexpr = buffer * self.group_tiles * 8
+        for i in cutlass.range_constexpr((count * chunks + 511) // 512):
+            chunk = tid + i * 512
+            token, col = chunk // chunks, chunk % chunks
+            dst = xs.iterator + ((rows + token) * self.x_stride + col) * 4
+            if cutlass.const_expr((i + 1) * 512 <= count * chunks):
+                copy_g2s(dst, at(x, first + token, col * 8))
+            else:
+                if chunk < count * chunks:  # noqa: SIM102
+                    copy_g2s(dst, at(x, first + token, col * 8))
+        cute.arch.cp_async_commit_group()
+
+    def x_words(self, xs, row, kofs, j):
+        return smem_words(xs.iterator + (row * self.x_stride + kofs // 8 + j * 4) * 4)
+
+    @cute.jit
+    def multiply(self, smem, wr, xr, first: Constexpr, buffer: Constexpr, M: Constexpr):
+        """FP32 chains of tokens [first, first + group); x from ``xr`` (registers) or
+        shared-memory rows (ring ``buffer``)."""
+        red, xs, staged = smem
+        _, gt, _ = self.schedule(M)
+        tiles: Constexpr = min(gt, (M - first + 7) // 8)
+        tid, _, _ = cute.arch.thread_idx()
+        lane, warp = tid % 32, tid // 32
+        group, part = lane // 4, lane % 4
+        kofs = warp % self.parts * self.kblocks * 32 + part * 8
+        acc = cute.make_rmem_tensor((gt, self.rows, self.chains, 4), Float32)
+        acc.fill(0)
+        for j in cutlass.range_constexpr(self.kblocks):
+            wa, wb = [], []
+            for t in cutlass.range_constexpr(self.rows):
+                if cutlass.const_expr(j < self.register_blocks):
+                    wa.append(tuple(wr[j, t, 0, v] for v in range(4)))
+                    wb.append(tuple(wr[j, t, 1, v] for v in range(4)))
+                else:
+                    sj: Constexpr = j - self.register_blocks
+                    unit = (warp * self.staged + sj) * self.rows + t
+                    ptr = staged.iterator + (unit * 2 * 32 + lane) * 4
+                    wa.append(smem_words(ptr))
+                    wb.append(smem_words(ptr + 128))
+            for nt in cutlass.range_constexpr(tiles):
+                if cutlass.const_expr(self.register_x):
+                    xv = tuple(xr[nt, j, v] for v in range(4))
+                else:
+                    row = buffer * gt * 8 + nt * 8
+                    if cutlass.const_expr(self.group_tiles):
+                        xv = self.x_words(xs, row + group, kofs, j)
+                    else:
+                        xv = (Uint32(0), Uint32(0), Uint32(0), Uint32(0))
+                        if nt * 8 + group < M:
+                            xv = self.x_words(xs, row + group, kofs, j)
+                for t in cutlass.range_constexpr(self.rows):
+                    for half in cutlass.range_constexpr(2):
+                        chain: Constexpr = (j * 2 + half) % self.chains
+                        a: tuple = (wa[t][2 * half], wb[t][2 * half])
+                        a += (wa[t][2 * half + 1], wb[t][2 * half + 1])
+                        b = (xv[2 * half], xv[2 * half + 1])
+                        mma_into(acc[nt, t, chain, None], a, b, bf16=True)
+        return acc
+
+    @cute.jit
+    def store_partials(self, red, acc, tiles: Constexpr):
+        tid, _, _ = cute.arch.thread_idx()
+        lane, warp = tid % 32, tid // 32
+        group, part = lane // 4, lane % 4
+        tile, kq = warp // self.parts * self.rows, warp % self.parts
+        for nt in cutlass.range_constexpr(tiles):
+            for t in cutlass.range_constexpr(self.rows):
+                for v in cutlass.range_constexpr(4):
+                    chains = [acc[nt, t, c, v] for c in range(self.chains)]
+                    token, row = 2 * part + v % 2, group + (v // 2) * 8
+                    red[nt, tile + t, kq, token, row] = tree_reduce(chains, add_rn)
+
+    @cute.jit
+    def publish(self, args, red, cta, first: Constexpr[int]):
+        """Sum the K parts of tokens [first, first + group) in order; multicast them."""
+        x, _, mailbox, mc, slot = args[:5]
+        _, gt, _ = self.schedule(x.shape[0])
+        count: Constexpr = min(x.shape[0] - first, gt * 8)
+        tid, _, _ = cute.arch.thread_idx()
+        if tid < 8 * count:
+            reduced = cute.make_rmem_tensor(8, Float32)
+            if cutlass.const_expr(self.group_tiles):
+                # Token-major lanes: conflict-free 16-byte partial loads.
+                token, out_half = tid % count, tid // count % 2
+                row_tile = tid // count // 2
+                tile, row = token // 8, token % 8
+                for i in cutlass.range_constexpr(2):
+                    column = out_half * 8 + i * 4
+                    words = smem_words(at(red, tile, row_tile, 0, row, column))
+                    sums = [bits_float(w) for w in words]
+                    for q in cutlass.range_constexpr(1, self.parts):
+                        words = smem_words(at(red, tile, row_tile, q, row, column))
+                        sums = [add_rn(s, bits_float(w)) for s, w in zip(sums, words)]
+                    for v in cutlass.range_constexpr(4):
+                        reduced[i * 4 + v] = sums[v]
+            else:
+                out_half, row_tile, token = tid % 2, (tid // 2) % 4, tid // 8
+                tile, row = token // 8, token % 8
+                for v in cutlass.range_constexpr(8):
+                    value = red[tile, row_tile, 0, row, out_half * 8 + v]
+                    for q in cutlass.range_constexpr(1, self.parts):
+                        part = red[tile, row_tile, q, row, out_half * 8 + v]
+                        value = add_rn(value, part)
+                    reduced[v] = value
+            words = pack8(reduced, "latent")
+            out_col = (cta * 4 + row_tile) * 16 + out_half * 8
+            if cutlass.const_expr(self.twoshot):
+                push(mailbox, mc, slot, first + token, out_col, words)
+            else:
+                offset = cute.crd2idx((slot, first + token, out_col), mailbox.layout)
+                multicast_words(mc + Int64(offset) * 2, words)
+
+    @cute.jit
+    def compute(self, args, smem, wr, cta, trigger_after_wait: cutlass.Constexpr[bool]):
+        x = args[0]
+        red, xs, staged = smem
+        M: Constexpr = x.shape[0]
+        NT, GT, NG = self.schedule(M)
+        tid, _, _ = cute.arch.thread_idx()
+        lane, warp = tid % 32, tid // 32
+        group, part = lane // 4, lane % 4
+        kofs = warp % self.parts * self.kblocks * 32 + part * 8
+        cute.arch.griddepcontrol_wait()
+        if cutlass.const_expr(trigger_after_wait):
+            cute.arch.griddepcontrol_launch_dependents()
+        xr = cute.make_rmem_tensor((NT, self.kblocks, 4), Uint32)
+        if cutlass.const_expr(self.group_tiles):
+            self.stage_x(x, xs, 0, 0)
+        if cutlass.const_expr(self.shared_x):
+            for chunk in range(tid, M * (self.k // 8), 512):
+                token, col = chunk // (self.k // 8), chunk % (self.k // 8)
+                words = load_vector(at(x, token, col * 8))
+                store_shared_words(
+                    xs.iterator + (token * self.x_stride + col) * 4, words
+                )
+            cute.arch.sync_threads()
+        if cutlass.const_expr(self.register_x):
+            xr.fill(0)
+            for nt in cutlass.range_constexpr(NT):
+                token = nt * 8 + group
+                if token < M:
+                    for j in cutlass.range_constexpr(self.kblocks):
+                        words = load_vector(at(x, token, kofs + j * 32))
+                        for v in cutlass.range_constexpr(4):
+                            xr[nt, j, v] = words[v]
+        if cutlass.const_expr(self.staged > 0 and not self.group_tiles):
+            cute.arch.cp_async_wait_group(0)
+        for g in cutlass.range_constexpr(NG):
+            if cutlass.const_expr(self.group_tiles):
+                # x(g) and group g - 1's partials are visible: stage x(g + 1), publish.
+                cute.arch.cp_async_wait_group(0)
+                cute.arch.sync_threads()
+                if cutlass.const_expr(g + 1 < NG):
+                    self.stage_x(x, xs, (g + 1) * GT * 8, (g + 1) % 2)
+                if cutlass.const_expr(g > 0):
+                    self.publish(args, red, cta, (g - 1) * GT * 8)
+            acc = self.multiply(smem, wr, xr, g * GT * 8, g % 2, M)
+            if cutlass.const_expr(self.group_tiles):
+                cute.arch.sync_threads()
+            self.store_partials(red, acc, min(GT, NT - g * GT))
+        cute.arch.sync_threads()
+        self.publish(args, red, cta, (NG - 1) * GT * 8)
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, 112, 512, min_blocks_per_mp=1)
+
+    @cute.kernel
+    def kernel(self, args):
+        cta, _, _ = cute.arch.block_idx()
+        smem = self.allocate(args[0].shape[0])
+        wr = self.load_weight(args[1], smem[2], cta)
+        self.compute(args, smem, wr, cta, True)
+
+
+class OprojUmma(OprojProducer):
+    """B>8 producer on tcgen05: 112 CTAs x 64 rows, tokens on N. Four loader warps
+    stream 128B-swizzled W [64, 64] and x [N, 64] K blocks through a cp.async ring
+    (W of the first stages before the PDL wait, the rest prefetched to L2), the MMA
+    warp accumulates D [64, N] FP32 in TMEM with M=64 kind::f16 MMAs, and the loader
+    warps round D to BF16 and publish it like ``OprojProducer``."""
+
+    def __init__(self, k, *, twoshot=False, budget=184 * 1024):
+        super().__init__(k, staged_blocks=0, twoshot=twoshot)
+        self.blocks, self.budget = k // 64, budget
+
+    def geometry(self, batches):
+        """(N, ring stage bytes, stages, TMEM columns)."""
+        n = (batches + 7) // 8 * 8
+        stage = 8192 + n * 128
+        stages = min(self.blocks, self.budget // stage)
+        return n, stage, stages, max(32, 1 << (n - 1).bit_length())
+
+    @cute.jit
+    def load(self, source, ring, rows: Constexpr, valid: Constexpr, row0, column):
+        """cp.async ``rows`` x 64 BF16 of ``source`` from (row0, column) into a K-major
+        128B-swizzled tile; rows >= ``valid`` are left alone. ``valid`` < 0: weights
+        (evict-first, every row)."""
+        tid, _, _ = cute.arch.thread_idx()
+        for i in cutlass.range_constexpr((rows * 8 + 127) // 128):
+            c = tid + 128 * i
+            row, chunk = c // 8, c % 8
+            dst = ring + row // 8 * 1024 + row % 8 * 128 + (chunk ^ row % 8) * 16
+            src = at(source, row0 + row, column + chunk * 8)
+            if cutlass.const_expr(valid < 0):
+                stage_weight(dst, src)
+            elif cutlass.const_expr((i + 1) * 128 <= valid * 8):
+                copy_g2s(dst, src)
+            else:
+                if c < valid * 8:  # noqa: SIM102
+                    copy_g2s(dst, src)
+
+    @cute.jit
+    def multiply(self, ring, bars, tmem, n: Constexpr, stage: Constexpr):
+        """The MMA warp: wait for each K block, issue its four K16 MMAs from the
+        elected lane and release its stage (the last block: the accumulator)."""
+        stages: Constexpr = self.geometry(n)[2]
+        # kind::f16 descriptor: FP32 D, BF16 A and B, both K-major, N, M = 64
+        instruction: Constexpr = 1 << 4 | 1 << 7 | 1 << 10 | n >> 3 << 17 | 4 << 24
+        for kb in cutlass.range_constexpr(self.blocks):
+            s: Constexpr = kb % stages
+            wait(bars.iterator + s, kb // stages % 2)
+            cute.arch.fence_proxy("async.shared", space="cta")
+            fence_tmem()
+            a = smem_desc(ring.iterator + s * stage, True)
+            b = smem_desc(ring.iterator + s * stage + 8192, True)
+            last: Constexpr = kb + 1 == self.blocks
+            release: Constexpr = last or kb + stages < self.blocks
+            barrier = bars.iterator + (2 * stages if last else stages + s)
+            mma_bf16_block(tmem, a, b, instruction, int(kb > 0), barrier, int(release))
+
+    @cute.jit
+    def epilogue(self, args, ring, tmem, cta, n: Constexpr):
+        """D rows 16 * warp + lane (lane < 16: M=64 TMEM layout) -> FP32 [token, row]
+        tile in the drained ring -> 8-row BF16 chunks to the mailbox / partials."""
+        x, _, mailbox, mc, slot = args[:5]
+        M: Constexpr = x.shape[0]
+        tid, warp, lane, _ = thread()
+        values = load_tmem(tmem + (32 * warp << 16), self.geometry(M)[3])
+        tile = view(ring.iterator, 0, Float32, (n, 68), (68, 1))  # [token, row]
+        if lane < 16:
+            for t in cutlass.range_constexpr(M):
+                tile[t, 16 * warp + lane] = values[t]
+        named_barrier(1, 128)
+        for i in cutlass.range_constexpr((M * 8 + 127) // 128):
+            c = tid + 128 * i
+            if c < M * 8:
+                token, part = c // 8, c % 8
+                reduced = cute.make_rmem_tensor(8, Float32)
+                for h in cutlass.range_constexpr(2):
+                    words = smem_words(at(tile, token, part * 8 + h * 4))
+                    for v in cutlass.range_constexpr(4):
+                        reduced[h * 4 + v] = bits_float(words[v])
+                words = pack8(reduced, "latent")
+                column = cta * 64 + part * 8
+                if cutlass.const_expr(self.twoshot):
+                    push(mailbox, mc, slot, token, column, words)
+                else:
+                    offset = cute.crd2idx((slot, token, column), mailbox.layout)
+                    multicast_words(mc + Int64(offset) * 2, words)
+
+    @cute.jit
+    def __call__(self, *args):
+        launch(self.kernel, args, 112, 160, min_blocks_per_mp=1)
+
+    @cute.kernel
+    def kernel(self, args):
+        x, weight = args[0], args[1]
+        n, stage, stages, columns = self.geometry(x.shape[0])
+        tid, warp, _, cta = thread()
+        smem = cutlass.utils.SmemAllocator()
+        ring = alloc(smem, Uint8, stages * stage, 1024)
+        bars = alloc(smem, Uint64, 2 * stages + 1, 8)
+        tmem_slot = alloc(smem, Uint32, 1)
+        if tid == 0:
+            for s in cutlass.range_constexpr(stages):
+                cute.arch.mbarrier_init(bars.iterator + s, 256)
+                cute.arch.mbarrier_init(bars.iterator + stages + s, 1)
+            cute.arch.mbarrier_init(bars.iterator + 2 * stages, 1)
+            cute.arch.mbarrier_init_fence()
+        if warp == 4:
+            cute.arch.alloc_tmem(columns, tmem_slot.iterator)
+            cute.arch.relinquish_tmem_alloc_permit()
+        fence_tmem(before=True)
+        cute.arch.sync_threads()
+        fence_tmem()
+        tmem = tmem_slot[0]
+        row0 = cta * 64
+        if warp < 4:  # weights do not depend on the previous kernel
+            if cutlass.const_expr(stages < self.blocks) and tid < 64:
+                rest = at(weight, row0 + tid, stages * 64)
+                prefetch_bytes(rest, (self.blocks - stages) * 128)
+            for s in cutlass.range_constexpr(stages):
+                self.load(weight, ring.iterator + s * stage, 64, -1, row0, s * 64)
+                cute.arch.cp_async_mbarrier_arrive_noinc(bars.iterator + s)
+        cute.arch.griddepcontrol_wait()
+        cute.arch.griddepcontrol_launch_dependents()
+        if warp < 4:
+            for kb in cutlass.range_constexpr(self.blocks):
+                slot: Constexpr = kb % stages
+                tile = ring.iterator + slot * stage
+                if cutlass.const_expr(kb >= stages):
+                    wait(bars.iterator + stages + slot, (kb // stages - 1) % 2)
+                    self.load(weight, tile, 64, -1, row0, kb * 64)
+                    cute.arch.cp_async_mbarrier_arrive_noinc(bars.iterator + slot)
+                self.load(x, tile + 8192, n, x.shape[0], 0, kb * 64)
+                cute.arch.cp_async_mbarrier_arrive_noinc(bars.iterator + slot)
+            wait(bars.iterator + 2 * stages, 0)
+            fence_tmem()
+            # The MMAs (async proxy) read the ring the epilogue now overwrites.
+            cute.arch.fence_proxy("async.shared", space="cta")
+            self.epilogue(args, ring, tmem, cta, n)
+        else:
+            self.multiply(ring, bars, tmem, n, stage)
+        release_tmem(tmem, columns, warp, 4)

vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi K3 specialized decode: a CuTe decode executor over the stock backbone, enabled
+by ``additional_config={"decode_backend": "latency_specialized"}``."""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DecodeConfig:
+    max_context: int
+    mla_core: str
+    dcp: int
+    # Tokens per request in a verify step (1 + draft tokens; 1 without speculation).
+    q: int
+    capture_sizes: tuple[int, ...]
+
+
+def specialized_decode_config(vllm_config) -> DecodeConfig | None:
+    """The deployment's specialized decode settings, None for the stock decode.
+    Raises when the specialized decode is requested for an unsupported deployment."""
+    extra = vllm_config.additional_config
+    backend = extra.get("decode_backend", "auto") if isinstance(extra, dict) else "auto"
+    if backend == "auto":
+        return None
+    if backend != "latency_specialized":
+        raise ValueError("Supported decode_backend values are auto/latency_specialized")
+    reason = unsupported(vllm_config)
+    if reason is not None:
+        raise ValueError(f"Kimi K3 specialized decode does not support {reason}")
+    spec = vllm_config.speculative_config
+    return DecodeConfig(
+        max_context=vllm_config.model_config.max_model_len,
+        mla_core=extra.get("mla_core", "auto"),
+        dcp=vllm_config.parallel_config.decode_context_parallel_size,
+        q=1 + spec.num_speculative_tokens if spec is not None else 1,
+        capture_sizes=tuple(
+            vllm_config.compilation_config.cudagraph_capture_sizes or ()
+        ),
+    )
+
+
+def kda_dtypes(vllm_config):
+    """(conv, recurrent) KDA state dtypes the deployment's cache config selects."""
+    from vllm.model_executor.layers.mamba.mamba_utils import MambaStateDtypeCalculator
+
+    cache = vllm_config.cache_config
+    return MambaStateDtypeCalculator.kda_state_dtype(
+        vllm_config.model_config.dtype,
+        cache.mamba_cache_dtype,
+        cache.mamba_ssm_cache_dtype,
+    )
+
+
+def unsupported(vllm_config) -> str | None:
+    import torch
+
+    from vllm import envs
+    from vllm.config import CompilationMode, CUDAGraphMode
+    from vllm.platforms import current_platform
+
+    parallel, spec = vllm_config.parallel_config, vllm_config.speculative_config
+    text = vllm_config.model_config.hf_text_config
+    moe = tuple(
+        getattr(text, name, None)
+        for name in (
+            "hidden_act",
+            "activation_situ_beta",
+            "activation_situ_linear_beta",
+            "num_experts",
+            "num_experts_per_token",
+            "num_expert_group",
+            "topk_group",
+            "moe_router_activation_func",
+        )
+    )
+    compilation = vllm_config.compilation_config
+    graphs = compilation.cudagraph_mode
+    checks = (
+        (not current_platform.is_device_capability_family(100), "this GPU (SM100/103)"),
+        (
+            moe != ("situ", 4.0, 25.0, 896, 16, 1, 1, "sigmoid"),
+            "MoE configs but the one-group sigmoid top-16 SiTU MoE of Kimi K3",
+        ),
+        (parallel.tensor_parallel_size not in (8, 16), "TP other than 8/16"),
+        (parallel.pipeline_parallel_size != 1, "pipeline parallelism"),
+        (parallel.enable_expert_parallel, "expert parallelism"),
+        (parallel.use_ubatching, "DBO"),
+        (
+            parallel.decode_context_parallel_size
+            not in (1, parallel.tensor_parallel_size),
+            "decode context parallel other than 1 or TP",
+        ),
+        (vllm_config.model_config.enable_sleep_mode, "sleep mode"),
+        (vllm_config.lora_config is not None, "LoRA"),
+        # The runner's ModelState hooks build the executor and prepare the KV layouts.
+        (not vllm_config.use_v2_model_runner, "the V1 model runner"),
+        # The decode loop runs in forward; a compiled forward would trace it.
+        (compilation.mode == CompilationMode.STOCK_TORCH_COMPILE, "torch.compile"),
+        # The loop runs only in FULL decode graphs; eager and PIECEWISE steps run stock.
+        (
+            graphs is None or graphs.decode_mode() != CUDAGraphMode.FULL,
+            "deployments without FULL decode CUDA graphs (e.g. enforce-eager)",
+        ),
+        (
+            spec is not None and spec.method != "dspark",
+            "speculative methods but DSpark",
+        ),
+        (
+            spec is not None and spec.enable_adaptive_verification,
+            "adaptive verification",
+        ),
+        # KdaSplit keeps a BF16 conv state and an FP32 recurrent state in place; it
+        # maintains neither other cache dtypes nor ReplaySSM / RecoverSSM records.
+        (
+            kda_dtypes(vllm_config) != (torch.bfloat16, torch.float32),
+            "KDA cache dtypes other than a BF16 conv state and an FP32 recurrent "
+            "state (mamba cache dtype)",
+        ),
+        (
+            vllm_config.cache_config.use_replayssm
+            or vllm_config.cache_config.use_kda_recoverssm,
+            "ReplaySSM / RecoverSSM (--use-replayssm)",
+        ),
+        # The verify kernels take at most 8 tokens per request (KdaSplit's barriers).
+        (
+            spec is not None and spec.num_speculative_tokens > 7,
+            "more than 7 speculative tokens",
+        ),
+        (
+            spec is not None and envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM,
+            "the AttnRes aux stream",
+        ),
+    )
+    return next((reason for failed, reason in checks if failed), None)

vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py

@@ -0,0 +1,280 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The specialized decode executor over the stock Kimi K3 backbone: one routing
+decision per step, then every layer's prepared providers in one loop."""
+
+from functools import partial
+from typing import Any, NamedTuple
+
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.distributed import get_tp_group
+from vllm.distributed import tensor_model_parallel_all_reduce as all_reduce
+from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.platforms import current_platform
+from vllm.utils.math_utils import round_up
+
+from ..cutedsl.attn_res.prepared import PreparedAttnRes
+from ..cutedsl.common.embedding import PreparedEmbedding
+from ..cutedsl.common.twoshot import TwoShot
+from ..cutedsl.moe.dense import PreparedDenseMLP
+from ..cutedsl.moe.experts_runtime import PersistentExperts
+from ..cutedsl.moe.front import SharedFront
+from ..cutedsl.moe.prepared import MoEScratch, PreparedMoE
+from ..cutedsl.moe.small_batch import chain as small_batch_chain
+from ..cutedsl.moe.tail_runtime import MoeTailBox, PreparedTail
+from ..cutedsl.oproj.prepared import OprojMailbox
+from .prefetch import PrefetchController
+from .providers import BATCHES, PREPARED_MAX, Attention, pack_front
+
+
+class Step(NamedTuple):
+    q: int  # tokens per request
+    metadata: list  # per layer
+    slot_mapping: Any
+
+
+def layer_name(module):
+    return getattr(module, "prefix", getattr(module, "layer_name", None))
+
+
+def decode_metadata(attention, context, q, b):
+    """One layer's metadata if this step is a pure decode (q = 1) or a uniform
+    verify step of b tokens, else None."""
+    md = context.attn_metadata
+    md = md.get(layer_name(attention)) if isinstance(md, dict) else None
+    if md is None:
+        return None
+    if q > 1:
+        return md if is_uniform_verify(md, q, b) else None
+    if md.num_prefills or getattr(md, "num_spec_decodes", 0):
+        return None
+    return md
+
+
+def is_uniform_verify(md, q, b):
+    """Whether every request of the b // q graph rows verifies q tokens (FULL graphs
+    pad whole requests: KDA stages NULL state rows, MLA zero-length decodes)."""
+    if md.num_prefills:
+        return False
+    if not hasattr(md, "num_spec_decodes"):
+        return md.num_decode_tokens == md.num_decodes * q == b
+    rows, tokens = md.spec_state_indices_tensor, md.num_spec_decode_tokens
+    full = rows is not None and rows.shape[0] * q == b
+    return full and not md.num_decodes and tokens == md.num_spec_decodes * q
+
+
+def caches_of(attention):
+    cache = getattr(attention, "kv_cache", ())
+    return [cache] if isinstance(cache, torch.Tensor) else list(cache)
+
+
+class SpecializedDecode:
+    """Built once the weights are loaded; owns every prepared provider, the
+    mailboxes and the decode buffers. Reads the backbone, never modifies it."""
+
+    def __init__(self, backbone, config):
+        self.backbone, self.buffers = backbone, {}
+        q = config.q
+        # Verify steps decode q tokens per request at M = requests * q tokens; FULL
+        # verify graphs round each capture size up to whole requests.
+        rounded = {round_up(s, q) for s in config.capture_sizes}
+        verify = {m for m in rounded if m <= PREPARED_MAX}
+        self.batches, self.q = tuple(sorted({*BATCHES, *verify})), q
+        self.queries = (1, q) if q > 1 else (1,)
+        # The longest sequence the deployment serves bounds the MLA split count.
+        self.max_context, self.mla_core = config.max_context, config.mla_core
+        # Decode context parallel over the TP ranks: MLA KV is sequence-sharded.
+        self.dcp, self.dcp_box, self.mla_layers = config.dcp, None, 0
+        group = get_tp_group()
+        self.tp, self.rank = group.world_size, group.rank_in_group
+        self.group = group.device_group
+        self.device = torch.device(current_platform.current_device())
+
+        layers = backbone.layers[backbone.start_layer : backbone.end_layer]
+        self.layers = layers = tuple(layers)
+        self.build_providers()
+        self.embedding = PreparedEmbedding(backbone.embed_tokens, self.tp, self.rank)
+        self.before = {
+            d.layer_idx: PreparedAttnRes.for_target(d, has_delta=d is not layers[0])
+            for d in layers
+        }
+        # The final AttnRes also applies the backbone's final norm.
+        self.final = PreparedAttnRes.for_target(backbone, final=True)
+        # The next (or final) AttnRes consumes each MoE output from its MoE tail output
+        # row.
+        self.fused_before = {}
+        followers = [self.before[f.layer_idx] for f in layers[1:]] + [self.final]
+        for layer, consumer in zip(layers, followers):
+            moe = self.feed_forward[layer.layer_idx]
+            if isinstance(moe, PreparedMoE):
+                consumer.bind_ring(moe.tail.tail_box.out.buffer)
+                self.fused_before[layer.layer_idx] = consumer
+        # The AttnRes entering each aux layer (or the final one) also stores the
+        # updated prefix: the stock aux stream (prefix sum plus the last MLP output).
+        self.aux_layers = tuple(backbone.aux_hidden_state_layers)
+        end = layers[-1].layer_idx + 1
+        for index in self.aux_layers:
+            (self.final if index == end else self.before[index]).bind_aux()
+        empty = partial(torch.empty, device=self.device, dtype=torch.bfloat16)
+        self.hidden = (empty((64, 7168)), empty((64, 7168)))
+        self.prefix = empty((64, 7168))
+        self.blocks = empty((64, backbone.num_attn_res_blocks, 7168))
+
+    def build_providers(self):
+        """Every layer's attention and feed-forward providers, the model's mailboxes
+        and the prefetch side stream. All collective allocations, weight packing and
+        kernel compiles finish here, before KV binding and CUDA graph capture."""
+        config, layers = self.backbone.config, self.layers
+        beta = config.activation_situ_beta
+        linear_beta = config.activation_situ_linear_beta
+        tp, rank, device, batches = self.tp, self.rank, self.device, self.batches
+        owner = self.group, device, tp, rank
+        packed = {
+            d.layer_idx: pack_front(d.mlp, rank, tp) for d in layers if d.is_moe_layer
+        }
+        prefetch = self.prefetch = PrefetchController(self, layers, packed)
+        state = MoEScratch(tp, device, batches)
+        mailbox, tail_box, moe_layers = (
+            OprojMailbox(*owner),
+            MoeTailBox(*owner),
+            list(packed),
+        )
+        twoshot = TwoShot(*owner)
+        front_box = OprojMailbox(*owner, (1, 64, 3584 + 2 * 896))
+        small_batch = small_batch_chain.SmallBatchBox(*owner)
+        self.attention: dict[int, Attention] = {}
+        self.feed_forward: dict[int, Any] = {}
+        for layer in layers:
+            index, module = layer.layer_idx, layer.mlp
+            boxes = (mailbox, small_batch, twoshot)
+            self.attention[index] = Attention(layer, self, prefetch, *boxes)
+            if not layer.is_moe_layer:
+                reduce = all_reduce if module.down_proj.reduce_results else None
+                self.feed_forward[index] = PreparedDenseMLP(module, batches, reduce)
+                continue
+            args = (module, packed[index], front_box, state.front, batches, beta)
+            shared = dict(plan=state.plan, counters=state.shared, q=self.q)
+            front = SharedFront(*args, linear_beta, **shared)
+            routed = module.experts.routed_experts
+            w13, w2 = (
+                w.view(torch.uint8) for w in (routed.w13_weight, routed.w2_weight)
+            )
+            weights = w13, routed.w13_weight_scale, w2, routed.w2_weight_scale
+            experts = PersistentExperts(*weights, beta=beta, linear_beta=linear_beta)
+            norm = module.routed_expert_norm.weight
+            up = module.routed_expert_up_proj.weight
+            op = module.experts._k3_latent_moe_tail_op
+            placement = tail_box, moe_layers.index(index), len(moe_layers)
+            tail = PreparedTail(op, norm, up, batches, placement)
+            moe = PreparedMoE(
+                front, experts, tail, state, prefetch=prefetch, layer_idx=index
+            )
+            moe.small_batch = small_batch_chain.SmallBatchChain(
+                small_batch, module, experts, tail, index, beta, linear_beta, prefetch
+            )
+            self.feed_forward[index] = moe
+
+    def prepare_kv(self):
+        """Compile each layer's cores for its bound KV cache layout (the profiling
+        cache, then the real one)."""
+        for layer in self.layers:
+            caches = caches_of(layer.self_attn)
+            self.attention[layer.layer_idx].prepare_layout(caches)
+
+    def embed_input_ids(self, ids):
+        rows = ids.shape[0] if ids.ndim == 1 and ids.is_contiguous() else 0
+        if not 1 <= rows <= PREPARED_MAX:
+            return self.backbone.embed_input_ids(ids)
+        table = self.backbone.embed_tokens.weight
+        return self.embedding.run(ids, table.new_empty((ids.shape[0], table.shape[1])))
+
+    def plan(self, positions, intermediate_tensors):
+        """This step's decode plan, or None to run the stock forward instead.
+
+        Steps that run the stock forward:
+        - pipeline-parallel stage inputs (``intermediate_tensors``);
+        - token counts without prepared providers (above 64 or not in ``batches``);
+        - steps outside a CUDA graph capture (profiling, warmup and the warmup run
+          before each capture), so every stock kernel runs before its capture;
+        - PIECEWISE steps, keyed by token count alone and so captured as stock;
+        - steps with prompt rows (prefills, chunked prefills, fresh prompts);
+        - speculative steps whose requests do not all verify q tokens (partial
+          drafts, verify rows mixed with plain decode rows).
+
+        Every predicate is the same on all TP ranks: the mailboxes of the specialized
+        decode need every rank to take the same path.
+        """
+        b = positions.shape[0]
+        if intermediate_tensors is not None or b not in self.batches:
+            return None
+        if not torch.cuda.is_current_stream_capturing():
+            return None
+        if not is_forward_context_available():
+            return None
+        context = get_forward_context()
+        if context.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE:
+            return None
+        by_layer = context.attn_metadata
+        rows: list[Any] = list(by_layer.values()) if isinstance(by_layer, dict) else []
+        q = self.q if any(getattr(md, "num_spec_decodes", 0) for md in rows) else 1
+        if b % q:
+            return None
+        metadata = [decode_metadata(d.self_attn, context, q, b) for d in self.layers]
+        if any(md is None for md in metadata):
+            return None
+        return Step(q, metadata, context.slot_mapping)
+
+    def run(self, step, input_ids, inputs_embeds):
+        """Normalized hidden states (and the aux streams) of a planned step."""
+        x = inputs_embeds
+        if x is None:
+            x = self.embed_input_ids(input_ids)
+        out = torch.empty_like(x)
+        # The stock aux streams: the AttnRes prefix entering each aux layer.
+        aux_out = [torch.empty_like(x) for _ in self.aux_layers]
+        aux = dict(zip(self.aux_layers, aux_out))
+        b, q, slots = x.shape[0], step.q, step.slot_mapping
+        prefix, blocks = self.prefix[:b], self.blocks[:b]
+        first, second = self.hidden[0][:b], self.hidden[1][:b]
+        hidden, pre_ready = x, False
+        for layer, md in zip(self.layers, step.metadata):
+            index, caches = layer.layer_idx, caches_of(layer.self_attn)
+            if not pre_ready:
+                self.before[index].run(hidden, prefix, blocks, first, aux.get(index))
+            self.attention[index].run(
+                first, second, prefix, blocks, caches, md, slots, q
+            )
+            consumer = self.fused_before.get(index)
+            pre_ready = consumer is not None
+            if consumer is not None:
+                target = out if consumer is self.final else first
+                stream = aux.get(index + 1)
+                consume = partial(consumer.run_ring, prefix, blocks, aux=stream)
+                hidden = self.feed_forward[index].run(second, target, consume=consume)
+            else:
+                hidden = self.feed_forward[index].run(second, first)
+                first, second = second, first
+        if not pre_ready:
+            self.final.run(hidden, prefix, blocks, out, aux.get(index + 1))
+        self.prefetch.join()
+        return (out, aux_out) if aux_out else out
+
+    def normalize(self, output):
+        """A stock step's output with the backbone's final norm applied, as the
+        specialized decode's last AttnRes applies it; the aux streams stay pre-norm.
+
+        Runs in chunks of at most 255 rows: the stock RMSNorm kernel reduces with
+        1024-thread blocks below 256 rows and 256-thread blocks above, so chunking keeps
+        the reduction order ``compute_logits`` uses for the few sampled rows, and
+        prefill logits match stock.
+        """
+        hidden, aux = output if isinstance(output, tuple) else (output, None)
+        norm = self.backbone.norm
+        if hidden.ndim == 2 and hidden.shape[0] >= 256:
+            chunks = range(0, hidden.shape[0], 255)
+            hidden = torch.cat([norm(hidden[i : i + 255], None) for i in chunks])
+        else:
+            hidden = norm(hidden, None)
+        return hidden if aux is None else (hidden, aux)

vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""L2 weight prefetch for the specialized decode.
+
+Small-batch decode is bound by HBM bandwidth, and each kernel that starts on weights
+not yet in L2 stalls on DRAM. Some kernels leave bandwidth idle, such as the attention
+core and the all-reduces. At fixed points of each layer, this module launches a small
+kernel (``cutedsl/common/prefetch.py``) on a side stream. It pulls the byte ranges of
+the weights the next kernels read into L2, so that they start warm.
+
+Prefetch points (``fork(point, layer, batch)``), each with its batches per TP:
+
+- ``oproj_down``: after the o-projection (B3-8), the MoE front's packed
+  router + latent-down rows.
+- ``attention_moe``: at TP8 only, after the KDA input projection (B5-8) or the MLA
+  core (B8), the same front rows plus the shared expert's weights.
+- ``next_attention``: once the routed experts are launched, the next layer's
+  attention weights, capped in MiB per batch (``NEXT_ATTENTION``). Not above B8,
+  where the HBM-bound experts lose more bandwidth than the attention gains.
+
+The front and shared-expert ranges are on the critical path. They launch with the
+highest stream priority and the maximum shared-memory carveout, which needs a driver
+launch because CuTe drops both attributes.
+
+Each fork records an event on the main stream that the side stream waits on, so a
+prefetch starts only after the work it overlaps. ``join()`` at the end of the step
+makes the main stream wait for the side stream, so a captured graph always rejoins
+it. This lives in ``runtime/`` rather than ``cutedsl/`` because it is host-side
+orchestration across layers: a stream, events, and per-(point, layer, batch) tables
+built from the model's layers.
+"""
+
+import ctypes
+from functools import partial
+
+import cuda.bindings.driver as cuda
+import torch
+
+from ...ops.cute_dsl.latent_moe_tail.primitives import to_cute
+from ..cutedsl.common.host import _checked, resolve_kernel
+from ..cutedsl.common.policy import batch_table
+from ..cutedsl.common.prefetch import compile_kernel
+
+# Per-TP batch -> MiB cap (or cap.carveout) of the next layer's attention weights,
+# prefetched once the routed experts finish.
+# Not at B>8: the HBM-bound experts lose more bandwidth than the attention gains.
+NEXT_ATTENTION = {8: "1:96.100,2-3:96,4-7:64,8:32", 16: "1-8:1024"}
+
+
+def attention_weights(attention):
+    """Next-attention weights in their consumption order."""
+    names = (
+        *("in_proj_qkvgfab", "fused_qkv_a_g_proj", "fused_qkv_a_proj", "q_proj"),
+        *("kv_a_proj_with_mqa", "f_b_proj", "q_b_proj", "g_proj"),
+    )
+    weights = [getattr(getattr(attention, n, None), "weight", None) for n in names]
+    weights += [getattr(attention, "W_UK_T", None), getattr(attention, "W_UV", None)]
+    weights = [w for w in weights if isinstance(w, torch.Tensor)]
+    return (*weights, attention.o_proj.weight)
+
+
+class PreparedWeightRanges:
+    """One owner retains weights, immutable range table, executable and launch ABI."""
+
+    def __init__(self, weights, *, byte_cap=0, grid=32, carveout=None, priority=False):
+        self.weights = weights = tuple(weights)
+        device = weights[0].device
+        ranges, remaining = [], byte_cap
+        for weight in weights:
+            if weight.numel() == 0:
+                continue
+            if weight.data_ptr() % 16:  # the kernel walks 16-byte aligned chunks
+                raise ValueError("L2 prefetch ranges must start 16-byte aligned")
+            span = 1 + sum((n - 1) * s for n, s in zip(weight.shape, weight.stride()))
+            size = span * weight.element_size()
+            if byte_cap:
+                size = min(size, remaining) // 16 * 16
+                remaining -= size
+            if size:
+                ranges.append((weight.data_ptr(), size))
+            if byte_cap and remaining < 16:
+                break
+        self.table = torch.tensor(ranges, dtype=torch.int64, device=device)
+        self.compiled = compile_kernel(len(ranges), grid, carveout, device.index)
+        self.config = None
+        if not priority and carveout is None:
+            self.views = to_cute(self.table, 16)
+            return
+        # Launch priority / carveout need a driver launch (CuTe drops both).
+        self.executor, self.function = resolve_kernel(self.compiled, device.index)
+        ids, self.attributes = cuda.CUlaunchAttributeID, []
+        if priority:  # the greatest stream priority
+            top = cuda.CUlaunchAttribute()
+            top.id = ids.CU_LAUNCH_ATTRIBUTE_PRIORITY
+            top.value.priority = _checked(cuda.cuCtxGetStreamPriorityRange())[1]
+            self.attributes.append(top)
+        if carveout is not None:
+            smem = cuda.CUlaunchAttribute()
+            smem.id = ids.CU_LAUNCH_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT
+            smem.value.sharedMemCarveout = carveout
+            self.attributes.append(smem)
+        config = self.config = cuda.CUlaunchConfig()
+        config.gridDimX, config.gridDimY, config.gridDimZ = grid, 1, 1
+        config.blockDimX, config.blockDimY, config.blockDimZ = 32, 1, 1
+        config.sharedMemBytes, config.attrs = 0, self.attributes
+        config.numAttrs = len(self.attributes)
+        self.pointer = ctypes.c_void_p(self.table.data_ptr())
+        self.arguments = (ctypes.c_void_p * 1)(ctypes.addressof(self.pointer))
+        self.address = ctypes.addressof(self.arguments)
+
+    def launch(self, stream):
+        if self.config is None:
+            self.compiled(self.views, stream)
+        else:
+            self.config.hStream = stream
+            _checked(cuda.cuLaunchKernelEx(self.config, self.function, self.address, 0))
+
+
+class PrefetchController:
+    """Prefetch launches per (point, layer, batch) on one side stream."""
+
+    def __init__(self, runtime, layers, packed):
+        self.device, tp = runtime.device, runtime.tp
+        self.stream = torch.cuda.Stream(device=self.device)
+        self.events, self.pending, self.tables = {}, False, {}
+        urgent = partial(PreparedWeightRanges, carveout=100, priority=True)
+        following = batch_table(NEXT_ATTENTION[tp])
+        for position, layer in enumerate(layers):
+            if not layer.is_moe_layer:
+                continue
+            index, moe, attention = layer.layer_idx, layer.mlp, layer.self_attn
+            # The front GEMM reads [router shard; down shard] from the packed copy.
+            ranges = urgent((packed[index],), grid=8)
+            self._add("oproj_down", index, range(3, 9), ranges)
+            if tp == 8:
+                shared = moe.shared_experts
+                shared = shared.gate_up_proj.weight, shared.down_proj.weight
+                ranges = urgent((packed[index], *shared), grid=32)
+                batches = range(5, 9) if hasattr(attention, "in_proj_qkvgfab") else (8,)
+                self._add("attention_moe", index, batches, ranges)
+            if position + 1 < len(layers):
+                weights = attention_weights(layers[position + 1].self_attn)
+                self._add_capped("next_attention", index, weights, following)
+
+    def _add_capped(self, point, layer, weights, table):
+        groups: dict[tuple, list[int]] = {}
+        for batch, spec in table.items():
+            cap, carveout = spec if isinstance(spec, tuple) else (spec, None)
+            if cap:
+                groups.setdefault((cap, carveout), []).append(batch)
+        for (cap, carveout), batches in groups.items():
+            ranges = PreparedWeightRanges(
+                weights, byte_cap=cap * 1024**2, grid=32, carveout=carveout
+            )
+            self._add(point, layer, batches, ranges)
+
+    def _add(self, point, layer, batches, ranges):
+        key = point, layer
+        if key not in self.events:
+            event = torch.cuda.Event()
+            event.record(torch.cuda.current_stream(self.device))
+            self.events[key] = event
+        for batch in batches:
+            self.tables[point, layer, batch] = ranges
+
+    def fork(self, point, layer, batch):
+        ranges = self.tables.get((point, layer, batch))
+        if ranges is None:
+            return False
+        event = self.events[point, layer]
+        event.record(torch.cuda.current_stream(self.device))
+        self.stream.wait_event(event)
+        self.pending = True
+        ranges.launch(cuda.CUstream(self.stream.cuda_stream))
+        return True
+
+    def join(self):
+        if self.pending:
+            torch.cuda.current_stream(self.device).wait_stream(self.stream)
+            self.pending = False

vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-layer attention providers (input projection, KDA/MLA core, o-projection) and
+the helpers ``SpecializedDecode.build_providers`` uses."""
+
+from typing import Any
+
+import torch
+
+from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
+    SkinnyGemmConfig,
+    shape_dynamic_skinny_gemm,
+)
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+
+from ..cutedsl.gemm.projection import PreparedProjection
+from ..cutedsl.kda.prepared import PreparedKda
+from ..cutedsl.mla.dcp import DcpBox
+from ..cutedsl.mla.prepared import PreparedMla
+from ..cutedsl.oproj.prepared import PreparedOproj
+
+# Decode batches (tokens) served by the prepared providers: every batch up to 8, the
+# padded CUDA-graph batches and each further captured size; others run stock.
+BATCHES = (*range(1, 9), 16, 32, 64)
+PREPARED_MAX = BATCHES[-1]
+
+
+def projection_policies(tp, rows, batches):
+    """Split-K (M tile, N tile, split, stages) per execution rows' next tested bin."""
+    policies = {}
+    for batch in batches:
+        n = next(b for b in (8, 16, 32, 64) if batch <= b)
+        if rows == 6288 and (tp == 8 or n > 8):
+            tactic = (64, 64, 1, 6) if n == 64 else (128, n, 2, 5 if n == 32 else 6)
+        else:  # TP16 KDA B64: a sixth stage, -0.7 us per layer in the KDA chain
+            stages = {8: 8, 16: 10, 32: 8, 64: 6 if rows == 3216 else 5}[n]
+            tactic = (64, n, 4 if n == 8 else 2, stages)
+        policies[batch] = tactic
+    return policies
+
+
+class SkinnyProjection:
+    """B1/B2 skinny GEMM; static K loads two weight tiles before the PDL wait."""
+
+    def __init__(self, weight, fallback):
+        self.weight, self.fallback, k = weight, fallback, weight.shape[1]
+        tactics = {1: (128, 4, 1), 2: (64, 4, 2)}
+        self.configs = {m: SkinnyGemmConfig(m, *t, 8, k) for m, t in tactics.items()}
+        for m, cfg in self.configs.items():  # compile before capture
+            shape_dynamic_skinny_gemm(weight.new_zeros(m, k), weight, cfg, None)
+
+    def run(self, x, out):
+        cfg = self.configs.get(x.shape[0])
+        if cfg is None:
+            return self.fallback.run(x, out)
+        return shape_dynamic_skinny_gemm(x, self.weight, cfg, None)
+
+
+def scratch(runtime, name, shape, dtype=torch.bfloat16):
+    if name not in runtime.buffers:
+        runtime.buffers[name] = torch.empty(shape, device=runtime.device, dtype=dtype)
+    return runtime.buffers[name]
+
+
+class Attention:
+    """Input projection -> KDA/MLA core -> o-projection with its fused AttnRes."""
+
+    def __init__(self, decoder, runtime, prefetch, mailbox, small_batch, twoshot=None):
+        attention, tp = decoder.self_attn, runtime.tp
+        self.kind = kind = "kda" if hasattr(attention, "in_proj_qkvgfab") else "mla"
+        name = "in_proj_qkvgfab" if kind == "kda" else "fused_qkv_a_g_proj"
+        weight = getattr(attention, name).weight
+        pitch = weight.shape[0]
+        policies = projection_policies(tp, pitch, runtime.batches)
+        split_k = PreparedProjection(weight, policies).prepare()
+        self.projection = SkinnyProjection(weight, split_k)
+        self.core: Any  # PreparedKda or PreparedMla
+        if kind == "kda":
+            mode = lambda b: (1 if b == 1 else 2 if b <= 8 else 0) if tp == 8 else 0
+            policies = {b: (pitch, mode(b)) for b in runtime.batches}
+            q = runtime.q  # verify steps: B requests x q tokens
+            verify = [b // q for b in runtime.batches if q > 1 and b % q == 0]
+            policies.update({(b, q): (pitch, int(tp == 8)) for b in verify})
+            self.core = PreparedKda(attention, policies, is_conv_state_dim_first())
+        else:
+            heads = attention.num_local_heads
+            query = scratch(runtime, "mla.query", (64, heads, 576), torch.uint8)
+            policy = dict(max_context=runtime.max_context, core=runtime.mla_core)
+            policy.update(queries=runtime.queries)
+            if runtime.dcp > 1:
+                policy.update(
+                    dcp=dcp_box(runtime, heads), parity=runtime.mla_layers % 2
+                )
+            runtime.mla_layers += 1
+            pitches = {b: pitch for b in runtime.batches}
+            self.core = PreparedMla(attention, pitches, query, **policy)
+        parity = decoder.layer_idx % 2
+        front: tuple | None = None
+        if decoder.is_moe_layer:
+            front = small_batch.front_views[parity], small_batch.hint_views[parity]
+            front += (small_batch.batches,)
+        self.output = PreparedOproj(
+            decoder, mailbox, runtime.batches, kind == "mla", front, twoshot
+        )
+        projected = scratch(runtime, kind + ".projected", (64, pitch))
+        local = scratch(runtime, "attention.local", (64, 12288 // tp))
+        self.prefetch, self.layer = prefetch, decoder
+        self.storage = {b: (projected[:b], local[:b]) for b in range(1, 65)}
+
+    def prepare_layout(self, caches):
+        if self.kind == "kda":
+            self.core.prepare_caches(caches)
+        else:
+            self.core.prepare_cache(caches[0])
+
+    def run(self, x, out, prefix, blocks, caches, metadata, slot_mapping, q=1):
+        b = x.shape[0]
+        projected, local = self.storage[b]
+        projected = self.projection.run(x, projected)
+        if self.kind == "kda":
+            self.prefetch.fork("attention_moe", self.layer.layer_idx, b)
+            if q > 1:  # uniform q-token verify rows
+                slots = metadata.spec_state_indices_tensor
+                accepted = metadata.num_accepted_tokens
+                self.core.run(
+                    projected, local, caches=caches, slots=slots, accepted=accepted
+                )
+            else:
+                slots = metadata.non_spec_state_indices_tensor
+                self.core.run(projected, local, caches=caches, slots=slots)
+        else:
+            decode = metadata.decode
+            tables, lengths = decode.block_table, decode.seq_lens
+            # Request metadata shorter than b: the core runs the prepared batch b
+            # with the missing rows as empty requests.
+            rows = min(b, q * min(lengths.numel(), tables.shape[0]))
+            if rows < b:
+                local[rows:].zero_()
+            slots = slot_mapping[self.layer.self_attn.layer_name]
+            mla = dict(cache=caches[0], block_table=tables, lengths=lengths)
+            # DCP: lengths are the local shard's, totals the full sequences'.
+            mla.update(totals=decode.dcp_tot_seq_lens)
+            self.core.run(projected, local, slots=slots, q=q, **mla)
+            self.prefetch.fork("attention_moe", self.layer.layer_idx, b)
+        self.output.run(local, prefix, blocks, out)
+        self.prefetch.fork("oproj_down", self.layer.layer_idx, b)
+        return out
+
+
+def dcp_box(runtime, heads):
+    """The model's DCP mailboxes (one set for every MLA layer)."""
+    if runtime.dcp_box is None:
+        from vllm.distributed import get_dcp_group
+
+        group = get_dcp_group()
+        assert group.world_size == runtime.dcp == runtime.tp
+        from vllm.config import get_current_vllm_config
+
+        args = group.device_group, runtime.device, group.world_size, group.rank_in_group
+        parallel = get_current_vllm_config().parallel_config
+        interleave = parallel.cp_kv_cache_interleave_size
+        runtime.dcp_box = DcpBox(*args, heads, interleave=interleave)
+    return runtime.dcp_box
+
+
+def pack_front(moe, rank, tp):
+    """[this rank's router rows (padded to 16); this rank's latent down shard]."""
+    router, down = moe.gate.weight, moe.routed_expert_down_proj.weight
+    rows, shard = 896 // tp, 3584 // tp
+    router = router[rank * rows : (rank + 1) * rows]
+    down = down[rank * shard : (rank + 1) * shard]
+    return torch.cat((router, router.new_zeros((-rows % 16, 7168)), down)).detach()