From 642c18d64c4cb6db3ded8fbe1786718cd4005fcb Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Thu, 15 Oct 2026 22:50:41 -0400 Subject: [PATCH] [Perf] Kimi K3 Specialized Decode: CuTeDSL Optimized GB300 (1.5-2x) (#61867) Signed-off-by: Benjamin Chislett Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- vllm/models/kimi_k3/nvidia/model.py | 51 +- vllm/models/kimi_k3/nvidia/model_state.py | 26 + .../nvidia/specialized_decode/__init__.py | 2 + .../specialized_decode/cutedsl/__init__.py | 2 + .../cutedsl/attn_res/__init__.py | 2 + .../cutedsl/attn_res/kernel.py | 276 +++ .../cutedsl/attn_res/prepared.py | 88 + .../cutedsl/common/__init__.py | 2 + .../cutedsl/common/embedding.py | 79 + .../specialized_decode/cutedsl/common/host.py | 126 ++ .../cutedsl/common/mailbox.py | 42 + .../cutedsl/common/policy.py | 18 + .../cutedsl/common/prefetch.py | 63 + .../cutedsl/common/primitives.py | 928 ++++++++++ .../cutedsl/common/twoshot.py | 115 ++ .../cutedsl/gemm/__init__.py | 2 + .../cutedsl/gemm/projection.py | 106 ++ .../specialized_decode/cutedsl/gemm/splitk.py | 380 ++++ .../cutedsl/gemm/ws_gemm.py | 293 +++ .../cutedsl/kda/__init__.py | 2 + .../specialized_decode/cutedsl/kda/kernel.py | 715 ++++++++ .../cutedsl/kda/prepared.py | 112 ++ .../cutedsl/mla/__init__.py | 2 + .../specialized_decode/cutedsl/mla/dcp.py | 378 ++++ .../cutedsl/mla/prepared.py | 255 +++ .../specialized_decode/cutedsl/mla/query.py | 474 +++++ .../specialized_decode/cutedsl/mla/value.py | 95 + .../cutedsl/moe/__init__.py | 2 + .../specialized_decode/cutedsl/moe/dense.py | 86 + .../specialized_decode/cutedsl/moe/experts.py | 1579 +++++++++++++++++ .../cutedsl/moe/experts_runtime.py | 207 +++ .../specialized_decode/cutedsl/moe/front.py | 538 ++++++ .../cutedsl/moe/prepared.py | 103 ++ .../specialized_decode/cutedsl/moe/routing.py | 254 +++ .../specialized_decode/cutedsl/moe/shared.py | 684 +++++++ .../cutedsl/moe/small_batch/__init__.py | 2 + .../cutedsl/moe/small_batch/chain.py | 168 ++ .../cutedsl/moe/small_batch/experts.py | 541 ++++++ .../cutedsl/moe/small_batch/front_tp16.py | 580 ++++++ .../cutedsl/moe/small_batch/front_tp8.py | 361 ++++ .../specialized_decode/cutedsl/moe/tail.py | 1031 +++++++++++ .../cutedsl/moe/tail_runtime.py | 241 +++ .../cutedsl/oproj/__init__.py | 2 + .../cutedsl/oproj/consumer.py | 423 +++++ .../cutedsl/oproj/prepared.py | 162 ++ .../cutedsl/oproj/producer.py | 443 +++++ .../specialized_decode/runtime/__init__.py | 135 ++ .../specialized_decode/runtime/decode.py | 280 +++ .../specialized_decode/runtime/prefetch.py | 180 ++ .../specialized_decode/runtime/providers.py | 172 ++ 50 files changed, 12805 insertions(+), 3 deletions(-) create mode 100644 vllm/models/kimi_k3/nvidia/model_state.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/kernel.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/prepared.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/embedding.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/policy.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/prefetch.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/twoshot.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/projection.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/splitk.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/ws_gemm.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/kernel.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/dcp.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/prepared.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/query.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/value.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/front.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/routing.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/shared.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/chain.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/experts.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp8.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail_runtime.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/consumer.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/producer.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 7bf54f35d4..fca74b4624 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/model_state.py b/vllm/models/kimi_k3/nvidia/model_state.py new file mode 100644 index 0000000000..d1d0e489b5 --- /dev/null +++ b/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() diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/kernel.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/kernel.py new file mode 100644 index 0000000000..5af01ab6fd --- /dev/null +++ b/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))) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/prepared.py new file mode 100644 index 0000000000..36c65a77c0 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/embedding.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/embedding.py new file mode 100644 index 0000000000..eebf56b41a --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py new file mode 100644 index 0000000000..2833c3a37d --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py new file mode 100644 index 0000000000..75da9e77c7 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/policy.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/policy.py new file mode 100644 index 0000000000..02938c1067 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/prefetch.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/prefetch.py new file mode 100644 index 0000000000..2db36e42bb --- /dev/null +++ b/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()) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py new file mode 100644 index 0000000000..66b1d6e0f0 --- /dev/null +++ b/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.``. 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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/twoshot.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/twoshot.py new file mode 100644 index 0000000000..adb036460a --- /dev/null +++ b/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)) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/projection.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/projection.py new file mode 100644 index 0000000000..f74500a128 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/splitk.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/splitk.py new file mode 100644 index 0000000000..c0591ad578 --- /dev/null +++ b/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]) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/ws_gemm.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/ws_gemm.py new file mode 100644 index 0000000000..260ad2b22e --- /dev/null +++ b/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") diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/kernel.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/kernel.py new file mode 100644 index 0000000000..7caab11e1a --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py new file mode 100644 index 0000000000..4c43cf6834 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/dcp.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/dcp.py new file mode 100644 index 0000000000..4cd94b4261 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/prepared.py new file mode 100644 index 0000000000..758f21d28f --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/query.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/query.py new file mode 100644 index 0000000000..45fff905b6 --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/value.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/value.py new file mode 100644 index 0000000000..e074498400 --- /dev/null +++ b/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),)) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py new file mode 100644 index 0000000000..b95720ca1d --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts.py new file mode 100644 index 0000000000..80aa8bcb2f --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts.py @@ -0,0 +1,1579 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Persistent routed experts: TMA producer, MXFP8 activations, FP4 x MXFP8 FC1/FC2 MMAs, +SiTU hidden exchange and output epilogue. Capacity <= 8 keeps the batch's activations in +SMEM; capacity > 8 (grouped) runs (expert, <= 8 tokens) units whose FC1 stages gather +8-row slices of the grid-quantized batch.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int8, Int16, Int32, Int64, Uint8, Uint32, Uint64 + +from ..common.primitives import ( + acquire_tensormap, + alloc, + asm, + bf16, + bf16x2, + bits_float, + commit, + copy_evict, + copy_g2s, + fence_tmem, + float_bits, + fma, + fp8_bits, + group_barrier, + launch, + load_tma3, + load_tmem, + load_volatile, + load_words, + multicast_words, + mx_exponent, + mx_rescale, + pack_fp8, + pack_pair, + release_tmem, + shared_address, + situ, + smem_desc, + smem_words, + store_shared_words, + store_words, + thread, + unpack_bf16, + view, + wait, + wait_tmem, +) +from ..common.twoshot import push + +shuffle = cute.arch.shuffle_sync + + +class ExpertsGeometry: + """Static plain-input geometry. No cluster-local or multicast storage. The workspace + is capacity-independent; capacity > 8 swaps the batch's SMEM activation image for + per-slot 8-row images, which frees a sixth weight stage, and keeps per-CTA windows + of the unit tables (expert, 8 pair rows) in the misc block.""" + + def __init__(self, tp, capacity=8): + self.tp = tp + self.tiles, self.w13_rows, self.h_chunks, self.fc2_stages = ( + (6, 768, 3, 2) if tp == 8 else (3, 512, 2, 1) + ) + self.grouped = capacity > 8 + self.stages = stages = 6 if self.grouped else 5 + self.units = 16 * capacity if self.grouped else 128 + self.w13_scale_bytes = self.w13_rows * 112 + self.w2_scale_columns = 4 * self.h_chunks + self.w2_scale_bytes = 3584 * self.w2_scale_columns + self.w2_tile_scale_bytes = 128 * self.w2_scale_columns + self.h_image_bytes = self.h_chunks * 1024 + self.h_scale_bytes = self.h_chunks * 128 + 384 + self.h_compact_bytes = self.h_image_bytes + 128 + self.off_sfa = stages * 32768 + self.off_x = self.off_sfa + stages * 1024 + self.off_h = self.off_x + (stages * 2048 if self.grouped else 28672) + self.off_xsf = self.off_h + 2 * self.h_image_bytes + self.off_hsf = self.off_xsf + (stages * 256 if self.grouped else 28 * 128) + 384 + self.off_misc = self.off_hsf + 2 * self.h_scale_bytes + self.smem_bytes = self.off_misc + 5120 + # Misc barriers (byte offsets): weight full/empty [stages], accumulator + # full/empty [16], hidden full/empty [2], x ready [14]. Grouped stage metadata + # and unit windows (4 x 16 units) use the B<=8 routing tables' space. + self.bar_full, self.bar_empty, self.bar_accfull = (0, 8 * stages, 16 * stages) + acc = self.bar_accfull + self.bar_accempty, self.bar_hfull = (acc + 128, acc + 256) + self.bar_hempty, self.bar_xready = (acc + 272, acc + 288) + self.off_meta = self.off_misc + (1152 if self.grouped else 520) + self.off_window = self.off_misc + 1280 + # Stealing: FC2 unit queue (producer -> hidden loader) [8], full/empty [8, 8]. + self.off_queue, self.bar_queue = (self.off_misc + 4608, 4640) + self.ws_partial = 2 * 128 * self.h_compact_bytes + self.ws_counter = self.ws_partial + 2 * 160 * 8 * 128 * 4 + self.ws_xflag = self.ws_counter + 160 * 4 + self.ws_stage = (self.ws_xflag + 160 * 4 + 255) // 256 * 256 + # Publish stage rows [64 tokens x 16 slots, 3584] BF16, then grouped: call + # counters [160] and unit counts [2], hidden images [2, 1024]. + self.ws_calls = self.ws_stage + 1024 * 3584 * 2 + self.ws_units = self.ws_calls + 1024 + # Pregrouped unit table (``ExpertsGroupKernel``): header (units, call, grouped + # call), experts [1024], rows [1024, 8]. + self.ws_group = self.ws_units + 2 * 1024 * self.h_compact_bytes + # FC2(B) tail claim counters [2 parities, 160] (``ExpertsKernel.steal``). + self.ws_steal = self.ws_group + 16 + 20 * 1024 + self.workspace_bytes = self.ws_steal + 2 * 160 * 4 + + def image(self, workspace, parity, unit): + """Compact hidden image of ``unit`` in the ``parity`` half of the workspace.""" + if self.grouped: + base = workspace.iterator + self.ws_units + return base + (Int64(parity) * 1024 + unit) * self.h_compact_bytes + return workspace.iterator + (Int64(parity) * 128 + unit) * self.h_compact_bytes + + +@cute.jit +def swizzled(chunk, row, base=0): + """Byte offset of 16-byte ``chunk`` of ``row`` in a 128B-swizzled 8-row image.""" + return base + chunk // 8 * 1024 + row * 128 + (chunk % 8 ^ row) * 16 + + +@cute.jit +def has_ff_byte(word: Uint32): + value = ~word + return value - Uint32(16843009) & ~value & Uint32(2155905152) != 0 + + +@cute.jit +def mx_block(words): + """MXFP8 of one 32-BF16 block (16 words): eight e4m3x4 words and the E8M0 scale.""" + values = cute.make_rmem_tensor(32, Float32) + for i in cutlass.range_constexpr(16): + pair = unpack_bf16(words[i]) + values[2 * i], values[2 * i + 1] = (pair[0], pair[1]) + maxima = cute.make_rmem_tensor(16, Float32) + for j in cutlass.range_constexpr(16): + maxima[j] = cute.arch.fmax(abs(values[j]), abs(values[j + 16])) + for width in (8, 4, 2, 1): + for j in cutlass.range_constexpr(width): + maxima[j] = cute.arch.fmax(maxima[j], maxima[j + width]) + scale = mx_exponent(maxima[0]) + reciprocal = mx_rescale(scale) + packed = [ + pack_fp8(tuple(values[4 * j + i] * reciprocal for i in range(4))) + for j in range(8) + ] + return packed, scale + + +# Grouped expert plan, written before the experts run (by the routing kernel's copy +# role, or ExpertsPlanKernel): MXFP8 rows [64, 3584] then scale words [64, 28, 16 B] +# (bytes 0-3 used), the layout of the experts' per-slot SMEM images. +PLAN_XSF = 64 * 3584 +PLAN_BYTES = PLAN_XSF + 64 * 448 + + +@cute.jit +def plan_chunk(plan, token, chunk, words): + """Write 8-element ``chunk`` (four BF16x2 ``words``) of ``token`` as MXFP8 into + the plan; lanes 4k..4k+3 hold one 32-element block, 16 lanes one scale word. + Bit-identical to ``mx_block`` (an exact max, the same scaling and packing).""" + lane, values = (thread()[2], cute.make_rmem_tensor(8, Float32)) + for i in cutlass.range_constexpr(4): + pair = unpack_bf16(words[i]) + values[2 * i], values[2 * i + 1] = (pair[0], pair[1]) + top = abs(values[0]) + for i in cutlass.range_constexpr(1, 8): + top = cute.arch.fmax(top, abs(values[i])) + for width in (1, 2): + top = cute.arch.fmax(top, cute.arch.shuffle_sync_bfly(top, width)) + scale = mx_exponent(top) + reciprocal = mx_rescale(scale) + packed = [ + pack_fp8(tuple(values[4 * j + i] * reciprocal for i in range(4))) + for j in range(2) + ] + store_words(plan.iterator + token * 3584 + chunk * 8, packed) + word = Uint32(scale) + for s in cutlass.range_constexpr(1, 4): + word |= shuffle(Uint32(scale), lane + 4 * s) << 8 * s + if chunk % 16 == 0: + dst = plan.iterator + PLAN_XSF + token * 448 + chunk // 16 * 16 + store_words(dst, (word, Uint32(0), Uint32(0), Uint32(0))) + + +class ExpertsPlanKernel: + """Standalone plan writer (no routing kernel in front): one CTA per token.""" + + def __init__(self, batch): + self.batch = batch + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.batch, 256) + + @cute.kernel + def kernel(self, args): + x, plan = args + tid, _, _, token = thread() + cute.arch.griddepcontrol_wait() + for i in cutlass.range_constexpr(2): + chunk = tid + 256 * i + if chunk < 448: # warp-uniform: chunks 256-447 are warps 0-5 + words = load_words( + x.iterator + token * 3584 + chunk * 8, count=4, mutable=True + ) + plan_chunk(plan, token, chunk, words) + cute.arch.griddepcontrol_launch_dependents() + + +@asm() +def stage(mma, weight_scales, activation_scales, release, done, *, kind): + """One FC1/FC2 stage. ``mma``: (TMEM accumulator, weight desc, activation desc); + scales: (TMEM column, desc[, activation copy flag]); ``release``: (FC1 accumulate + or FC2 prior-H release, weight empty barrier); ``done``: accumulator barrier.""" + half = kind in ("fc2b", "fc2p") + steps = 6 if kind == "fc2" else 4 if half else 8 + cp = "tcgen05.cp.cta_group::1.32x128b.warpx4" + arrive = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64" + code = [ + "{ .reg .pred e,p0,pt,pc,pa,ph; .reg .b32 i<4>; .reg .b64 as1,bs1,bs2;", + f".reg .b32 af1,bf1,bf2; .reg .b64 xa<{steps}>,xb<{steps}>;", + "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 pc,$7,0;", + "and.pred pc,pc,e; setp.ne.b32 pa,$10,0; and.pred pa,pa,e;", + ] + if kind == "fc1": + code.append("setp.ne.b32 p0,$8,0;") + elif half: + code.append("setp.eq.b32 p0,0,0;") + else: + code.append("setp.eq.b32 p0,0,1; setp.ne.b32 ph,$8,0; and.pred ph,ph,e;") + code.append(f"@ph {arrive} [$8];") + code += [f"mov.b32 i{k},{142738048 | k << 29 | k << 4};" for k in range(4)] + if not half: + code.append("add.s64 as1,$4,32; add.s64 bs1,$6,8;") + code.append("add.u32 af1,$3,4; add.u32 bf1,$5,4;") + if kind == "fc2a": + code.append("add.s64 bs2,$6,16; add.u32 bf2,$5,8;") + code += [f"@pc {cp} [$5],$6;", f"@pc {cp} [bf1],bs1;"] + if kind == "fc2a": + code.append(f"@pc {cp} [bf2],bs2;") + code.append(f"@e {cp} [$3],$4;") + if not half: + code.append(f"@e {cp} [af1],as1;") + for k in range(steps): + a = k // 4 * 1024 + k % 4 * 2 + (4 if kind == "fc2" and k >= 4 else 0) + b = k // 4 * 64 + k % 4 * 2 + code.append(f"add.s64 xa{k},$1,{a}; add.s64 xb{k},$2,{b};") + for k in range(steps): + af, bf = ("$3", "$5") if k < 4 else ("af1", "bf1") + code.append( + "@e tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale " + f"[$0],xa{k},xb{k},i{k % 4},[{af}],[{bf}],{'p0' if k == 0 else 'pt'};" + ) + if kind != "fc2p": + code.append(f"@e {arrive} [$9];") + if kind != "fc2a": + code.append(f"@pa {arrive} [$10];") + code.append("}") + kinds = "rllrlrlrrrr" + operands = (*mma, *weight_scales, *activation_scales, *release, done) + values = [(Uint64 if k == "l" else Uint32)(v) for k, v in zip(kinds, operands)] + return "\n".join(code), ",".join(kinds) + ",~{memory}", values + + +@asm(Uint32) +def claim(pointer, value): + """Relaxed GPU-scope fetch-add (FC2 tail claims): the counter before the add.""" + code = "atom.relaxed.gpu.global.add.u32 $0,[$1],$2;" + return code, "=r,l,r,~{memory}", (pointer, Uint32(value)) + + +def fc2_pairs(tp, capacity, streamed): + """TP8 FC2 jobs are tile pairs on the incremental producer from capacity 4 (the + grouped engine always streams): chunks 0-1 of tiles t and t + 1, then both tiles' + chunk 2 through one 256-row box (map 5), so every stage is 16 KB.""" + return tp == 8 and streamed and capacity >= 4 + + +class ExpertsKernel(ExpertsGeometry): + """The engine for one capacity and persistent grid. CTAs wait on other clusters' + partials (``load_hidden``, ``partials``, ``finalize``), so + ``PersistentExperts.prepare`` checks the grid is co-resident. It is not launched + cooperatively: that forfeits the PDL prologue overlap (1-2% decode at B4-B8). + ``publish`` stages FC2 rows for ``finalize``: routing weights replace the output + arg, and (latent multicast, parity, rank) follow the SiTU scalars; ``"twoshot"``: + (receive slots [3 x TP, 64, 3584/TP], their owners' bases, parity x TP + rank) of + ``common.twoshot`` instead. Capacity > 8 (grouped) takes the plan (``PLAN_BYTES``: + MXFP8 x) in place of x. ``stream`` runs capacity <= 8 on the incremental producer + and CTA-scoped waits (grouped always). ``pregrouped`` (grouped) reads the unit table + ``ExpertsGroupKernel`` left in the workspace instead of grouping the routes itself. + ``steal`` (grouped) balances the end: each CTA's last ``zone`` FC2(B) grains (TP8 a + tile pair, TP16 three tiles) are claimed by the owner from the front and by the + ``victims`` CTAs before it from the back, through one counter per CTA re-armed by + call parity; a stolen grain is computed exactly as by its owner.""" + + def __init__( + self, tp, capacity, experts, grid, publish=False, + stream=True, pregrouped=False, steal=False, + ): # fmt: skip + super().__init__(tp, capacity) + self.pregrouped = pregrouped and self.grouped + self.steal = steal and self.grouped + self.streamed = stream or self.grouped + self.paired = fc2_pairs(tp, capacity, self.streamed) + self.jobs = 14 if self.paired else 28 # FC2 jobs per unit + self.grain, self.zone, self.victims = (1 if self.paired else 3, 8, 3) + self.threads = 384 + self.capacity, self.experts, self.grid = (capacity, experts, grid) + self.publish = publish + assert not self.grouped or grid >= capacity <= 64 + self.width = min(capacity, 8) + self.dense = capacity <= 4 + self.rounds = (112 * capacity + 319) // 320 if self.dense else 3 + + def sync(self, barrier, phase=0, tmem=False): + """Parity wait (then tcgen05 ordering if ``tmem``). Every barrier here is + CTA-local, so grouped and streamed waits use CTA scope.""" + scope = "cta" if self.streamed else "cluster" + (wait_tmem if tmem else wait)(barrier, phase, scope=scope) + + @cute.jit + def route(self, sm, ids, counters, batch: Int32): + """Dedupe after the dependency wait. Capacity <= 2: one work item per routing + pair; else first occurrences across tokens in order, keeping token slots.""" + tid, warp, lane, cta = thread() + misc = self.off_misc + table = view(sm.iterator, 0, Int32, self.experts) + expert, uof = (view(sm.iterator, misc + o, Int32, 128) for o in (1140, 1652)) + slots = view(sm.iterator, misc + 2176, Int8, 1024) + slot_words = view(sm.iterator, misc + 2176, Uint32, 256) + counts = view(sm.iterator, misc + 868, Int32, 4) + distinct, parity = (view(sm.iterator, misc + o, Int32, 1) for o in (860, 864)) + for i in range(tid, self.experts, 384): + table[i] = Int32(2147483647) + myid, pairs = (Int32(-1), batch * 16) + if tid < 128: + lo, hi = (Uint32(4294967295), Uint32(4294967295)) + if tid < pairs: + myid = ids[tid] + if cutlass.const_expr(self.capacity <= 2): + expert[tid] = myid + if tid < 16: + lo = lo & Uint32(4294967040) | Uint32(tid % 16) + else: + lo = lo & Uint32(4294902015) | Uint32(tid % 16) << 8 + slot_words[2 * tid], slot_words[2 * tid + 1] = (lo, hi) + if tid == 160: + parity[0] = counters[cta] + if cutlass.const_expr(self.capacity <= 2) and tid == 0: + distinct[0] = pairs + cute.arch.sync_threads() + if cutlass.const_expr(self.capacity > 2): + if tid < pairs: + cute.arch.atomic_min( + table.iterator + myid, tid, sem="relaxed", scope="cta" + ) + cute.arch.sync_threads() + if tid < 128: + firstk = Int32(-1) + if tid < pairs: + firstk = table[myid] + first = (tid < pairs) & (firstk == tid) + ballot = Uint32(cute.arch.vote_ballot_sync(first)) + if lane == 0: + counts[warp] = Int32(cute.arch.popc(ballot)) + cute.arch.barrier(barrier_id=1, number_of_threads=128) + base = Int32(0) + for w in cutlass.range_constexpr(3): + if warp > w: + base += counts[w] + rank = base + Int32(cute.arch.popc(ballot & (Uint32(1) << lane) - 1)) + if first: + uof[tid] = rank + expert[rank] = myid + if tid == 0: + distinct[0] = counts[0] + counts[1] + counts[2] + counts[3] + cute.arch.barrier(barrier_id=1, number_of_threads=128) + if tid < pairs: + slots[uof[firstk] * 8 + tid // 16] = Int8(tid % 16) + # Retire every thread's hash accesses before async weight-ring reuse. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + return (distinct[0], parity[0], expert) + + def tables(self, sm, window=True): + """Grouped unit tables: (expert, eight pair rows (Int16, -1 padded), the rows + as four words per unit) of every unit in the weight-ring scratch during + routing, or of this CTA's windows. Views all start from the byte pointer.""" + base, units = (self.off_window, 64) if window else (12288, self.units) + experts = view(sm.iterator, base, Int32, units) + rows = view(sm.iterator, base + 4 * units, Int16, 8 * units) + return experts, rows, view(sm.iterator, base + 4 * units, Uint32, 4 * units) + + @cute.jit + def first_unit(self, ranges, fc2: cutlass.Constexpr, group): + """First unit of this CTA's FC1 or FC2 range in ``group``.""" + first = ranges[group, 4] // self.jobs + if cutlass.const_expr(not fc2): + first = ranges[group, 0] + ranges[group, 2] // 14 // self.tiles + return first + + @cute.jit + def local(self, ranges, unit, fc2: cutlass.Constexpr, group): + """Window index of ``unit`` (windows: FC1 A, FC1 B, FC2 A, FC2 B).""" + return (2 * fc2 + group) * 16 + unit - self.first_unit(ranges, fc2, group) + + @cute.jit + def windows(self, sm, ranges): + """Copy this CTA's four unit windows (<= 16 units each) out of the scratch.""" + tid, _, _, cta = thread() + experts, _, words = self.tables(sm, False) + window_experts, _, window_words = self.tables(sm) + for w in cutlass.range_constexpr(4): + fc2, group = (w // 2, w % 2) + lo, hi = (ranges[group, 2 + 2 * fc2], ranges[group, 3 + 2 * fc2]) + if cutlass.const_expr(self.steal and w == 3): + hi = self.stolen_end(ranges, cta) + last = (hi - 1) // self.jobs + if cutlass.const_expr(not fc2): + last = ranges[group, 0] + (hi - 1) // 14 // self.tiles + unit = self.first_unit(ranges, fc2, group) + tid % 16 + if (tid // 16 == w) & (lo < hi) & (unit <= last): + window_experts[tid] = experts[unit] + for i in cutlass.range_constexpr(4): + window_words[4 * tid + i] = words[4 * unit + i] + + @cute.jit + def stolen_end(self, ranges, cta): + """End of the FC2(B) jobs of CTAs up to ``cta + victims`` (window 3's reach).""" + stop = cta + self.victims + 1 + if stop > self.grid: + stop = Int32(self.grid) + count = ranges[1, 1] // (self.tiles * 14) + return self.jobs * ranges[1, 0] + stop * self.jobs * count // self.grid + + @cute.jit + def group(self, sm, ids, workspace, batch: Int32, scalars=None): + """Grouped routing: per-expert token masks; expert e's + tokens, in token order, fill ceil(n / 8) units numbered in expert order. Unit + rows hold pair indices (token * 16 + slot), -1 when padded. Misc scalars: + units, call, grouped call (``scalars`` given: the units only).""" + tid, warp, lane, cta = thread() + masks = view(sm.iterator, 0, Uint32, 2 * self.experts) + base = view(sm.iterator, 8 * self.experts, Int32, self.experts) + totals = view(sm.iterator, 12 * self.experts, Int32, self.threads // 32) + experts, rows, row_words = self.tables(sm, False) + counted = scalars is None + if cutlass.const_expr(counted): + scalars = view(sm.iterator, self.off_misc + 1140, Int32, 3) + threads: cutlass.Constexpr = self.threads + rounds: cutlass.Constexpr = (16 * self.capacity + threads - 1) // threads + mine = cute.make_rmem_tensor(rounds, Int32) + for k in cutlass.range_constexpr(rounds): + mine[k] = Int32(0) + if tid + self.threads * k < batch * 16: + mine[k] = ids[tid + self.threads * k] + if cutlass.const_expr(counted): + if tid == 0: + scalars[1] = view(workspace.iterator, self.ws_counter, Int32, 160)[cta] + if tid == 1: + scalars[2] = view(workspace.iterator, self.ws_calls, Int32, 160)[cta] + for i in range(tid, 2 * self.experts, self.threads): + masks[i] = Uint32(0) + for i in range(tid, 4 * self.units, self.threads): + row_words[i] = Uint32(4294967295) + cute.arch.sync_threads() + for k in cutlass.range_constexpr(rounds): + pair = tid + self.threads * k + if pair < batch * 16: + token = pair >> 4 + bit = Uint32(1) << (token & 31) + word = masks.iterator + 2 * mine[k] + (token >> 5) + cute.arch.atomic_or(word, bit, sem="relaxed", scope="cta") + cute.arch.sync_threads() + span: cutlass.Constexpr = (self.experts + self.threads - 1) // self.threads + units = cute.make_rmem_tensor(span, Int32) + count = Int32(0) + for j in cutlass.range_constexpr(span): + units[j], e = (Int32(0), span * tid + j) + if e < self.experts: + n = cute.arch.popc(masks[2 * e]) + cute.arch.popc(masks[2 * e + 1]) + units[j] = Int32(n + 7) >> 3 + count += units[j] + total = count + for shift in cutlass.range_constexpr(5): + other = cute.arch.shuffle_sync_up(total, 1 << shift) + if lane >= 1 << shift: + total += other + if lane == 31: + totals[warp] = total + cute.arch.sync_threads() + start = total - count + for w in cutlass.range_constexpr(self.threads // 32 - 1): + if warp > w: + start += totals[w] + if tid == self.threads - 1: + scalars[0] = start + count + for j in cutlass.range_constexpr(span): + e = span * tid + j + if e < self.experts: + base[e] = start + for v in range(units[j]): + experts[start + v] = e + start += units[j] + cute.arch.sync_threads() + for k in cutlass.range_constexpr(rounds): + pair = tid + self.threads * k + if pair < batch * 16: + e, token = (mine[k], pair >> 4) + low, word = (masks[2 * e], masks[2 * e + 1]) + rank = Int32(0) + if token < 32: + word = low + else: + rank = Int32(cute.arch.popc(low)) + below = (Uint32(1) << (token & 31)) - 1 + rank += Int32(cute.arch.popc(word & below)) + rows[(base[e] + (rank >> 3)) * 8 + (rank & 7)] = Int16(pair) + cute.arch.sync_threads() + + @cute.jit + def fetch_windows(self, sm, workspace): + """Pregrouped: the misc scalars from the unit table's header (every CTA's call + counters are equal between calls), then thread w * 16 + i copies unit i of + window w (FC1 A, FC1 B, FC2 A, FC2 B as ``partition``) out of the table. One + compact block (this code runs cold at kernel start).""" + tid, _, _, cta = thread() + scalars = view(sm.iterator, self.off_misc + 1140, Int32, 3) + table = workspace.iterator + self.ws_group + if tid == 0: + header = load_words(table, count=4, mutable=True) + for i in cutlass.range_constexpr(3): + scalars[i] = Int32(header[i]) + cute.arch.sync_threads() + if tid < 64: + units, fc2, group = (scalars[0], tid >= 32, tid // 16 % 2 == 1) + middle = units + if units > 40: + middle = (units * 50 + 99) // 100 + start, count = (Int32(0), middle) + if group: + start, count = (middle, units - middle) + per = self.tiles * 14 + lo = cta * per * count // self.grid + hi = (cta + 1) * per * count // self.grid + first, last = (start + lo // per, start + (hi - 1) // per) + if fc2: + lo = cta * self.jobs * count // self.grid + hi = (cta + 1) * self.jobs * count // self.grid + if cutlass.const_expr(self.steal): + stop = cta + self.victims + 1 + if stop > self.grid: + stop = Int32(self.grid) + if group: + hi = stop * self.jobs * count // self.grid + first, last = (start + lo // self.jobs, start + (hi - 1) // self.jobs) + unit = first + tid % 16 + if (lo < hi) & (unit <= last): + expert = load_words(table + 16 + 4 * unit, count=1, mutable=True) + row = load_words(table + 16 + 4096 + 16 * unit, count=4, mutable=True) + window_experts, _, window_words = self.tables(sm) + window_experts[tid] = Int32(expert[0]) + for i in cutlass.range_constexpr(4): + window_words[4 * tid + i] = row[i] + + @cute.jit + def gather_x(self, sm, plan, ranges, batch): + """Warp 3 (grouped): cp.async each FC1 stage's 8-row MXFP8 slice and scale + rows out of the plan (complete before the dependency wait returns) into its + slot; full[slot] counts the 32 lanes besides the producer (which stands in for + them on FC2 stages and the end). Padded rows repeat the unit's first token + (their outputs are dropped).""" + lane, ws = (thread()[2], plan.iterator) + stages, rows = (self.stages, self.tables(sm)[1]) + full, empty = ( + view(sm.iterator, self.off_misc + o, Uint64, stages) + for o in (self.bar_full, self.bar_empty) + ) + chunk, lanes = (lane & 15, lane // 16) + image = [self.off_x + swizzled(chunk, 2 * j + lanes) for j in range(4)] + scale = self.off_xsf + lane // 8 * 128 + (lane & 7) * 16 + sources = cute.make_rmem_tensor(5, Int32) + slot, phase, issued = (Int32(0), Int32(0), Int32(0)) + for group in cutlass.range_constexpr(2): + lo, hi = (ranges[group, 2], ranges[group, 3]) + ks, r, index = (lo % 14, lo // 14 % self.tiles, Int32(16 * group)) + if lo < hi: + self.row_sources(rows, index, sources, batch) + for s in cutlass.range(lo, hi, unroll=1): + if issued >= stages: + self.sync(empty.iterator + slot, phase ^ 1) + for j in cutlass.range_constexpr(4): + src = sources[j] + ks * 256 + chunk * 16 + copy_g2s(sm.iterator + image[j] + slot * 2048, ws + src) + if lane < 16: + src = PLAN_XSF + sources[4] + (2 * ks + lane // 8) * 16 + copy_g2s(sm.iterator + scale + slot * 256, ws + src) + cute.arch.cp_async_mbarrier_arrive_noinc(full.iterator + slot) + slot, issued = (slot + 1, issued + 1) + if slot == stages: + slot, phase = (Int32(0), phase ^ 1) + ks += 1 + if ks == 14: + ks, r = (Int32(0), r + 1) + if r == self.tiles: + r, index = (Int32(0), index + 1) + if s + 1 < hi: + self.row_sources(rows, index, sources, batch) + + @cute.jit + def row_sources(self, rows, index, sources, batch): + """Byte offsets of this lane's four MXFP8 rows and scale row of window unit + ``index`` (padded rows repeat its first token).""" + lane = thread()[2] + for j in cutlass.range_constexpr(5): + row = 2 * j + lane // 16 if j < 4 else lane & 7 + pair = Int32(rows[index * 8 + row]) + if pair < 0: + pair = Int32(rows[index * 8]) + sources[j] = (pair >> 4) * (3584 if j < 4 else 448) + + @cute.jit + def partition(self, distinct: Int32, cta: Int32): + """FC1(A), FC1(B), FC2(A), FC2(B) work partitions.""" + middle = distinct + if distinct > 40: + middle = (distinct * 50 + 99) // 100 + ranges = cute.make_rmem_tensor((2, 6), Int32) + for group in cutlass.range_constexpr(2): + start, end = (Int32(0), middle) if group == 0 else (middle, distinct) + count = end - start + size = self.tiles * 14 * count + ranges[group, 0], ranges[group, 1] = (start, size) + ranges[group, 2] = cta * size // self.grid + ranges[group, 3] = (cta + 1) * size // self.grid + n = self.jobs + ranges[group, 4] = n * start + cta * n * count // self.grid + ranges[group, 5] = n * start + (cta + 1) * n * count // self.grid + return ranges + + @cute.jit + def schedule(self, ranges, experts, n: Int32): + """Returns the stage's (kind, a, b, c) metadata and its TMA arguments.""" + n1a = ranges[0, 3] - ranges[0, 2] + n1b = ranges[1, 3] - ranges[1, 2] + n2a = ranges[0, 5] - ranges[0, 4] + kind, a, b, c = (Int32(0), Int32(0), Int32(0), Int32(0)) + row, chunk, scale_bytes, scale_offset = (Int32(0), Int32(0), Int32(0), Int64(0)) + if n < n1a + n1b: + group = Int32(n >= n1a) + s = ranges[group, 2] + (n - n1a if group == 1 else n) + tile = self.tiles * ranges[group, 0] + s // 14 + ks = s % 14 + expert, r = (experts[tile // self.tiles], tile % self.tiles) + a, b = (tile, ks) + opens = Int32((s == ranges[group, 2]) | (ks == 0)) + c = opens | Int32((s == ranges[group, 3] - 1) | (ks == 13)) << 1 + row, chunk = (expert * self.w13_rows + r * 128, 2 * ks) + scale_offset = Int64(expert) * self.w13_scale_bytes + r * 14336 + ks * 1024 + scale_bytes = Int32(1024) + else: + step = n - n1a - n1b + unit, half = (step // self.fc2_stages, step % self.fc2_stages) + group = Int32(unit >= n2a) + job = ranges[group, 4] + (unit - n2a if group == 1 else unit) + u, tile = (job // 28, job % 28) + expert = experts[u] + kind = Int32(1) if half == 0 else Int32(3) + a, b = (u, tile) + row, chunk = (expert * 3584 + tile * 128, 2 * half) + scale_offset = Int64(expert) * self.w2_scale_bytes + scale_offset += tile * self.w2_tile_scale_bytes + scale_offset += half * 1024 + scale_bytes = Int32(1024 if half == 0 else self.w2_tile_scale_bytes - 1024) + return (kind, a, b, c), (row, chunk, scale_offset, scale_bytes) + + @cute.jit + def poll_hidden(self, sm, source, buffer: Int32): + """Poll real K data and scale bytes, preserving TP16 padding semantics.""" + lane = thread()[2] + chunks: cutlass.Constexpr = self.h_image_bytes // 512 + values, scale = (cute.make_rmem_tensor((chunks, 4), Uint32), Uint32(0)) + valid_bytes = min(4, max(0, self.tiles * 2 - 4 * (lane & 3))) + mask = Uint32(4294967295) + if valid_bytes < 4: + mask = (Uint32(1) << 8 * valid_bytes) - 1 + ready = False + while not ready: + bad = False + for v in cutlass.range_constexpr(chunks): + index = lane + 32 * v + words = load_volatile(source + index * 16, count=4) + for w in cutlass.range_constexpr(4): + values[v, w] = words[w] + atom, token, position = (index >> 6, index >> 3 & 7, index & 7) + if 8 * atom + (position ^ token) < self.tiles * 4: + for w in cutlass.range_constexpr(4): + bad |= has_ff_byte(words[w]) + scale = load_volatile(source + self.h_image_bytes + lane * 4)[0] + if valid_bytes > 0: + bad |= has_ff_byte(scale & mask) + ready = cute.arch.vote_all_sync(not bad) + image = sm.iterator + self.off_h + buffer * self.h_image_bytes + for v in cutlass.range_constexpr(chunks): + words = tuple(values[v, w] for w in range(4)) + store_shared_words(image + (lane + 32 * v) * 16, words) + word = cute.arch.shuffle_sync(scale, 4 * (lane & 7) + (lane >> 3)) + if lane < 8 * self.h_chunks: + sf = sm.iterator + self.off_hsf + buffer * self.h_scale_bytes + store_shared_words(sf + (lane >> 3) * 128 + (lane & 7) * 16, (word,) * 4) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + + @cute.jit + def load_hidden(self, sm, workspace, ranges, parity): + lane = thread()[2] + full = view(sm.iterator, self.off_misc + self.bar_hfull, Uint64, 2) + empty = view(sm.iterator, self.off_misc + self.bar_hempty, Uint64, 2) + ordinal = Int32(-1) + if cutlass.const_expr(self.steal): + unit = self.pop_unit(sm, Int32(0)) + while unit >= 0: + ordinal += 1 + buffer = ordinal & 1 + if ordinal >= 2: + self.sync(empty.iterator + buffer, (ordinal >> 1) - 1 & 1) + cute.arch.fence_proxy("async.shared", space="cta") # MMA read it + image = self.image(workspace, parity, unit) + self.poll_hidden(sm, image, buffer) + if lane == 0: + cute.arch.mbarrier_arrive(full.iterator + buffer) + unit = self.pop_unit(sm, ordinal + 1) + for g in cutlass.range_constexpr(2 * (not self.steal)): + if ranges[g, 5] > ranges[g, 4]: + lo, hi = (ranges[g, 4] // self.jobs, (ranges[g, 5] - 1) // self.jobs) + for u in range(lo, hi + 1): + ordinal += 1 + buffer = ordinal & 1 + if ordinal >= 2: + self.sync(empty.iterator + buffer, (ordinal >> 1) - 1 & 1) + cute.arch.fence_proxy("async.shared", space="cta") + image = self.image(workspace, parity, u) + self.poll_hidden(sm, image, buffer) + if lane == 0: + cute.arch.mbarrier_arrive(full.iterator + buffer) + + @cute.jit + def queue(self, sm): + """The unit queue: (units [8], full [8], empty [8] barriers).""" + bars = self.off_misc + self.bar_queue + units = view(sm.iterator, self.off_queue, Int32, 8) + full = view(sm.iterator, bars, Uint64, 8) + return units, full, view(sm.iterator, bars + 64, Uint64, 8) + + @cute.jit + def push_unit(self, sm, pushed, unit): + """Producer: queue ``unit`` (-1 ends) for the hidden loader.""" + units, full, empty = self.queue(sm) + slot = pushed % 8 + if pushed >= 8: + self.sync(empty.iterator + slot, (pushed // 8 - 1) & 1) + units[slot] = unit + cute.arch.mbarrier_arrive(full.iterator + slot) + return pushed + 1 + + @cute.jit + def pop_unit(self, sm, popped): + """Hidden loader warp: the ``popped``-th queued unit.""" + units, full, empty = self.queue(sm) + slot = popped % 8 + self.sync(full.iterator + slot, popped // 8 & 1) + unit = units[slot] + cute.arch.sync_warp() + if thread()[2] == 0: + cute.arch.mbarrier_arrive(empty.iterator + slot) + return unit + + @cute.jit + def fc2_job(self, ring, state, maps, experts, ranges, j, group: cutlass.Constexpr): + """FC2 job ``j`` (absolute) of window ``group``: queue a new unit, then issue + its stages (TP8 a tile pair, TP16 a tile) exactly as the static walk does.""" + slot, phase, issued, pushed, current = state + unit, job = (j // self.jobs, j % self.jobs) + if unit != current: + pushed = self.push_unit(ring[0], pushed, unit) + current = unit + index = 32 + 16 * group + unit - self.first_unit(ranges, 1, group) + expert = experts[index] + if cutlass.const_expr(self.paired): + work = (maps.iterator, expert, unit, 2 * job) + slot, phase, issued = self.fc2_pair(ring, slot, phase, issued, *work) + else: + row = expert * 3584 + job * 128 + offset = Int64(expert) * self.w2_scale_bytes + offset += job * self.w2_tile_scale_bytes + tma = (maps.iterator + 128, ring[3], row, Int32(0), offset, Int32(1024)) + meta = (Int32(1), unit, job, Int32(0)) + slot, phase, issued = self.issue(ring, slot, phase, issued, meta, tma, True) + return slot, phase, issued, pushed, current + + @cute.jit + def fc2_jobs(self, ring, state, maps, experts, ranges, lo, hi, group): + for j in cutlass.range(lo, hi, unroll=1): + state = self.fc2_job(ring, state, maps, experts, ranges, j, group) + return state + + @cute.jit + def claims(self, ring, state, maps, experts, ranges, counter, base, grains, back): + """Grains of one zone (``grains`` from job ``base``) claimed through + ``counter`` (baseline -1: low half front claims, high half back claims) from + the front or the ``back``; the next claim is in flight while one issues. + Returns the state and the grains taken.""" + step = (1 << 16) if back else 1 + pending = claim(counter, step) + active, taken = (grains > 0, Int32(0)) + while active: + value = pending + 1 + front, rear = (Int32(value & 65535), Int32(value >> 16)) + active = front + rear < grains + if active: + ahead = claim(counter, step) + index = front + if back: + index = grains - 1 - rear + lo = base + index * self.grain + hi = lo + self.grain + state = self.fc2_jobs(ring, state, maps, experts, ranges, lo, hi, 1) + pending, taken = (ahead, taken + 1) + return state, taken + + @cute.jit + def fc2_steal(self, ring, state, maps, experts, ranges, workspace, parity): + """FC2 with tail stealing: FC2(A) and the static head of FC2(B), the owner's + zone from the front, then the zones of CTAs cta + 1 .. cta + victims from the + back while their units fit window 3; the queue ends with -1.""" + cta = thread()[3] + n, g = (self.jobs, self.grain) + head = (ranges[0, 4], ranges[0, 5], 0) + state = self.fc2_jobs(ring, state, maps, experts, ranges, *head) + counters = workspace.iterator + self.ws_steal + 640 * parity + lo, hi = (ranges[1, 4], ranges[1, 5]) + grains = (hi - lo) // g + if grains > self.zone: + grains = Int32(self.zone) + base = hi - grains * g + state = self.fc2_jobs(ring, state, maps, experts, ranges, lo, base, 1) + own = (counters + 4 * cta, base, grains, False) + state, _ = self.claims(ring, state, maps, experts, ranges, *own) + start, count = (ranges[1, 0], ranges[1, 1] // (self.tiles * 14)) + first = ranges[1, 4] // n + victim = cta + 1 + while (victim <= cta + self.victims) & (victim < self.grid): + lo = n * start + victim * n * count // self.grid + hi = n * start + (victim + 1) * n * count // self.grid + grains = (hi - lo) // g + if grains > self.zone: + grains = Int32(self.zone) + if (hi - 1) // n - first > 15: + victim = cta + self.victims + 1 # beyond window 3: stop + else: + counter = counters + 4 * victim + work = (counter, hi - grains * g, grains, True) + state, taken = self.claims(ring, state, maps, experts, ranges, *work) + victim = victim + 1 + slot, phase, issued, pushed, current = state + pushed = self.push_unit(ring[0], pushed, Int32(-1)) + return slot, phase, issued + + @cute.jit + def fetch(self, sm, desc, scales, full, slot, row, chunk, offset, size, more=None): + """Weight box and scale bytes (``more``: a second ``size``-byte scale run).""" + load_tma3(sm.iterator + slot * 32768, desc, full.iterator + slot, row, chunk) + dst = sm.iterator + self.off_sfa + slot * 1024 + copy_evict(dst, scales.iterator + offset, size, full.iterator + slot) + if cutlass.const_expr(more is not None): + copy_evict(dst + 512, scales.iterator + more, size, full.iterator + slot) + + @cute.jit + def produce(self, sm, maps, w13_scales, w2_scales, ranges, experts): + """Producer lane; dynamic scale bytes: 1024, or 512 for the TP8 FC2 tail.""" + stages = self.stages + barriers = view(sm.iterator, self.off_misc, Uint64, 2 * stages) + metadata = view(sm.iterator, self.off_meta, Int32, 4 * stages) + fc1 = ranges[0, 3] - ranges[0, 2] + ranges[1, 3] - ranges[1, 2] + fc2 = ranges[0, 5] - ranges[0, 4] + ranges[1, 5] - ranges[1, 4] + total = fc1 + self.fc2_stages * fc2 + for n in cutlass.range(total, unroll=1): + slot, phase = (n % stages, n // stages % 2) + meta, tma = self.schedule(ranges, experts, n) + if n >= stages: + empty = barriers.iterator + stages + slot + self.sync(empty, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = meta[i] + cute.arch.mbarrier_arrive_and_expect_tx( + barriers.iterator + slot, 16384 + tma[3] + ) + if meta[0] == 0: + self.fetch(sm, maps.iterator, w13_scales, barriers, slot, *tma) + else: + self.fetch(sm, maps.iterator + 128, w2_scales, barriers, slot, *tma) + slot, phase = (total % stages, total // stages % 2) + if total >= stages: + self.sync(barriers.iterator + stages + slot, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = Int32(2 if i == 0 else 0) + cute.arch.mbarrier_arrive(barriers.iterator + slot) + + @cute.jit + def stream(self, sm, maps, w13_scales, w2_scales, ranges, experts, steal=None): + """Producer lane walking its FC1 then FC2 stages incrementally (no per-stage + divisions; the expert is re-read only when the unit changes). Grouped expert + tables are this CTA's windows; else they hold every distinct expert.""" + tiles, w13, w2 = (self.tiles, maps.iterator, maps.iterator + 128) + barriers = view(sm.iterator, self.off_misc, Uint64, 2 * self.stages) + metadata = view(sm.iterator, self.off_meta, Int32, 4 * self.stages) + ring = (sm, barriers, metadata, w2_scales) + slot, phase, issued = (Int32(0), Int32(0), Int32(0)) + for group in cutlass.range_constexpr(2): + lo, hi = (ranges[group, 2], ranges[group, 3]) + ks, tile = (lo % 14, tiles * ranges[group, 0] + lo // 14) + r, index = (tile % tiles, self.first_unit(ranges, 0, group)) + if cutlass.const_expr(self.grouped): + index = Int32(16 * group) + expert = experts[index] + for s in cutlass.range(lo, hi, unroll=1): + opens, closes = ((s == lo) | (ks == 0), (s == hi - 1) | (ks == 13)) + c = Int32(opens) | Int32(closes) << 1 + row = expert * self.w13_rows + r * 128 + offset = Int64(expert) * self.w13_scale_bytes + r * 14336 + ks * 1024 + tma = (w13, w13_scales, row, 2 * ks, offset, Int32(1024)) + slot, phase, issued = self.issue( + ring, slot, phase, issued, (Int32(0), tile, ks, c), tma, False + ) + ks += 1 + if ks == 14: + ks, tile, r = (Int32(0), tile + 1, r + 1) + if r == tiles: + r, index = (Int32(0), index + 1) + expert = experts[index] + if cutlass.const_expr(self.steal): + state = (slot, phase, issued, Int32(0), Int32(-1)) + work = (ring, state, maps, experts, ranges, *steal) + slot, phase, issued = self.fc2_steal(*work) + for group in cutlass.range_constexpr(2 * (not self.steal)): + lo, hi = (ranges[group, 4], ranges[group, 5]) + unit, tile, index = (lo // self.jobs, lo % self.jobs, lo // self.jobs) + if cutlass.const_expr(self.grouped): + index = Int32(32 + 16 * group) + expert = experts[index] + for _ in cutlass.range(lo, hi, unroll=1): + if cutlass.const_expr(self.paired): + job = (maps.iterator, expert, unit, 2 * tile) + slot, phase, issued = self.fc2_pair(ring, slot, phase, issued, *job) + else: + for half in cutlass.range_constexpr(self.fc2_stages): + size = 1024 if half == 0 else self.w2_tile_scale_bytes - 1024 + row = expert * 3584 + tile * 128 + offset = Int64(expert) * self.w2_scale_bytes + half * 1024 + offset += tile * self.w2_tile_scale_bytes + tma = (w2, w2_scales, row, 2 * half, offset, Int32(size)) + meta = (Int32(1 if half == 0 else 3), unit, tile, Int32(0)) + slot, phase, issued = self.issue( + ring, slot, phase, issued, meta, tma, True + ) + tile += 1 + if tile == self.jobs: + unit, tile, index = (unit + 1, Int32(0), index + 1) + expert = experts[index] + if issued >= self.stages: + self.sync(barriers.iterator + self.stages + slot, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = Int32(2 if i == 0 else 0) + arrivals = 33 if self.grouped else 1 + cute.arch.mbarrier_arrive(barriers.iterator + slot, arrive_count=arrivals) + + @cute.jit + def fc2_pair(self, ring, slot, phase, issued, maps, expert, unit, tile): + """The three 16 KB stages of FC2 tiles ``tile`` and ``tile + 1``: chunks 0-1 of + each (1 KB scales), then chunk 2 of both (map 5, two 512 B scale runs).""" + w2, scales, step = (maps + 128, ring[3], self.w2_tile_scale_bytes) + row = expert * 3584 + tile * 128 + offset = Int64(expert) * self.w2_scale_bytes + tile * step + for t in cutlass.range_constexpr(2): + tma = (w2, scales, row + 128 * t, Int32(0), offset + t * step, Int32(1024)) + meta = (Int32(1), unit, tile + t, Int32(0)) + slot, phase, issued = self.issue(ring, slot, phase, issued, meta, tma, True) + tma = (maps + 640, scales, row, Int32(2), offset + 1024, Int32(512)) + meta = (Int32(3), unit, tile, Int32(0)) + more = offset + step + 1024 + return self.issue(ring, slot, phase, issued, meta, tma, True, more) + + @cute.jit + def issue( + self, ring, slot, phase, issued, meta, tma, fc2: cutlass.Constexpr, more=None + ): + """One grouped stage: wait for its slot, publish the metadata, expect and issue + the weight box and scales (FC2: standing in for the gather warp's 32 + arrivals); returns the advanced (slot, phase, issued).""" + sm, barriers, metadata = ring[:3] + desc, scales, row, chunk, offset, size = tma + if issued >= self.stages: + empty = barriers.iterator + self.stages + slot + self.sync(empty, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = meta[i] + full = barriers.iterator + slot + total = 16384 + size if more is None else 16384 + 2 * size + cute.arch.mbarrier_arrive_and_expect_tx(full, total) + self.fetch(sm, desc, scales, barriers, slot, row, chunk, offset, size, more) + if cutlass.const_expr(fc2 and self.grouped): + cute.arch.mbarrier_arrive(full, arrive_count=32) + slot += 1 + if slot == self.stages: + slot, phase = (Int32(0), phase ^ 1) + return slot, phase, issued + 1 + + @cute.jit + def ownership(self, index: cutlass.Constexpr[int], batch): + xtid = thread()[0] - 64 + if cutlass.const_expr(self.dense): + block = xtid + 320 * index + q = block // 8 + stage = q // batch + token = q - stage * batch + kblock = block & 7 + valid = block < 112 * batch + else: + stage = 5 * index + xtid // 64 + token = (xtid & 63) // 8 + kblock = xtid & 7 + valid = (stage < 14) & (token < batch) + return (stage, token, kblock, valid) + + @cute.jit + def load_x(self, x, raw, batch): + """Warps 2..11 preload their 32-BF16 x blocks before the dependent launch.""" + raw.fill(0) + for index in cutlass.range_constexpr(self.rounds): + stage, token, kblock, valid = self.ownership(index, batch) + if valid: + source = x.iterator + token * 3584 + stage * 256 + kblock * 32 + for fragment in cutlass.range_constexpr(4): + words = load_words( + source + fragment * 8, count=4, cache_policy="", mutable=True + ) + for j in cutlass.range_constexpr(4): + raw[index, fragment, j] = words[j] + + @cute.jit + def quantize_x(self, sm, raw, barriers, batch): + """Each thread quantizes its block (two uint4 fragments, one scale byte) and + arrives once at xready[stage]: eight arrivals per active token, no DSMEM.""" + for index in cutlass.range_constexpr(self.rounds): + stage, token, kblock, valid = self.ownership(index, batch) + if valid: + packed, scale = mx_block([raw[index, i // 4, i % 4] for i in range(16)]) + for half, words in enumerate((packed[:4], packed[4:])): + offset = swizzled(16 * stage + 2 * kblock + half, token, self.off_x) + store_shared_words(sm.iterator + offset, words) + kb = 8 * stage + kblock + sm[self.off_xsf + kb // 4 * 128 + token * 16 + kb % 4] = Uint8(scale) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.mbarrier_arrive(barriers + stage) + + @cute.jit + def acquire(self, empty, acc: Int32, uses: Int32): + if uses >= 16: + self.sync(empty.iterator + acc, uses // 16 - 1 & 1, True) + + @cute.jit + def signal(self, metadata, full, acc, values, arrivals: cutlass.Constexpr = 1): + if thread()[2] == 0: + for i in cutlass.range_constexpr(4): + metadata[4 * acc + i] = values[i] + for _ in cutlass.range_constexpr(arrivals): + cute.arch.mbarrier_arrive(full.iterator + acc) + + @cute.jit + def consume(self, sm, tmem): + lane = thread()[2] + stages = self.stages + bars: tuple = (self.bar_full, self.bar_empty, self.bar_accfull) + bars += (self.bar_accempty, self.bar_hfull, self.bar_hempty, self.bar_xready) + full, empty, accfull, accempty, hfull, hempty, xready = ( + view(sm.iterator, self.off_misc + o, Uint64, n) + for o, n in zip(bars, (stages, stages, 16, 16, 2, 2, 14)) + ) + meta = view(sm.iterator, self.off_meta, Int32, 4 * stages) + accmeta = view(sm.iterator, self.off_misc + 600, Int32, 64) + offsets = (0, self.off_sfa, self.off_x, self.off_xsf, self.off_h, self.off_hsf) + a0, asf0, x0, xsf0, h0, hsf0 = ( + smem_desc(sm.iterator + o, i % 2 == 0) for i, o in enumerate(offsets) + ) + xmask, ended = (Uint32(0), False) + slot, uses, acc, count = (Int32(0), Int32(0), Int32(0), Int32(0)) + current, ordinal, hbuffer = (Int32(-1), Int32(-1), Int32(0)) + while not ended: + slot = count % stages + ready = (full.iterator + slot, count // stages & 1, True) + self.sync(*ready) + kind, a, b, c = (meta[4 * slot + i] for i in range(4)) + if kind == 2: + if (lane == 0) & (current >= 0) & (hbuffer < 2): + commit(hempty.iterator + hbuffer) + for _ in cutlass.range(2, unroll=1): + self.acquire(accempty, acc, uses) + self.signal(accmeta, accfull, acc, (kind, a, b, c), 2) + uses += 1 + acc = uses % 16 + ended = True + else: + publish = False + if kind == 0: + first, last = (c & 1 != 0, c & 2 != 0) + if first: + self.acquire(accempty, acc, uses) + copy_scales = Uint32(0) + if cutlass.const_expr(self.grouped): + # The slot's own image (its cp.async completed full[slot]). + cute.arch.fence_proxy("async.shared", space="cta") + copy_scales = Uint32(1) + elif xmask >> b & 1 == 0: + self.sync(xready.iterator + b, 0, True) + xmask |= Uint32(1) << b + copy_scales = Uint32(1) + completion = Uint32(0) + if last: + completion = shared_address(accfull.iterator + acc) + xi = slot if self.grouped else b + stage( + (tmem + acc * 8, a0 + slot * 2048, x0 + xi * 128), + (tmem + 128 + slot * 8, asf0 + slot * 64), + (tmem + 256 + xi * 8, xsf0 + xi * 16, copy_scales), + (Uint32(not first), shared_address(empty.iterator + slot)), + completion, + kind="fc1", + ) + cute.arch.sync_warp() + publish = last + else: + copy_scales, release = (Uint32(0), Uint32(0)) + if a != current: + if (current >= 0) & (hbuffer < 2): + release = shared_address(hempty.iterator + hbuffer) + current = a + ordinal += 1 + hbuffer = ordinal & 1 + ready = (hfull.iterator + hbuffer, ordinal >> 1 & 1, True) + self.sync(*ready) + copy_scales = Uint32(1) + hd = h0 + hbuffer * (self.h_image_bytes // 16) + hs = hsf0 + hbuffer * (self.h_scale_bytes // 16) + sfbc = tmem + 128 + 8 * stages + hbuffer * self.w2_scale_columns + tp16 = cutlass.const_expr(self.tp == 16) + if cutlass.const_expr(self.paired): + weights = (tmem + 128 + slot * 8, asf0 + slot * 64) + released = shared_address(empty.iterator + slot) + if kind == 1: # chunks 0-1 of the pair's tile b + target, owner = (acc, uses) + if b % 2 == 1: + target, owner = ((uses + 1) % 16, uses + 1) + self.acquire(accempty, target, owner) + stage( + (tmem + target * 8, a0 + slot * 2048, hd), + weights, + (sfbc, hs, copy_scales), + (release, released), + 0, + kind="fc2a", + ) + else: # chunk 2 of tiles b (rows 0-127) and b + 1 (128-255) + for t in cutlass.range_constexpr(2): + rows = a0 + slot * 2048 + 1024 * t + stage( + (tmem + acc * 8, rows, hd + 128), + (weights[0] + 4 * t, weights[1] + 32 * t), + (sfbc + 8, 0, 0), + (0, released), + shared_address(accfull.iterator + acc), + kind="fc2b" if t else "fc2p", + ) + if cutlass.const_expr(t == 0): + cute.arch.sync_warp() + first = (Int32(1), a, b, c) + self.signal(accmeta, accfull, acc, first) + uses += 1 + acc = uses % 16 + cute.arch.sync_warp() + publish, kind, b = (True, Int32(1), b + 1) + else: + if tp16 or kind == 1: + self.acquire(accempty, acc, uses) + stage( + (tmem + acc * 8, a0 + slot * 2048, hd), + (tmem + 128 + slot * 8, asf0 + slot * 64), + (sfbc, hs, copy_scales), + (release, shared_address(empty.iterator + slot)), + shared_address(accfull.iterator + acc) if tp16 else 0, + kind="fc2" if tp16 else "fc2a", + ) + publish = tp16 + else: + stage( + (tmem + acc * 8, a0 + slot * 2048, hd + 128), + (tmem + 128 + slot * 8, asf0 + slot * 64), + (sfbc + 8, 0, 0), + (0, shared_address(empty.iterator + slot)), + shared_address(accfull.iterator + acc), + kind="fc2b", + ) + publish = True + kind = Int32(1) + cute.arch.sync_warp() + if publish: + self.signal(accmeta, accfull, acc, (kind, a, b, c)) + uses += 1 + acc = uses % 16 + count += 1 + + @cute.jit + def partials(self, workspace, values, tile, ranges, batch): + """Publish non-owner fragments, or reduce the owner's remaining pieces.""" + ws, live = (workspace.iterator, 8 if self.grouped else batch) + _, warp, lane, cta = thread() + row = 32 * (warp & 3) + lane + group = Int32(tile >= self.tiles * ranges[1, 0]) + start = (tile - self.tiles * ranges[group, 0]) * 14 + owner = start >= ranges[group, 2] + if not owner: + offset = self.ws_partial + ((Int64(group) * 160 + cta) * 1024 + row) * 4 + for t in cutlass.range_constexpr(self.width): + if t < live: + store_words(ws + offset + t * 512, (float_bits(values[t]),)) + else: + peer, size = (cta + 1, ranges[group, 1]) + while (peer < self.grid) & (peer * size // self.grid < start + 14): + if (peer + 1) * size // self.grid != peer * size // self.grid: + source = ws + self.ws_partial + source += ((Int64(group) * 160 + peer) * 1024 + row) * 4 + words = cute.make_rmem_tensor(self.width, Uint32) + ready = False + while not ready: + bad = False + for t in cutlass.range_constexpr(self.width): + words[t] = Uint32(0) + if t < live: + words[t] = load_volatile(source + t * 512)[0] + bad |= words[t] == Uint32(4294967295) + ready = cute.arch.vote_all_sync(not bad) + for t in cutlass.range_constexpr(self.width): + if t < live: + values[t] += bits_float(words[t]) + store_words(source + t * 512, (Uint32(4294967295),)) + peer += 1 + return owner + + @cute.jit + def store_hidden(self, sm, workspace, values, tile, parity, scalars): + _, warp, lane, _ = thread() + ew, eg = (warp & 3, (warp - 4) // 4) + expert, chunk = (tile // self.tiles, tile % self.tiles) + up = lane & 8 == 0 + h, maxima = (cute.make_rmem_tensor(self.width, Float32) for _ in range(2)) + shared_max = view(sm.iterator, self.off_misc + 884, Float32, 64) + for t in cutlass.range_constexpr(self.width): + gate = cute.arch.shuffle_sync_bfly(values[t], 8) + h[t] = situ(gate, values[t], *scalars) if up else Float32(0) + maxima[t] = abs(h[t]) + for t in cutlass.range_constexpr(self.width): + for off in cutlass.range_constexpr(5): + other = cute.arch.shuffle_sync_bfly(maxima[t], 1 << off) + maxima[t] = cute.arch.fmax(maxima[t], other) + if lane == 0: + shared_max[(eg * 4 + ew) * 8 + t] = maxima[t] + group_barrier(eg) + stage_base, scale_token = (self.off_misc + 3200 + (eg * 4 + ew) * 128, Int32(0)) + for t in cutlass.range_constexpr(self.width): + peer = shared_max[(eg * 4 + (ew ^ 1)) * 8 + t] + scale = mx_exponent(cute.arch.fmax(maxima[t], peer), floor=self.tp == 8) + if lane == t: + scale_token = scale + if up: + bits = fp8_bits(h[t] * mx_rescale(scale)) + sm[stage_base + t * 16 + 2 * (lane % 8) + lane // 16] = bits + cute.arch.sync_warp() + if lane < 8: + kc, data = (4 * chunk + ew, (Uint32(0), Uint32(0), Uint32(0), Uint32(0))) + if lane < self.width: + data = smem_words(sm.iterator + stage_base + lane * 16) + image = self.image(workspace, parity, expert) + store_words(image + swizzled(kc, lane), data) + if ew & 1 == 0: + byte = view(image, 0, Uint8, self.h_compact_bytes) + value = Uint8(scale_token) + byte[self.h_image_bytes + lane * 16 + 2 * chunk + ew // 2] = value + group_barrier(eg) + + def pair(self, t, slot): + """Expanded row of unit row ``t``: its pair (grouped), or token t's ``slot``.""" + return Int64(slot) if self.grouped else Int64(t) * 16 + slot + + @cute.jit + def fc2(self, sm, workspace, output, values, unit, tile, ranges): + """Rows of ``unit``: slots of tokens 0..7, or grouped pair indices.""" + _, warp, lane, _ = thread() + if cutlass.const_expr(self.grouped): + slots = self.tables(sm)[1] + unit = self.local(ranges, unit, 1, Int32(unit >= ranges[1, 0])) + else: + slots = view(sm.iterator, self.off_misc + 2176, Int8, 1024) + for t in cutlass.range_constexpr(self.width): + slot = Int32(slots[unit * 8 + t]) + if slot >= 0: + other = cute.arch.shuffle_sync_down(values[t], 8) + if cutlass.const_expr(self.publish): + # Stage row [token, slot]: lane q < 4 stores this warp's 8q..8q+7. + word, q = pack_pair(values[t], other, policy="stage"), lane & 3 + words = [shuffle(word, s + 2 * q) for s in (0, 16, 1, 17)] + row = self.pair(t, slot) * 1792 + tile * 64 + (warp & 3) * 16 + stage = workspace.iterator + self.ws_stage + (row + 4 * q) * 4 + if lane < 4: + store_words(stage, words) + elif lane & 8 == 0: + column = tile * 128 + (warp & 3) * 32 + 4 * (lane % 8) + lane // 8 + dst = output.iterator + self.pair(t, slot) * 3584 + column + store_words(dst, (bf16x2(values[t], other),)) + + @cute.jit + def finalize(self, workspace, weights, batch, box): + """Each epilogue warp polls (and re-arms) the 16 staged rows of its 64-column + lines and multicasts their weighted sum into latent slot (parity, rank) of the + [3, TP, 64, 3584] mailbox (two-shot: pushes every 8 columns to their owner); + lines stride over every CTA's epilogue warps.""" + _, warp, lane, cta = thread() + for line in range(cta * 8 + warp - 4, batch * 56, self.grid * 8): + token, column = line // 56, line % 56 * 64 + source = workspace.iterator + self.ws_stage + source += (Int64(token) * 16 * 1792 + column // 2 + lane) * 4 + words = cute.make_rmem_tensor(16, Uint32) + ready = False + while not ready: + bad = False + for slot in cutlass.range_constexpr(16): + words[slot] = load_volatile(source + slot * 7168)[0] + bad |= words[slot] == Uint32(4294967295) + ready = cute.arch.vote_all_sync(not bad) + low, high = Float32(0), Float32(0) + for slot in cutlass.range_constexpr(16): + store_words(source + slot * 7168, (Uint32(4294967295),)) + weight = Float32(weights[token * 16 + slot]) + low = fma(bf16(words[slot]), weight, low) + high = fma(bf16(words[slot], True), weight, high) + word = pack_pair(low, high, policy="latent") + packed = [shuffle(word, 4 * (lane & 7) + i) for i in range(4)] + if lane < 8: + if cutlass.const_expr(self.publish == "twoshot"): + push(*box, token, column + 8 * lane, packed) + else: + mc, parity, rank = box + offset = ((parity * self.tp + rank) * 64 + token) * 3584 + column + offset += 8 * lane + multicast_words(mc + Int64(offset) * 2, packed) + + @cute.jit + def epilogue(self, sm, workspace, output, tmem, ranges, batch, parity, scalars): + _, warp, lane, _ = thread() + ew, eg = (warp & 3, (warp - 4) // 4) + full = view(sm.iterator, self.off_misc + self.bar_accfull, Uint64, 16) + empty = view(sm.iterator, self.off_misc + self.bar_accempty, Uint64, 16) + meta = view(sm.iterator, self.off_misc + 600, Int32, 64) + acc, phase, ended = (eg, Int32(0), False) + while not ended: + self.sync(full.iterator + acc, phase, True) + kind, a, b = (meta[4 * acc], meta[4 * acc + 1], meta[4 * acc + 2]) + if kind == 2: + ended = True + else: + raw = load_tmem(tmem + (Uint32(32 * ew) << 16) + acc * 8, 8) + values = cute.make_rmem_tensor(self.width, Float32) + for t in cutlass.range_constexpr(self.width): + values[t] = raw[t] + fence_tmem(before=True) + cute.arch.sync_warp() + if lane == 0: + cute.arch.mbarrier_arrive(empty.iterator + acc) + acc += 2 + if acc >= 16: + acc -= 16 + phase ^= 1 + if kind == 0: + if self.partials(workspace, values, a, ranges, batch): + self.store_hidden(sm, workspace, values, a, parity, scalars[:3]) + else: + self.fc2(sm, workspace, output, values, a, b, ranges) + if cutlass.const_expr(self.publish): + self.finalize(workspace, output, batch, scalars[3:]) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, 384, cluster=1) + + @cute.jit + def initialize(self, sm, maps, batch): + tid, _, _ = cute.arch.thread_idx() + barriers = view(sm.iterator, self.off_misc, Uint64, 65) + tmem_base = view(sm.iterator, self.off_misc + 856, Uint32, 1) + if tid == 0: + receive = range(528, 533) if self.capacity == 4 else () + for i in (*range(65), *receive): + acc = 2 * self.stages + count = 1 if not acc <= i < acc + 32 else 2 if i < acc + 16 else 4 + count = batch * 8 if acc + 36 <= i < acc + 50 else count + if self.grouped and i < self.stages: + count = 33 + cute.arch.mbarrier_init(barriers.iterator + i, count) + if cutlass.const_expr(self.steal): + queue = view(sm.iterator, self.off_misc + self.bar_queue, Uint64, 16) + for i in cutlass.range_constexpr(16): + cute.arch.mbarrier_init(queue.iterator + i, 1) + cute.arch.mbarrier_init_fence() + acquire_tensormap(maps.iterator) + acquire_tensormap(maps.iterator + 128) + if cutlass.const_expr(self.paired): + acquire_tensormap(maps.iterator + 640) + if tid // 32 == 3: + cute.arch.alloc_tmem(512, tmem_base.iterator) + cute.arch.relinquish_tmem_alloc_permit() + if cutlass.const_expr(not self.grouped): + self.clear_x(sm, tid, batch) + cute.arch.fence_proxy("async.shared", space="cta") + return tmem_base + + @cute.jit + def clear_x(self, sm, tid, batch): + """Zero the padded token rows and the scale image (capacity <= 8).""" + zero = (Uint32(0), Uint32(0), Uint32(0), Uint32(0)) + for i in range(tid, (8 - batch) * 224, 384): + token, chunk = (batch + i // 224, i % 224) + store_shared_words(sm.iterator + swizzled(chunk, token, self.off_x), zero) + xsf = sm.iterator + self.off_xsf + for i in range(tid, 224, 384): + store_shared_words(xsf + (i >> 3) * 128 + (i & 7) * 16, zero) + + @cute.jit + def rearm(self, workspace, call_count): + """Re-arm the other parity's images: all 128, or the last grouped call's.""" + _, _, lane, cta = thread() + items: cutlass.Constexpr = 128 * self.h_compact_bytes // 16 + lo, hi = (cta * items // self.grid, (cta + 1) * items // self.grid) + if cutlass.const_expr(self.grouped): + last = view(workspace.iterator, self.ws_calls + 640, Int32, 2) + used = max(last[call_count & 1 ^ 1], 0) * (self.h_compact_bytes // 16) + lo, hi = (cta * used // self.grid, (cta + 1) * used // self.grid) + other = self.image(workspace, call_count & 1 ^ 1, 0) + for i in range(lo + lane, hi, 32): + store_words(other + i * 16, (Uint32(4294967295),) * 4) + if cutlass.const_expr(self.steal): + # Every one of the 160 counters: grids change with the batch. + base = self.ws_steal + 640 * (call_count & 1 ^ 1) + counters = view(workspace.iterator, base, Int32, 160) + if lane == 0: + counters[cta] = Int32(-1) + if cta < 160 - self.grid: + counters[self.grid + cta] = Int32(-1) + + @cute.kernel + def kernel(self, args): + x, ids, maps, w13_scales, w2_scales, workspace, output, batch, *scalars = args + tid, warp, lane, cta = thread() + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem_bytes, 1024) + tmem_base = self.initialize(sm, maps, batch) + cute.arch.griddepcontrol_wait() + counters = view(workspace.iterator, self.ws_counter, Int32, 160) + if cutlass.const_expr(self.grouped): + # Images and rearm alternate over grouped calls; ``x`` is the plan. + routed = view(sm.iterator, self.off_misc + 1140, Int32, 3) + if cutlass.const_expr(self.pregrouped): + self.fetch_windows(sm, workspace) + else: + self.group(sm, ids, workspace, batch) + self.windows(sm, self.partition(routed[0], cta)) + # Retire the routing scratch's generic accesses before weight-ring reuse. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + distinct, call_count, generation = (routed[0], routed[1], routed[2]) + experts = self.tables(sm)[0] + else: + distinct, call_count, experts = self.route(sm, ids, counters, batch) + generation = call_count + layout = cute.make_layout((self.rounds, 4, 4), stride=(16, 4, 1)) + raw_x = cute.make_rmem_tensor(layout, Uint32) + if warp >= 2: + self.load_x(x, raw_x, batch) + fence_tmem() + tmem = tmem_base[0] + cute.arch.griddepcontrol_launch_dependents() + ranges = self.partition(distinct, cta) + if warp == 0: + if lane == 0: + weights = (sm, maps, w13_scales, w2_scales, ranges, experts) + if cutlass.const_expr(self.steal): + self.stream(*weights, (workspace, generation & 1)) + elif cutlass.const_expr(self.streamed): + self.stream(*weights) + else: + self.produce(*weights) + else: + if cutlass.const_expr(not self.grouped): + ready = view(sm.iterator, self.off_misc + self.bar_xready, Uint64, 14) + if warp >= 2: + self.quantize_x(sm, raw_x, ready.iterator, batch) + if warp == 1: + self.consume(sm, tmem) + elif warp == 2: + self.load_hidden(sm, workspace, ranges, generation & 1) + elif warp == 3: + if cutlass.const_expr(self.grouped): + self.gather_x(sm, x, ranges, batch) + self.rearm(workspace, generation) + else: + self.epilogue( + sm, workspace, output, tmem, ranges, batch, generation & 1, scalars + ) + release_tmem(tmem, 512, warp, 3) + if tid == 0: + counters[cta] = call_count + 1 + if cta < 160 - self.grid: + counters[self.grid + cta] = call_count + 1 + if cutlass.const_expr(self.grouped): + grouped = view(workspace.iterator, self.ws_calls, Int32, 162) + grouped[cta] = generation + 1 + if cta < 160 - self.grid: + grouped[self.grid + cta] = generation + 1 + if cta == 0: + grouped[160 + (generation & 1)] = distinct + + +class ExpertsGroupKernel(ExpertsKernel): + """One CTA groups the routes into the workspace unit table that pregrouped + ``ExpertsKernel`` read (launched after routing, beside the shared expert). Every + experts call advances all 160 call counters, so entry 0 stands for each CTA's.""" + + def __init__(self, *args, early=False): + super().__init__(*args) + self.threads = 1024 + # early: release the experts at entry, so their launch and prologue overlap + # this kernel's dependency wait and grouping (they still wait for its end). + self.early = early + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, 1, self.threads) + + @cute.kernel + def kernel(self, args): + ids, workspace, batch = args + tid = thread()[0] + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, 32768 + 16, 1024) + units = view(sm.iterator, 32768, Int32, 3) + if cutlass.const_expr(self.early): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not self.early): + cute.arch.griddepcontrol_launch_dependents() + self.group(sm, ids, workspace, batch, units) + experts, _, words = self.tables(sm, False) + table = workspace.iterator + self.ws_group + out_experts = view(table, 16, Int32, 1024) + for i in range(tid, units[0], self.threads): + out_experts[i] = experts[i] + row = tuple(words[4 * i + k] for k in range(4)) + store_words(table + 16 + 4096 + 16 * i, row) + if tid == 0: + offsets = (self.ws_counter, self.ws_calls) + calls = [Uint32(view(workspace.iterator, o, Int32, 1)[0]) for o in offsets] + store_words(table, (Uint32(units[0]), *calls, Uint32(0))) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py new file mode 100644 index 0000000000..02b46523b1 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/front.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/front.py new file mode 100644 index 0000000000..0067ed5add --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py new file mode 100644 index 0000000000..12d03fe2df --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/routing.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/routing.py new file mode 100644 index 0000000000..23bf073d23 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/shared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/shared.py new file mode 100644 index 0000000000..75e9762074 --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/chain.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/chain.py new file mode 100644 index 0000000000..35383df906 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/experts.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/experts.py new file mode 100644 index 0000000000..3376e77839 --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py new file mode 100644 index 0000000000..7e80d43048 --- /dev/null +++ b/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() diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp8.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp8.py new file mode 100644 index 0000000000..47a8a0f7c9 --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail.py new file mode 100644 index 0000000000..0ae448631d --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail.py @@ -0,0 +1,1031 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clustered MoE tail (TP8 B1..4, TP16 B1..8): routed finalization/publication, +all-reduce, RMS tree, K-split up projection, shared reduce-scatter and output multicast. +TMA weight loads overlap the mailbox wait. The caller owns the three rotating mailboxes +and must use different parities on successive calls. TP8 B5+ and TP16 B9+ run the +one-wave WaveMoeTail over latent rows published by the routed experts and shards +pushed by WaveMoeTail.push.""" + +from functools import reduce +from operator import or_ + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint32, Uint64 + +from ..common.primitives import ( + acquire_tensormap, + alloc, + bf16, + copy_g2s, + fence_tmem, + fma, + fma_bf16, + front_dirty, + init_barriers, + inv_rms, + launch, + load_f32x4, + load_tma2, + load_tma3, + load_tmem, + load_volatile, + load_words, + mailbox_store, + mma, + mma_bf16_block, + mul2, + multicast_words, + no_neg_zero, + pack_pair, + push_float, + release_tmem, + round_pair_square, + scatter_reduce, + sleep_ns, + smem_desc, + smem_words, + store_relaxed, + store_shared_words, + store_words, + thread, + view, + wait, + wait_tmem, + zeros, +) + +shuffle = cute.arch.shuffle_sync + + +def scatter_sums(values, lane, width): + """``warp_sum(values[g], width)`` of every g (4 or 8 values) as a reduce-scatter: + each butterfly level keeps half of the values and receives the partner's other + half, so every partial sum (and the total) is the butterfly's. Returns the lane's + total; value g's sits in lanes g << (log2(width) - log2(len(values))).""" + values, offset = list(values), width // 2 + while len(values) > 1: + half, upper = len(values) // 2, (lane & offset) != 0 + kept = [] + for j in range(half): + low, high = values[j], values[half + j] + keep = Float32(cutlass.select_(upper, high, low)) + send = Float32(cutlass.select_(upper, low, high)) + kept.append(keep + cute.arch.shuffle_sync_bfly(send, offset)) + values, offset = kept, offset // 2 + total = values[0] + while offset >= 1: + total += cute.arch.shuffle_sync_bfly(total, offset) + offset //= 2 + return total + + +def gather_sums(total, count, width, padded): + """Every lane's copy of ``scatter_sums``' first ``count`` totals.""" + shift = width.bit_length() - padded.bit_length() + return [shuffle(total, g << shift) for g in range(count)] + + +def spread(pointer, lane, swap): + first = smem_words(pointer + 8 * (2 * lane + swap)) + return first, smem_words(pointer + 8 * (2 * lane + 1 - swap)) + + +class MoeTail: + """One MoE tail launch over caller-owned persistent buffers. ``front_poll`` inputs + have three parity slots armed with 0xffff BF16; latent/shared mailboxes [3, TP, + capacity, 3584 or 7168/TP] are armed with 0x80000000. Only complete chains may use + the early B1 trigger, which waits for the predecessor before releasing + successors.""" + + threads = 128 + options: dict = {} + + def __init__(self, tp: int, batch: int, *, mode="front", trigger=2): + self.tp, self.batch, self.mode, self.trigger = tp, batch, mode, trigger + self.cluster, self.tokens = (7 if tp == 8 else 8), min(batch, 2) + self.poll_ns, self.front_poll_ns = (0, 0) if tp == 8 else (64, 128) + self.cols, self.shard = 3584 // self.cluster, 7168 // tp + self.frags, self.nc = self.cols // 8, self.shard // 32 + self.grid = self.cluster * self.nc * ((batch + self.tokens - 1) // self.tokens) + + @cute.jit + def __call__(self, *args): + common = dict(cluster=self.cluster, preferred_smem_carveout=100) + launch(self.kernel, args, self.grid, self.threads, **common, **self.options) + + @cute.jit + def poll(self, base, stride, second, second_stride, dual): + values = cute.make_rmem_tensor((16, 4), Uint32) + sources = (base, stride, second, second_stride) + pending = Uint32(0) + for source in cutlass.range_constexpr(16): + if self.tp == 16 or source < 8 or dual: + pending |= self.poll_source(values, source, sources) + while pending != 0: + if cutlass.const_expr(self.poll_ns): + sleep_ns(Uint32(self.poll_ns)) + next_pending = Uint32(0) + for source in cutlass.range_constexpr(16): + if (pending & (Uint32(1) << source)) != 0: + next_pending |= self.poll_source(values, source, sources) + pending = next_pending + return values + + @cute.jit + def poll_source(self, values, source: cutlass.Constexpr[int], sources): + base, stride, second, second_stride = sources + if cutlass.const_expr(self.tp == 16 or source < 8): + word = load_volatile(base + source * stride, count=4) + else: + word = load_volatile(second + (source - 8) * second_stride, count=4) + for k in cutlass.range_constexpr(4): + values[source, k] = word[k] + return Uint32(reduce(or_, (v == 0x80000000 for v in word))) << source + + @cute.jit + def sum_sources(self, values, start: cutlass.Constexpr[int]): + acc = zeros(8, Float32) + for source in cutlass.range_constexpr(self.tp): + for k in cutlass.range_constexpr(8): + acc[k] = acc[k] + bf16(values[source + start, k // 2], k % 2) + return acc + + def poll_strides(self, io): + if self.mode != "front_poll": + return 0, 0, 0 + return tuple(t.shape[0] // 3 for t in io[2:5]) + + @cute.jit + def front(self, io, latent_mc, par, rank, capacity, block, warp, lane): + _, _, g, weights, s, peers = io + if cutlass.const_expr(self.mode == "front"): + cute.arch.griddepcontrol_wait() + gs, ws, ss = self.poll_strides(io) + start = block + self.grid * warp + for task in cutlass.range(start, self.batch * 168, self.grid * 4, unroll=1): + if task < self.batch * 56: + token, line = task // 56, task % 56 + src = g.iterator + par * gs + token * 16 * 3584 + 64 * line + 2 * lane + wp = weights.iterator + par * ws + token * 16 + w = cute.make_rmem_tensor(8, Uint32) + rows = cute.make_rmem_tensor(16, Uint32) + dirty = Int32(1) + while dirty != 0: + for h in cutlass.range_constexpr(2): + for k, v in enumerate(load_volatile(wp + 8 * h, count=4)): + w[4 * h + k] = v + for j in cutlass.range_constexpr(16): + rows[j] = load_volatile(src + j * 3584)[0] + if cutlass.const_expr(self.mode == "front_poll"): + words = [w[k] for k in range(8)] + [rows[j] for j in range(16)] + bad = reduce(or_, map(front_dirty, words)) + dirty = Int32(cute.arch.vote_any_sync(bad)) + if self.front_poll_ns > 0 and dirty != 0: + sleep_ns(Uint32(self.front_poll_ns)) + else: + dirty = 0 + a0, a1 = Float32(0), Float32(0) + for j in cutlass.range_constexpr(16): + weight = bf16(w[j // 2], j % 2) + a0 = fma(bf16(rows[j]), weight, a0) + a1 = fma(bf16(rows[j], True), weight, a1) + word = pack_pair(a0, a1, policy="latent") + packed = tuple(shuffle(word, (4 * lane + k) & 31) for k in range(4)) + if lane < 8: + row = (par * self.tp + rank) * capacity + token + offset = row * 3584 + 64 * line + 8 * lane + multicast_words(latent_mc + Int64(offset) * 2, packed) + elif lane < 8: + u = task - self.batch * 56 + token, line = u // 112, u % 112 + destination = line // (self.shard // 64) + own_line = line % (self.shard // 64) + src = s.iterator + par * ss + token * 7168 + 64 * line + 8 * lane + word = load_volatile(src, count=4) + if cutlass.const_expr(self.mode == "front_poll"): + while reduce(or_, map(front_dirty, word)): + if cutlass.const_expr(self.front_poll_ns): + sleep_ns(Uint32(self.front_poll_ns)) + word = load_volatile(src, count=4) + packed = tuple(no_neg_zero(v) for v in word) + row = (par * self.tp + rank) * capacity + token + offset = row * self.shard + 64 * own_line + 8 * lane + store_relaxed(peers[destination] + Int64(offset) * 2, packed) + + @cute.jit + def gemv(self, sw, sx, partial, bars, q, warp, lane, mt: cutlass.Constexpr[int]): + weight, xs, out = sw.iterator, sx.iterator, partial.iterator + barrier = bars.iterator + 2 + acc = zeros(8 * mt, Float32) + if cutlass.const_expr(self.tp == 8): + # Two bank-spread loads; columns 16*lane .. 16*lane+15 accumulate in order. + swap = (lane >> 2) & 1 + xl = cute.make_rmem_tensor((mt, 8), Uint32) + for token in cutlass.range_constexpr(mt): + aa, bb = spread(xs + token * self.cols, lane, swap) + for k in cutlass.range_constexpr(4): + xl[token, k] = bb[k] if swap != 0 else aa[k] + xl[token, k + 4] = aa[k] if swap != 0 else bb[k] + for row in cutlass.range_constexpr(8): + wa, wb = spread(weight + (warp * 8 + row) * self.cols, lane, swap) + for token in cutlass.range_constexpr(mt): + value = acc[row * mt + token] + for k in cutlass.range_constexpr(8): + ww = (wb + wa)[k] if swap != 0 else (wa + wb)[k] + value = fma_bf16(ww, xl[token, k], value) + value = fma_bf16(ww, xl[token, k], value, high=True) + acc[row * mt + token] = value + elif lane < 28: + for half in cutlass.range_constexpr(2): + chunk = lane + 28 * half + x = [smem_words(xs + t * self.cols + 8 * chunk) for t in range(mt)] + for row in cutlass.range_constexpr(8): + ww = smem_words(weight + (warp * 8 + row) * self.cols + 8 * chunk) + for token in cutlass.range_constexpr(mt): + value = acc[row * mt + token] + for k in cutlass.range_constexpr(4): + value = fma_bf16(ww[k], x[token][k], value) + value = fma_bf16(ww[k], x[token][k], value, high=True) + acc[row * mt + token] = value + # One final row/token per lane group, immediately pushed to the cluster leader. + value, selected = scatter_reduce(acc, lane, 8 * mt) + if (lane & ((4 if mt == 1 else 2) - 1)) == 0: + row, token = selected // mt, selected % mt + push_float(out + (q * 32 + warp * 8 + row) * mt + token, value, barrier, 0) + + @cute.kernel + def kernel(self, args): + latent, shared, weight_map, gamma, gemm2, weights, shared_output = args[0:7] + peers, output, latent_mc, output_mc, par, rank, capacity, eps = args[7:15] + tid, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + q = cute.arch.block_idx_in_cluster() + warp, lane = tid // 32, tid % 32 + cluster = block // self.cluster + group, row_cluster = cluster // self.nc, cluster % self.nc + token0, row0 = group * self.tokens, row_cluster * 32 + local_tokens = min(self.tokens, self.batch - token0) + allocator = cutlass.utils.SmemAllocator() + sw = alloc(allocator, BFloat16, 32 * self.cols, 128) + sg = alloc(allocator, BFloat16, self.cols) + sx = alloc(allocator, BFloat16, self.tokens * self.cols) + partial = alloc(allocator, Float32, self.cluster * 32 * self.tokens) + squares = alloc(allocator, Float32, self.cluster * 2 * self.tokens) + shared_sum = alloc(allocator, Float32, self.tokens * 32) + bars = alloc(allocator, Uint64, 4, 8) + bar, expect = bars.iterator, cute.arch.mbarrier_arrive_and_expect_tx + if tid == 127: + init_barriers(bar, 4) + expect(bar + 3, self.cols * 2) + gamma_q = gamma.iterator + q * self.cols + copy_g2s(sg.iterator, gamma_q, size=self.cols * 2, barrier=bar + 3) + expect(bar + 1, self.cluster * 2 * local_tokens * 4) + expect(bar + 2, self.cluster * 32 * local_tokens * 4 if q == 0 else 0) + expect(bar, 32 * self.cols * 2) + acquire_tensormap(weight_map.iterator) + load_tma3(sw.iterator, weight_map.iterator, bar, 2 * q, row0) + cute.arch.cluster_arrive_relaxed() + io = (latent, shared, gemm2, weights, shared_output, peers) + self.front(io, latent_mc, par, rank, capacity, block, warp, lane) + pt, pk = tid // 64, tid % 64 + lat_thread = (pt < local_tokens) & (pk < self.frags) + rs_j = tid - 64 if self.tokens == 1 else pk if self.tp == 8 else pk - self.frags + rs_thread = (q == 0) & (rs_j >= 0) & (rs_j < 4) + if cutlass.const_expr(self.tokens != 1): + rs_thread &= pt < local_tokens + token = Int32(0) if cutlass.const_expr(self.tokens == 1) else pt + s2 = Float32(0) + if lat_thread | rs_thread: + slot = par * self.tp * capacity + token0 + token + lat_ptr = latent.iterator + (slot * 3584 + q * self.cols + 8 * pk) + rs_ptr = shared.iterator + (slot * self.shard + row0 + 8 * max(rs_j, 0)) + base = lat_ptr if lat_thread else rs_ptr + stride = capacity * (3584 if lat_thread else self.shard) + dual = rs_thread & lat_thread + values = self.poll(base, stride, rs_ptr, capacity * self.shard, dual) + sums = self.sum_sources(values, 0) + if lat_thread: + packed = [] + for k in cutlass.range_constexpr(4): + packed.append(pack_pair(sums[2 * k], sums[2 * k + 1])) + square = round_pair_square(packed[k]) + s2 += bf16(square) + s2 += bf16(square, True) + store_shared_words(sx.iterator + token * self.cols + pk * 8, packed) + else: + for k in cutlass.range_constexpr(8): + shared_sum[token * 32 + 8 * rs_j + k] = Float32(BFloat16(sums[k])) + if cutlass.const_expr(self.tp == 8 and self.tokens == 2): # noqa: SIM102 + if dual: + rs = self.sum_sources(values, 8) + for k in cutlass.range_constexpr(8): + shared_sum[token * 32 + 8 * rs_j + k] = Float32(BFloat16(rs[k])) + cute.arch.sync_warp() + for stage in cutlass.range_constexpr(5): + s2 += cute.arch.shuffle_sync_bfly(s2, 16 >> stage) + cute.arch.cluster_wait() + if pt < local_tokens and lane < self.cluster: + sq = squares.iterator + (2 * q + (warp & 1)) * local_tokens + pt + push_float(sq, s2, bar + 1, lane) + cute.arch.sync_warp() + wait(bar + 1) + cute.arch.sync_warp() + if cutlass.const_expr(self.trigger == 3): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + wait(bar + 3) + if lat_thread: + total = Float32(0) + for peer in cutlass.range_constexpr(2 * self.cluster): + total += squares[peer * local_tokens + pt] + inverse = inv_rms(total, eps) + gx = smem_words(sg.iterator + pk * 8) + xx = smem_words(sx.iterator + pt * self.cols + pk * 8) + scaled = lambda k, h: bf16(xx[k], h) * inverse * bf16(gx[k], h) + packed = [pack_pair(scaled(k, 0), scaled(k, 1)) for k in range(4)] + store_shared_words(sx.iterator + pt * self.cols + pk * 8, packed) + wait(bar) + cute.arch.sync_warp() + cute.arch.sync_threads() + if self.tokens == 1 or self.batch % 2 == 0 or local_tokens == 2: + self.gemv(sw, sx, partial, bars, q, warp, lane, self.tokens) + else: + self.gemv(sw, sx, partial, bars, q, warp, lane, 1) + if q == 0: # noqa: SIM102 - Keep the owner CTA branch separate from its warp selection. + if warp < local_tokens: + wait(bar + 2) + cute.arch.sync_warp() + value = Float32(0) + for peer in cutlass.range_constexpr(self.cluster): + value += partial[(peer * 32 + lane) * local_tokens + warp] + rounded = Float32(BFloat16(value)) + out = Float32(BFloat16(rounded + shared_sum[warp * 32 + lane])) + # Pack paired rows, then four pairs into a 16-byte writer. + pair = pack_pair(out, shuffle(out, (lane + 1) & 31), policy="latent") + words = tuple(shuffle(pair, (8 * lane + 2 * k) & 31) for k in range(4)) + if lane < 4: + offset = (token0 + warp) * 7168 + rank * self.shard + row0 + multicast_words(output_mc + Int64(offset + 8 * lane) * 2, words) + if cutlass.const_expr(self.trigger == 2): + # front_poll: the next AttnRes reads the prefix before any wait, and no + # kernel since the o-proj consumer that wrote it (the MoE front and + # experts) waits before releasing: this grid must, at every batch. + if cutlass.const_expr(self.mode == "front_poll"): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + self.rearm(io, capacity, *args[15:17], block, tid) + # Orders deferred rearm, workspace reuse and producer lifetimes (all modes). + cute.arch.griddepcontrol_wait() + + @cute.jit + def rearm(self, io, capacity, parity, batch, block, tid, reduced=None): + """Re-arms ``parity``'s first ``batch`` tokens (and WaveMoeTail's reduced + rows).""" + latent, shared, g, w, s, _ = io + if parity >= 0: + nl, nr = self.tp * batch * 448, self.tp * batch * (self.shard // 8) + ng, nw, ns = Int32(0), Int32(0), Int32(0) + if cutlass.const_expr(self.mode == "front_poll"): + gs, ws, ss = self.poll_strides(io) + ng = min(batch, gs // (16 * 3584)) * 16 * 448 + nw = min(batch, ws // 16) * 2 + ns = min(batch, ss // 7168) * 896 + elif cutlass.const_expr(reduced is not None): + ng = batch * 448 + ring = (parity, capacity, batch) + start, step = block * self.threads + tid, self.grid * self.threads + for i in cutlass.range(start, nl + nr + ng + nw + ns, step, unroll=1): + if i < nl: + self.rearm_row(latent, 3584, i, ring) + elif i < nl + nr: + self.rearm_row(shared, self.shard, i - nl, ring) + elif cutlass.const_expr(self.mode == "front_poll"): + j = i - nl - nr + fill = (Uint32(0xFFFFFFFF),) * 4 + if j < ng: + store_words(g.iterator + parity * gs + 8 * j, fill) + elif j < ng + nw: + store_words(w.iterator + parity * ws + 8 * (j - ng), fill) + else: + store_words(s.iterator + parity * ss + 8 * (j - ng - nw), fill) + elif cutlass.const_expr(reduced is not None): + rows = reduced.iterator + parity * capacity * 3584 + store_words(rows + 8 * (i - nl - nr), (Uint32(0x80000000),) * 4) + + @cute.jit + def rearm_row(self, box, width: cutlass.Constexpr[int], i, ring): + parity, capacity, batch = ring + row, frag = i // (width // 8), i % (width // 8) + source, token = row // batch, row % batch + dst = box.iterator + ((parity * self.tp + source) * capacity + token) * width + store_words(dst + 8 * frag, (Uint32(0x80000000),) * 4) + + +def pow2_spans(n): + """(start, count) spans of n columns, power-of-two counts, largest first.""" + bits = [1 << i for i in reversed(range(n.bit_length())) if n >> i & 1] + return tuple((sum(bits[:i]), b) for i, b in enumerate(bits)) + + +class WaveMoeTail(MoeTail): + """One-wave tail over published inputs (TP8 B5+, TP16 B9+): seven 16-CTA clusters + split K into 224-column slices of a 128-row (TP16: 64-row) tile. Weight fragments + stay in registers through the latent poll, RMS all-gather and normalization; row + warps (TP16: two sets taking alternate groups) run two seven-step BF16 MMA chains + per 8-token group, and each 64-row output line has one owner CTA and one multicast + writer. Past one group (and at TP16) each cluster instead all-reduces every seventh + token of its slices once and shares them as BF16 rows through ``reduced``, so all + 112 CTAs must be co-resident; RS lines spread over warp pairs, every group's RMS + sums travel as one vector per peer, all groups normalize before any MMA, the + groups' MMA chains interleave per K step, and the previous layer's rows are + re-armed after the output multicast. ``twoshot`` (B9+): the latent + mailbox is this rank's ``common.twoshot`` receive slots [3, TP, capacity, 3584/TP] + (pushed by the experts), the CTAs of its owned slices reduce them like TP16 but + multicast the rows into every rank's ``reduced`` (the gather; its multicast + address replaces the latent one), so each GPU receives 2 x B x 7 KB of latent + instead of TP x B x 7 KB.""" + + threads, options = 256, dict(min_blocks_per_mp=1) + + def __init__(self, batch, tp=8, twoshot=False, early=False, umma=False): + super().__init__(tp, batch, mode="publish") + assert not twoshot or batch > 8 + assert not umma or batch > 8 + # ``early``: release the next AttnRes at entry, so its block statistics and + # exchange overlap this kernel (it reads only earlier layers' data before its + # Lamport poll of the output rows), instead of after the output multicast. + self.early = early + self.lines, self.groups = 2 if tp == 8 else 1, (batch + 7) // 8 + self.slots = (batch * self.lines + 15) // 16 # owned output lines per CTA + self.twoshot, self.width = twoshot, 3584 // tp if twoshot else 3584 + # Past one group (and at TP16) clusters share reduced latent rows; TP8 B<=8 + # polls its sources directly. + self.reduce = self.groups > 1 or tp == 16 + self.grid, self.poll_ns, self.cluster = 112, 64, 16 + # Shared memory: per group x [8, 224] BF16 and MMA partials [8 warps, 8, 16] + # (TP16 pads to whole warp sets); gathered squares (one group keeps each + # slice's 28 lane sums, more one sum); per owned line its partials [16 slices, + # 64] and RS rows [64]; two barriers. + n, s = self.groups + self.groups % (2 // self.lines), self.slots + # Past one group each (warp, slice) row holds the groups' sums, padded to v4. + self.padded = -(-self.groups // 4) * 4 + self.squares = 3584 if n == 1 else 128 * max(n, self.padded) + partials = 7680 * n + 4 * self.squares + self.offsets: tuple = (3584 * n, 7680 * n, partials, partials + 4096 * s) + self.offsets += (partials + 4352 * s + 32,) + self.smem = 116 * 1024 # one CTA per SM + # ``umma`` (past one group): the CTA's [M, 224] x [N, 224] up projection on + # tcgen05 (M = 64 x lines rows, N = the padded tokens, K padded to 256) from + # 128B-swizzled W and x tiles (W: four TMA boxes [M, 64] at entry, the last + # running into the next slice, whose columns meet the zeroed x padding; the + # weight argument is then the shard's tensor map), D in TMEM, staged as + # [token, row] in the drained W tiles for the owners' pushes. FP32 + # accumulation in the tensor core's order. + self.umma = umma + if umma: + self.m = 64 * self.lines + self.n = -(-8 * self.groups // (16 if self.lines == 2 else 8)) + self.n *= 16 if self.lines == 2 else 8 + self.columns = max(32, 1 << (self.n - 1).bit_length()) + w_bytes, x_bytes = self.m * 512, self.n * 512 + squares = w_bytes + x_bytes + partials = squares + 4 * self.squares + self.offsets = (0, squares, partials, partials + 4096 * s) + self.offsets += (partials + 4352 * s + 32,) + self.smem = max(self.smem, self.offsets[4] + 40) + + def group_smem(self, sm, g): + """Group g's x (BF16 pointer) and MMA partials.""" + x = view(sm, 3584 * g, BFloat16, 1792).iterator + return x, view(sm, self.offsets[0] + 4096 * g, Float32, 1024) + + def line_smem(self, sm, s): + """Owned line s's partials [16 slices, 64] and RS rows.""" + partial, rs = self.offsets[2:4] + cells = view(sm, partial + 4096 * s, Float32, 1024) + return cells, view(sm, rs + 256 * s, Float32, 64) + + @cute.kernel + def kernel(self, args): + latent, shared, weight, gamma = args[0:4] + output_mc, par, rank, capacity, eps = args[10:15] + tid, warp, lane, block = thread() + if cutlass.const_expr(self.early): + cute.arch.griddepcontrol_launch_dependents() + q = cute.arch.block_idx_in_cluster() + n, batch, lines, rows = self.groups, self.batch, self.lines, 64 * self.lines + row0, gq, tq, owner = block // 16 * rows, lane // 4, lane % 4, q < lines * batch + own_token, own_line = q // lines, q % lines + rs_j, lat = 4 * warp + lane - 28, lane < 28 + owns = lambda s: owner if s == 0 else q < lines * batch - 16 * s + rs_lane = lambda g: (lane >= 28) & owns(g) & (rs_j < 8) + align = 1024 if self.umma else 128 + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem, align).iterator + groups = [self.group_smem(sm, g) for g in range(n)] + owned = [self.line_smem(sm, s) for s in range(self.slots)] + squares = view(sm, self.offsets[1], Float32, self.squares) + bar = view(sm, self.offsets[4], Uint64, 4).iterator + tmem_slot = view(sm, self.offsets[4] + 32, Uint32, 1) + if cutlass.const_expr(self.umma): # noqa: SIM102 + if warp == 0: + cute.arch.alloc_tmem(self.columns, tmem_slot.iterator) + cute.arch.relinquish_tmem_alloc_permit() + if tid == 0: + init_barriers(bar, 4 if self.umma else 2) + if cutlass.const_expr(self.umma): # W [M, 64] x 4 by TMA on barrier 3 + acquire_tensormap(weight.iterator) + cute.arch.mbarrier_arrive_and_expect_tx(bar + 3, self.m * 512) + for kb in cutlass.range_constexpr(4): + tile = sm + kb * (self.m * 128) + load_tma2(tile, weight.iterator, bar + 3, 224 * q + 64 * kb, row0) + # One group gathers 28 lane sums per (token, slice), more one sum for + # every warp's token. + gathered = 16 * batch * 28 * 4 if n == 1 else 16 * 8 * self.padded * 4 + cute.arch.mbarrier_arrive_and_expect_tx(bar, gathered) + size = 4096 if owns(self.slots - 1) else 0 # earlier lines are all owned + size = size + 4096 * (self.slots - 1) if self.slots > 1 else size + cute.arch.mbarrier_arrive_and_expect_tx(bar + 1, size) + cute.arch.cluster_arrive_relaxed() + xs = groups[0][0] + if cutlass.const_expr(not self.umma): + for i in cutlass.range(tid, (8 * n - batch) * 28, 256, unroll=1): + store_shared_words(xs + batch * 224 + 8 * i, (Uint32(0),) * 4) + gx, gamma_q = cute.make_rmem_tensor(4, Uint32), gamma.iterator + q * 224 + if lat & (warp < batch): + for k, v in enumerate(load_words(gamma_q + 8 * lane, count=4)): + gx[k] = v + # Rows 16 warp + gq (+ 8) of the cluster's tile (TP16: warp mod 4), columns + # 224 q + 32 k + 8 tq. + rw = warp if lines == 2 else warp % 4 + base = weight.iterator + (row0 + 16 * rw + gq) * 3584 + 224 * q + 8 * tq + fetch = lambda p: load_words(p, count=4, evict_first=True, mutable=True) + a = None + if cutlass.const_expr(not self.umma): + a = [ + [fetch(base + h * 8 * 3584 + 32 * k) for h in range(2)] + for k in range(7) + ] + xwords = [cute.make_rmem_tensor(4, Uint32) for _ in range(n)] + reduced = args[17] if self.reduce else None + ring = (*args[15:17], block, tid, reduced) + if cutlass.const_expr(self.reduce): + s2 = zeros(n, Float32) + ids: tuple = (q, warp, lane, block // 16) + places = (row0, owned[0][1].iterator) + gather = (args[9], rank) if self.twoshot else None + boxes = (latent, reduced, shared) + self.poll_inputs(boxes, (par, capacity), ids, places, gather) + if lat: + self.read_reduced(reduced, par * capacity, xwords, s2, ids) + else: + # Lanes 0..27 of warp t poll token t's latent; owners' lanes 28..31 the RS. + sq = Float32(0) + if lat & (warp < batch) | rs_lane(0): + slot = par * self.tp * capacity + rs = (slot + own_token) * self.shard + row0 + 64 * own_line + rs += 8 * max(rs_j, 0) + ptr = latent.iterator + (slot + warp) * 3584 + 224 * q + 8 * lane + ptr = ptr if lat else shared.iterator + rs + stride = capacity * (3584 if lat else self.shard) + sums = self.sum_sources(self.poll(ptr, stride, ptr, 0, False), 0) + if lat: + sq = self.latent_words(sums, xwords[0]) + else: + self.rs_words(sums, owned[0][1], rs_j) + s2 = [sq] + cute.arch.sync_warp() + cute.arch.cluster_wait() + if cutlass.const_expr(n == 1): + self.push_squares(squares, s2[0], bar, (q, warp, lane)) + wait(bar) + cute.arch.sync_warp() + row = (xwords[0], gx, eps) + self.normalize(squares, row, groups[0][0], (warp, lane, lat)) + else: + self.exchange_squares(squares, s2, bar, (q, warp, lane)) + wait(bar) + cute.arch.sync_warp() + outs = [groups[g][0] for g in range(n)] + if cutlass.const_expr(self.umma): # x tiles [N, 64] x 4 after W + outs = [sm + self.m * 512 + 1024 * g for g in range(n)] + self.normalize_groups(squares, (xwords, gx, eps), outs, (warp, lane, lat)) + if cutlass.const_expr(self.umma): # x stores -> tcgen05 + cute.arch.fence_proxy("async.shared", space="cta") + fence_tmem(before=True) + cute.arch.sync_warp() + cute.arch.sync_threads() + ids = (q, warp, lane, rw, gq, tq) + if cutlass.const_expr(n == 1): + self.mma_group(a, groups[0], owned, bar, 0, ids) + elif cutlass.const_expr(self.umma): + tmem = tmem_slot[0] + self.umma_partials(sm, tmem, bar, (q, warp, lane, tid)) + else: + size = (self.offsets[1] - self.offsets[0]) // 4 + products = view(sm, self.offsets[0], Float32, size) + self.mma_groups(a, xs, products, owned, bar, ids) + line = (lane, rank, row0, own_token, own_line) + for s in cutlass.range_constexpr(self.slots): + if owns(s) & (warp == s): + self.write_line(owned[s], bar, output_mc, line, s) + if cutlass.const_expr(not self.early): + cute.arch.griddepcontrol_launch_dependents() + # The previous layer's rows, off the output path. + if cutlass.const_expr(n == 1): + self.rearm(args[0:2] + args[4:8], capacity, *ring) + else: + self.rearm_rows(args[0:2], capacity, *ring) + if cutlass.const_expr(self.umma): + release_tmem(tmem, self.columns, warp, 0) + cute.arch.griddepcontrol_wait() + + @cute.jit + def load_sources(self, values, ptr, stride): + """Issue the TP source loads of a lane's 16-byte column without waiting.""" + for source in cutlass.range_constexpr(self.tp): + word = load_volatile(ptr + source * stride, count=4) + for k in cutlass.range_constexpr(4): + values[source, k] = word[k] + + @cute.jit + def settle(self, values, ptr, stride): + """Re-poll loaded sources until no word is the 0x80000000 sentinel.""" + pending = Uint32(0) + for source in cutlass.range_constexpr(self.tp): + words = (values[source, k] for k in range(4)) + pending |= Uint32(reduce(or_, (v == 0x80000000 for v in words))) << source + while pending != 0: + sleep_ns(Uint32(self.poll_ns)) + again = Uint32(0) + for source in cutlass.range_constexpr(self.tp): + if (pending & (Uint32(1) << source)) != 0: + again |= self.poll_row(values, source, ptr + source * stride) + pending = again + + @cute.jit + def column_sums(self, values): + """``sum_sources`` (same order and rounding) as one BF16 FMA by 1.0 per term.""" + acc, one = zeros(8, Float32), Uint32(0x3F803F80) + for source in cutlass.range_constexpr(self.tp): + for k in cutlass.range_constexpr(8): + acc[k] = fma_bf16(values[source, k // 2], one, acc[k], high=k % 2 == 1) + return acc + + @cute.jit + def rearm_rows(self, boxes, capacity, parity, batch, block, tid, reduced): + """Division-free ``rearm``: each source's latent and RS rows, then the reduced + rows, of the previous layer (complete before this launch).""" + if parity >= 0: + start, step = block * self.threads + tid, self.grid * self.threads + fill, widths = (Uint32(0x80000000),) * 4, (self.width, self.shard) + for s in cutlass.range_constexpr(self.tp): + first = (parity * self.tp + s) * capacity + for b in cutlass.range_constexpr(2): + row, w = boxes[b].iterator + first * widths[b], widths[b] + for i in cutlass.range(start, batch * (w // 8), step, unroll=1): + store_words(row + 8 * i, fill) + if cutlass.const_expr(reduced is not None): + row = reduced.iterator + parity * capacity * 3584 + for i in cutlass.range(start, batch * 448, step, unroll=1): + store_words(row + 8 * i, fill) + + @cute.jit + def poll_inputs(self, boxes, ring, ids, places, gather=None): + """Phase 1 and the RS lines in one convergent pass (every round's loads issued + first). Lanes 0..27 all-reduce the cluster's tokens tile + 7 j of the CTA's + slice once into BF16 ``reduced`` rows (``gather`` (multicast address, rank): + only the CTAs of owned slices, from the receive slots into every rank's); + lanes 28..31 of warps 2 p and 2 p + 1 sum owned lines p, p + 4's RS shard + (line 16 i + q: token (16 i + q) / lines).""" + (latent, reduced, shared), (par, capacity) = boxes, ring + (q, warp, lane, tile), (row0, rs_rows) = ids, places + lat, width = lane < 28, self.width + col = 224 * q + 8 * lane + local = col if gather is None else col - gather[1] * width + src = latent.iterator + (par * self.tp * capacity + tile) * width + local + dst = reduced.iterator + (par * capacity + tile) * 3584 + col + reduces = q >= 0 if gather is None else q * self.tp // 16 == gather[1] + rs_j = 4 * (warp % 2) + lane - 28 + rs = shared.iterator + par * self.tp * capacity * self.shard + row0 + 8 * rs_j + stride = capacity * width if lat else capacity * self.shard + latent_rounds = ((self.batch - 1) // 7 + 8) // 8 + rounds = max(latent_rounds, (self.slots + 3) // 4) + tokens = [7 * (warp + 8 * r) for r in range(rounds)] + lines = [4 * r + warp // 2 for r in range(rounds)] + live, ptrs = [], [] + for r in cutlass.range_constexpr(rounds): + token, line = tokens[r], lines[r] + on = lane < 0 # no latent round r + if cutlass.const_expr(r < latent_rounds): + on = reduces & (tile + token < self.batch) + rs_on = (line < self.slots) & (16 * line + q < self.lines * self.batch) + live.append(on if lat else rs_on) + rs_row = rs + (16 * line + q) // self.lines * self.shard + rs_row += 64 * ((16 * line + q) % self.lines) + ptrs.append(src + token * width if lat else rs_row) + values = [cute.make_rmem_tensor((self.tp, 4), Uint32) for _ in range(rounds)] + for r in cutlass.range_constexpr(rounds): + if live[r]: + self.load_sources(values[r], ptrs[r], stride) + for r in cutlass.range_constexpr(rounds): + if live[r]: + self.settle(values[r], ptrs[r], stride) + sums = self.column_sums(values[r]) + if lat: + if cutlass.const_expr(gather is None): + self.store_reduced(sums, dst + tokens[r] * 3584) + else: + offset = (par * capacity + tile + tokens[r]) * 3584 + col + self.store_reduced(sums, gather[0] + Int64(offset) * 2, True) + else: + rows = view(rs_rows, 64 * lines[r], Float32, 64) + self.rs_words(sums, rows, rs_j) + + @cute.jit + def store_reduced(self, sums, dst, multicast: cutlass.Constexpr[bool] = False): + pack = lambda j: pack_pair(sums[2 * j], sums[2 * j + 1], policy="latent") + if cutlass.const_expr(multicast): + multicast_words(dst, [pack(j) for j in range(4)]) + else: + mailbox_store(dst, [pack(j) for j in range(4)]) + + @cute.jit + def read_reduced(self, reduced, slot, xwords, s2, ids): + """Phase 2: the lane's words of tokens 8 g + warp and their squares.""" + q, warp, lane, _ = ids + base = reduced.iterator + (slot + warp) * 3584 + 224 * q + 8 * lane + values = self.poll_rows(base, 8 * 3584, (self.batch - warp + 7) // 8) + for g in cutlass.range_constexpr(self.groups): + for k in cutlass.range_constexpr(4): + xwords[g][k] = values[g, k] + s2[g] = self.square_sum(xwords[g]) + + @cute.jit + def poll_rows(self, base, stride, count): + """Lamport-poll the first ``count`` 16-byte rows base + r stride.""" + values = cute.make_rmem_tensor((self.groups, 4), Uint32) + pending = Uint32(0) + for r in cutlass.range_constexpr(self.groups): + if r < count: + pending |= self.poll_row(values, r, base + r * stride) + while pending != 0: + sleep_ns(Uint32(self.poll_ns)) + again = Uint32(0) + for r in cutlass.range_constexpr(self.groups): + if (pending & (Uint32(1) << r)) != 0: + again |= self.poll_row(values, r, base + r * stride) + pending = again + return values + + @cute.jit + def poll_row(self, values, r: cutlass.Constexpr[int], ptr): + word = load_volatile(ptr, count=4) + for k in cutlass.range_constexpr(4): + values[r, k] = word[k] + return Uint32(reduce(or_, (v == 0x80000000 for v in word))) << r + + @cute.jit + def mma_group(self, a, group, owned, bar, g: cutlass.Constexpr[int], ids): + """Group g's MMA chains; each 4-row partial goes to its 64-row line's owner.""" + (q, warp, lane, rw, gq, tq), (xg, pg) = ids, group + b = [smem_words(xg + gq * 224 + 8 * tq + 32 * k) for k in range(7)] + d = e = (Float32(0),) * 4 + for k in cutlass.range_constexpr(7): + (w0, w1), x = a[k], b[k] + d = mma((w0[0], w1[0], w0[1], w1[1]), x[0:2], d, bf16=True) + e = mma((w0[2], w1[2], w0[3], w1[3]), x[2:4], e, bf16=True) + for i in cutlass.range_constexpr(4): # [warp, token, row] partial products + pg[warp * 128 + (2 * tq + i % 2) * 16 + gq + 8 * (i // 2)] = d[i] + e[i] + cute.arch.sync_warp() + token, row = lane // 4, 16 * rw + 4 * (lane % 4) + if token < self.batch - 8 * g: + values = load_f32x4(pg.iterator + warp * 128 + 4 * lane, shared=True) + # Line lines (8 g + token) + row // 64 is owned by CTA (line % 16). + peer = self.lines * token + row // 64 + 8 * g * self.lines % 16 + cells = owned[8 * g * self.lines // 16][0] + push_float(cells.iterator + q * 64 + row % 64, values, bar + 1, peer) + + @cute.jit + def mma_groups(self, a, xs, products, owned, bar, ids): + """``mma_group`` over this warp's groups g = sets j + set (TP16: two warp + sets): every MMA chain (the groups' chains interleaved per K step), one warp + sync, then the owners' pushes (slot j).""" + q, warp, lane, rw, gq, tq = ids + sets = 2 // self.lines + h = 0 if sets == 1 else warp // 4 + mine = (self.groups + sets - 1) // sets + x = [xs + (sets * j + h) * 1792 + gq * 224 + 8 * tq for j in range(mine)] + d = [(Float32(0),) * 4 for _ in range(mine)] + e = [(Float32(0),) * 4 for _ in range(mine)] + for k in cutlass.range_constexpr(7): + w0, w1 = a[k] + for j in cutlass.range_constexpr(mine): + xk = smem_words(x[j] + 32 * k) + d[j] = mma((w0[0], w1[0], w0[1], w1[1]), xk[0:2], d[j], bf16=True) + e[j] = mma((w0[2], w1[2], w0[3], w1[3]), xk[2:4], e[j], bf16=True) + acc = [[d[j][i] + e[j][i] for i in range(4)] for j in range(mine)] + for j in cutlass.range_constexpr(mine): + cell = (sets * j + h) * 1024 + warp * 128 + 2 * tq * 16 + gq + for i in cutlass.range_constexpr(4): + products[cell + i % 2 * 16 + 8 * (i // 2)] = acc[j][i] + cute.arch.sync_warp() + token, row = lane // 4, 16 * rw + 4 * (lane % 4) + peer = self.lines * token + row // 64 + (8 * h if sets > 1 else 0) + for j in cutlass.range_constexpr(mine): + if token < self.batch - 8 * (sets * j + h): + src = products.iterator + (sets * j + h) * 1024 + warp * 128 + 4 * lane + values = load_f32x4(src, shared=True) + cells = owned[j][0] + push_float(cells.iterator + q * 64 + row % 64, values, bar + 1, peer) + + @cute.jit + def write_line(self, line, bar, output_mc, ids, s: cutlass.Constexpr[int]): + """Owned line s: sum the 16 slices' partials, add the RS rows and multicast the + 64 output rows.""" + (cells, rs_sum), (lane, rank, row0, own_token, own_line) = line, ids + wait(bar + 1) + cute.arch.sync_warp() + out = [] + for h in cutlass.range_constexpr(2): + parts = (cells[64 * p + 2 * lane + h] for p in range(16)) + value = Float32(BFloat16(sum(parts, Float32(0)))) + out.append(value + rs_sum[2 * lane + h]) + pair = pack_pair(*out, policy="latent") + words = tuple(shuffle(pair, (4 * lane + k) & 31) for k in range(4)) + if lane < 8: + t = 16 // self.lines * s + own_token + offset = t * 7168 + rank * self.shard + row0 + 64 * own_line + multicast_words(output_mc + Int64(offset + 8 * lane) * 2, words) + + @cute.jit + def latent_words(self, sums, words): + """BF16 latent words of a lane's eight columns; returns their square sum.""" + s2 = Float32(0) + for j in cutlass.range_constexpr(4): + words[j] = pack_pair(sums[2 * j], sums[2 * j + 1]) + square = round_pair_square(words[j]) + s2 += bf16(square) + s2 += bf16(square, True) + return s2 + + @cute.jit + def square_sum(self, words): + s2 = Float32(0) + for j in cutlass.range_constexpr(4): + square = round_pair_square(words[j]) + s2 += bf16(square) + s2 += bf16(square, True) + return s2 + + @cute.jit + def rs_words(self, sums, rs_sum, rs_j): + for j in cutlass.range_constexpr(8): + rs_sum[8 * rs_j + j] = Float32(BFloat16(sums[j])) + + @cute.jit + def push_squares(self, squares, s2, bar, ids): + """One group: the warp's token's 28 lane sums of every slice to every peer.""" + q, warp, lane = ids + if warp < self.batch: + chunk, first = lane % 7, lane // 7 + vector = tuple(shuffle(s2, 4 * chunk + i) for i in range(4)) + if first < 4: + dst = squares.iterator + warp * 448 + 28 * q + 4 * chunk + for peer in cutlass.range(first, 16, 4, unroll=1): + push_float(dst, vector, bar, peer) + + @cute.jit + def exchange_squares(self, squares, s2, bar, ids): + """Past one group: every group's warp sum for the warp's token; lane p < 16 + pushes them to peer p's (warp, slice) row as v4 words.""" + q, warp, lane = ids + values = [s2[g] for g in range(self.groups)] + values += [Float32(0)] * (self.padded - self.groups) + total = scatter_sums(values, lane, 32) + totals = gather_sums(total, self.padded, 32, self.padded) + if lane < 16: + dst = squares.iterator + (warp * 16 + q) * self.padded + for c in cutlass.range_constexpr(self.padded // 4): + push_float(dst + 4 * c, tuple(totals[4 * c : 4 * c + 4]), bar, lane) + + @cute.jit + def normalize_groups(self, squares, row, xs, ids): + """Past one group: every group's RMS (the 16 slices' sums, ``square_total``'s + order) first, then the scaled latent words (padding tokens too).""" + (words, gx, eps), (warp, lane, lat) = row, ids + base = squares.iterator + (warp * 16 + lane % 16) * self.padded + sums = [] + for c in cutlass.range_constexpr(self.padded // 4): + sums += list(load_f32x4(base + 4 * c, shared=True)) + inverse = inv_rms(scatter_sums(sums, lane, 16), eps, m8=True) + inverse = gather_sums(inverse, self.groups, 16, self.padded) + if lat: + for g in cutlass.range_constexpr(self.groups): + self.scale_words(words[g], gx, inverse[g], xs[g], warp, lane) + else: + if cutlass.const_expr(self.umma): # K columns 224..255 of the x tiles + for g in cutlass.range_constexpr(self.groups): + dst = self.x_chunk(xs[g], warp, lane) + store_shared_words(dst, (Uint32(0),) * 4) + + @cute.jit + def normalize(self, squares, row, xs, ids): + """One group: RMS-normalize the warp's token into ``xs``.""" + (words, gx, eps), (warp, lane, _) = row, ids + if warp < self.batch: + total = self.square_total(squares, warp, lane) + self.scale(words, gx, inv_rms(total, eps, m8=True), xs, ids) + + @cute.jit + def square_total(self, squares, warp, lane): + groups = cute.make_rmem_tensor(16, Float32) + for c in cutlass.range_constexpr(16): + groups[c] = Float32(0) + if lane < (16 if c % 4 == 3 else 32): + groups[c] = squares[warp * 448 + c // 4 * 112 + c % 4 * 32 + lane] + group_sum, _ = scatter_reduce(groups, lane, 16) + return sum((shuffle(group_sum, 2 * c) for c in range(16)), Float32(0)) + + @cute.jit + def scale(self, words, gx, inverse, xs, ids): + warp, lane, lat = ids + if lat: + self.scale_words(words, gx, inverse, xs, warp, lane) + + def scale_words(self, words, gx, inverse, xs, warp, lane): + """(x * inverse) * gamma per column, as packed FP32x2 multiplies.""" + pair = lambda w: (bf16(w), bf16(w, True)) + scaled = lambda k: mul2(mul2(pair(words[k]), (inverse, inverse)), pair(gx[k])) + packed = [pack_pair(*scaled(k)) for k in range(4)] + dst = xs + warp * 224 + 8 * lane + if cutlass.const_expr(self.umma): + dst = self.x_chunk(xs, warp, lane) + store_shared_words(dst, packed) + + def x_chunk(self, tile, row, chunk): + """16-byte chunk ``chunk`` (K columns 8 chunk ..) of row ``row`` in 128B- + swizzled [rows, 64] K tiles of ``self.n`` rows (``tile``: the 8-row block).""" + atom, part = chunk // 8, chunk % 8 + return tile + atom * (self.n * 128) + row * 128 + (part ^ row) * 16 + + @cute.jit + def umma_partials(self, sm, tmem, bar, ids): + """D [M, N] = W x^T on tcgen05 (warp 0 issues 4 K tiles), staged [token, row] + in the W tiles, then the 4-row partials pushed to their lines' owners.""" + q, warp, lane, tid = ids + m, n = self.m, self.n + done = bar + 2 + instruction: cutlass.Constexpr = 1 << 4 | 1 << 7 | 1 << 10 + instruction |= n >> 3 << 17 | m >> 4 << 24 + if warp == 0: + wait(bar + 3) # W tiles + fence_tmem() + for kb in cutlass.range_constexpr(4): + a = smem_desc(sm + kb * (m * 128), True) + b = smem_desc(sm + m * 512 + kb * (n * 128), True) + mma_bf16_block(tmem, a, b, instruction, int(kb > 0), done, int(kb == 3)) + wait_tmem(done) + # The MMAs (async proxy) read the W tiles the staging now overwrites. + cute.arch.fence_proxy("async.shared", space="cta") + # Warp w reads TMEM lanes 32 (w mod 4) .. (M = 64: the first 16 of each), + # half w / 4 of the columns. + quad, half, cols = warp % 4, warp // 4, n // 2 + values = () + for start, count in pow2_spans(cols): # TMEM loads: power-of-two columns + got = load_tmem(tmem + (32 * quad << 16) + half * cols + start, count) + values += got if count > 1 else (got,) + row = 32 * quad + lane if m == 128 else 16 * quad + lane + stage = view(sm, 0, Float32, n * m) + if lane < (32 if m == 128 else 16): + for t in cutlass.range_constexpr(cols): + stage[(half * cols + t) * m + row] = values[t] + fence_tmem(before=True) + cute.arch.sync_threads() + partials = view(sm, self.offsets[2], Float32, 1024 * self.slots).iterator + for p in cutlass.range(tid, self.batch * (m // 4), 256, unroll=1): + token, r = p // (m // 4), p % (m // 4) * 4 + line = self.lines * token + r // 64 + values4 = load_f32x4(stage.iterator + token * m + r, shared=True) + dst = partials + line // 16 * 1024 + q * 64 + r % 64 + push_float(dst, values4, bar + 1, line % 16) + + @cute.jit + def push(self, *args): + """WaveMoeTail's shared inputs: peer-store this rank's shared-expert rows [B, + 7168] into the RS mailboxes (-0 made +0); an ordinary launch after the down + GEMM.""" + launch(self.push_kernel, args, self.batch, 128, pdl=False) + + @cute.kernel + def push_kernel(self, args): + shared, peers, parity, rank, capacity = args + tid, _, _, token = thread() + for i in cutlass.range_constexpr(7): + column = 8 * (tid + 128 * i) + words = load_words(shared.iterator + token * 7168 + column, count=4) + offset = ((parity * self.tp + rank) * capacity + token) * self.shard + offset += column % self.shard + packed = tuple(no_neg_zero(w) for w in words) + store_relaxed(peers[column // self.shard] + Int64(offset) * 2, packed) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail_runtime.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail_runtime.py new file mode 100644 index 0000000000..a1aa106e27 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/consumer.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/consumer.py new file mode 100644 index 0000000000..ad5f2a4cf9 --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py new file mode 100644 index 0000000000..5f753fc868 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/producer.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/producer.py new file mode 100644 index 0000000000..e48397ce4f --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py new file mode 100644 index 0000000000..dccfb9a5cb --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py new file mode 100644 index 0000000000..e1563d38cb --- /dev/null +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py new file mode 100644 index 0000000000..8ca3e1e527 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py new file mode 100644 index 0000000000..63e96bfb79 --- /dev/null +++ b/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()