diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu index 8267f3f44f9f..796a81c5e0d5 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu @@ -30,6 +30,9 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { constexpr uint32_t copyBlockCtaSize = 128; constexpr uint32_t copyBlocknbBufs = 2; +// Loads per thread of one block-offset copy CTA: each CTA copies 1024 entries of a row, +// so a long row is read from host memory by many CTAs in parallel. +constexpr uint32_t copyBlockItersPerCta = 2; __device__ __host__ inline uint32_t divUp(uint32_t a, uint32_t b) { @@ -78,6 +81,7 @@ void copyBasePageRowsToDevice(ITensor const& input, ITensor& output, SizeType32 // dst_tensor[:, :num_seqs, 0] = src_tensor[:, copy_idx] // dst_tensor[:, :num_seqs, 1] = dst_tensor[:, :num_seqs, 0] + 1 +// Grid: (pool, sequence, chunk of copyBlockItersPerCta loads per thread along the row). __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict__ srcPtr, SizeType32* __restrict__ dstPtr, SizeType32 const srcMaxNumSequences, SizeType32 const dstMaxNumSequences, SizeType32 numBlocksPerSeq, SizeType32 const* __restrict__ copyIndex, SizeType32 const* __restrict__ indexScales, @@ -89,17 +93,21 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict __shared__ PackedInt data[copyBlocknbBufs][copyBlockCtaSize]; auto const iterPerSeq = divUp(numBlocksPerSeq * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize); + auto const iterBeg = blockIdx.z * copyBlockItersPerCta; + auto const numIters = min(copyBlockItersPerCta, iterPerSeq - iterBeg); auto const tid = threadIdx.x; auto const poolIdx = blockIdx.x; auto const seqIdx = blockIdx.y; auto const seqDimStride = kvFactor * numBlocksPerSeq; - uint32_t const srcIdxBeg = tid * elemPerAccess + (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride; - uint32_t const dstIdxKBeg = tid * elemPerAccess + (poolIdx * dstMaxNumSequences + seqIdx) * seqDimStride; + uint32_t const chunkBeg = iterBeg * copyBlockCtaSize * elemPerAccess; + uint32_t const srcIdxBeg + = chunkBeg + tid * elemPerAccess + (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride; + uint32_t const dstIdxKBeg = chunkBeg + tid * elemPerAccess + (poolIdx * dstMaxNumSequences + seqIdx) * seqDimStride; uint32_t const dstIdxVBeg = dstIdxKBeg + numBlocksPerSeq; uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksPerSeq; - for (uint32_t i = 0; i < iterPerSeq + copyBlocknbBufs; i++) + for (uint32_t i = 0; i < numIters + copyBlocknbBufs; i++) { uint32_t const idxBuf = i % copyBlocknbBufs; if (i >= copyBlocknbBufs) @@ -129,7 +137,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict PackedInt* const dst = &data[idxBuf][tid]; uint32_t const srcIdx = srcIdxBeg + copyBlockCtaSize * ldIter * elemPerAccess; PackedInt const* const src = reinterpret_cast(srcPtr + srcIdx); - if (srcIdx < srcIdxEnd) + if (ldIter < numIters && srcIdx < srcIdxEnd) { uint32_t const size = sizeof(PackedInt); asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"l"(__cvta_generic_to_shared(dst)), @@ -174,7 +182,16 @@ void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITenso "Not implemented case: numBlocksPerSeq * sizeof(SizeType32) = %zu must be a multiple of %zu.", static_cast(numBlocksPerSeq * sizeof(SizeType32)), static_cast(sizeof(PackedInt))); - dim3 gridDim(numPools, numSeqs, 1); + auto const iterPerSeq = divUp(static_cast(numBlocksPerSeq * sizeof(SizeType32)), + static_cast(sizeof(PackedInt) * copyBlockCtaSize)); + auto const numChunks = divUp(iterPerSeq, copyBlockItersPerCta); + if (numChunks == 0) + { + return; + } + TLLM_CHECK_WITH_INFO( + numChunks <= 65535, "numBlocksPerSeq = %d needs %u CTAs per row, more than 65535.", numBlocksPerSeq, numChunks); + dim3 gridDim(numPools, numSeqs, numChunks); dim3 blockDim(copyBlockCtaSize); copyBatchBlockOffsetsToDeviceKernel<<>>(srcPtr, dstPtr, srcMaxNumSequences, diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index e0b9e0ee76d3..3a0f49d9f165 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -1906,6 +1906,27 @@ void invokeMLAContextFp8Quantize(MlaParams& params, int total_kv_len, cudaStr params.quant_scale_o, params.dequant_scale_q, params.dequant_scale_kv, params.host_bmm1_scale); } } + else if (params.fp8_kv_input) + { + // K and V are FP8 already: quantize Q only, or nothing when it is FP8 too. Block 0 still + // writes the BMM scales. + constexpr int threads_per_block = 384; + constexpr int num_tokens_per_block = threads_per_block * 16 / 192 * sizeof(T); + int const q_len = params.fp8_q_input ? 0 : params.acc_q_len; + dim3 grid(std::max(1, int(tensorrt_llm::common::divUp(q_len, num_tokens_per_block))), 1, + params.fp8_q_input ? 1 : params.head_num); + + TLLM_LOG_DEBUG( + "Launching quantizeCopyInputToFp8Kernel (FP8 K/V input) with grid_size: (%d, %d, %d), " + "threads_per_block: %d, acc_q_len: %d, fp8_q_input: %d", + grid.x, grid.y, grid.z, threads_per_block, params.acc_q_len, params.fp8_q_input); + + quantizeCopyInputToFp8Kernel + <<>>(params.q_buf, static_cast<__nv_fp8_e4m3*>(params.quant_q_buf), + nullptr, nullptr, nullptr, nullptr, q_len, 0, srcVTokenStride, params.quant_scale_qkv, + params.bmm1_scale, params.bmm2_scale, params.quant_scale_o, params.dequant_scale_q, + params.dequant_scale_kv, params.host_bmm1_scale); + } else { // The Q or K tensor has layout of [num_tokens, head_num, 192] in the non-absorption mode. diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h index f480f0010a99..359599dd77b7 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -150,6 +150,14 @@ struct MlaParams // `quant_q_buf`. Context preprocessing must still rotate/cache K. bool q_rope_applied = false; + // Non-absorption context MLA: `k_buf` / `v_buf` point to contiguous FP8 E4M3 K + // [total_kv_len, head_num * (qk_nope_head_dim + qk_rope_head_dim)] and V + // [total_kv_len, head_num * v_head_dim], which the FMHA reads as they are. + bool fp8_kv_input = false; + // With `fp8_kv_input`: `quant_q_buf` already holds the FP8 Q, so the context + // quantization only writes the BMM scales. + bool fp8_q_input = false; + // DSv4 fused inverse-RoPE + FP8 quant epilogue parameters. Dsv4EpilogueFusionParams dsv4_epilogue_fusion; diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 295e82b847bb..acb693f3366e 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -2881,6 +2881,34 @@ void MLA::prepareInputs(FmhaParams& p, bool isGen) "Separate-K context FMHA requires contiguous K features within each token"); TORCH_CHECK(p.value_input->strides()[1] == 1, "Separate-V context FMHA requires contiguous V features within each token"); + auto const& k = p.key_input.value(); + auto const& v = p.value_input.value(); + TORCH_CHECK(k.scalar_type() == v.scalar_type(), "MLA context k and v must have the same dtype."); + // FP8 k / v (and an FP8 Q in quant_q_buffer, if given with quant_scale_qkv): the caller already + // quantized them for the FP8 context FMHA, which reads them instead of quantizing BF16 inputs. + if (k.scalar_type() == torch::kFloat8_e4m3fn) + { + TORCH_CHECK(mFP8ContextMLA, "FP8 MLA context k / v require the FP8 context MLA FMHA."); + TORCH_CHECK(!p.fwd.latent_cache.has_value(), + "FP8 MLA context k / v are for cached-KV context attention, without latent_cache."); + TORCH_CHECK(k.is_contiguous() && v.is_contiguous(), "FP8 MLA context k and v must be contiguous."); + TORCH_CHECK(k.size(1) == num_heads * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim) + && v.size(1) == num_heads * mMLAParams.v_head_dim, + "FP8 MLA context k / v must be [kv_tokens, num_heads * qk_head_dim] / [kv_tokens, num_heads * " + "v_head_dim]."); + if (p.fwd.quant_q_buffer.has_value() && p.fwd.quant_scale_qkv.has_value()) + { + auto const& quantQ = p.fwd.quant_q_buffer.value(); + TORCH_CHECK(quantQ.scalar_type() == torch::kFloat8_e4m3fn && quantQ.is_contiguous() + && quantQ.numel() >= p.getQueryInput().numel(), + "The FP8 MLA context quant_q_buffer must be a contiguous FP8 copy of q."); + } + } + else + { + TORCH_CHECK(k.scalar_type() == p.getQueryInput().scalar_type(), + "MLA context k / v must be FP8 E4M3 or have the dtype of q."); + } } } else @@ -3291,6 +3319,8 @@ int MLA::enqueueContext(FmhaParams const& p, cudaStream_t stream) { mla.k_buf = p.getK(); mla.v_buf = p.getV(); + mla.fp8_kv_input = p.key_input.has_value() && p.key_input->scalar_type() == torch::kFloat8_e4m3fn; + mla.fp8_q_input = mla.fp8_kv_input && p.fwd.quant_q_buffer.has_value() && p.fwd.quant_scale_qkv.has_value(); mla.helix_position_offsets = p.getHelixPositionOffsets(); mla.helix_is_inactive_rank = p.getHelixIsInactiveRank(); } @@ -3453,8 +3483,21 @@ int MLA::enqueueContext(FmhaParams const& p, cudaStream_t stream) "FP8 v buffer is required for fp8 context MLA in non-sparse mode"); fmhaParams.qPtr = reinterpret_cast(fp8QBuf); - fmhaParams.kPtr = useSparseMLA() ? nullptr : reinterpret_cast(workspaceViews.fp8KBuf); - fmhaParams.vPtr = useSparseMLA() ? nullptr : reinterpret_cast(workspaceViews.fp8VBuf); + if (useSparseMLA()) + { + fmhaParams.kPtr = nullptr; + fmhaParams.vPtr = nullptr; + } + else if (mla.fp8_kv_input) + { + fmhaParams.kPtr = reinterpret_cast(mla.k_buf); + fmhaParams.vPtr = reinterpret_cast(mla.v_buf); + } + else + { + fmhaParams.kPtr = reinterpret_cast(workspaceViews.fp8KBuf); + fmhaParams.vPtr = reinterpret_cast(workspaceViews.fp8VBuf); + } } else { diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 94213b82efb5..d1dfe5c3cdb5 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -180,23 +180,37 @@ target pages. The draft layers keep the drafter's own dtype, so an FP8 target KV cache (`kv_cache_config.dtype: fp8`) is paired with a BF16 draft cache. In disaggregated serving, direct worker launches require -`TRTLLM_DISAGG_SHARED_DRAFT_KV=1` on both the context and generation workers -to enable shared draft-KV transfer. The `disaggregated_mpi_worker` launcher -enables this setting automatically when both servers enable compatible -speculation and use the same `tensor_parallel_size` and -`enable_attention_dp`. +`TRTLLM_DISAGG_SHARED_DRAFT_KV=1` (or `auto`, which shares wherever the layout +allows it) on both the context and generation workers to enable shared +draft-KV transfer. The `disaggregated_mpi_worker` launcher enables this setting +automatically when its config lists both context and generation servers, both +enable compatible speculation, and they shard attention alike: the same +`tensor_parallel_size`, where `enable_attention_dp` counts as 1 because every +rank keeps all KV heads. A config that lists only one side (for example, a +context-only MPI world) leaves sharing off unless the variable is `1`. When sharing is enabled, the context server's draft pages are transferred to the generation server with the target pages, so the drafter starts generation with the prompt context instead of an empty one. This requires the draft layers to have the target's rank-local KV-head count; the draft dtype and head dimension may differ, except that an NVFP4 target KV cache requires an -identical draft layout. Both servers must use the same layout. +identical draft layout unless both caches are K-only (MLA). Both servers must +use the same layout. + +Hybrid recurrent/attention targets whose cache manager opts in (Kimi K3) also +hold a standalone DSpark drafter's layers, after their own. In every case the +drafter decides whether it shares: it declares whether its block attention reads +the manager's pages at the run's `tokens_per_block` +(`DFlashForCausalLM.reads_managed_ctx_pages`). By default only +`attention_backend="TRTLLM"` does, at the page sizes its trtllm-gen kernels +support; a drafter whose other backends read paged context declares that itself. The drafter keeps a separate draft KV cache manager instead when any of the -following applies: context parallelism larger than 1, a per-layer KV-head count -on the target or drafter, a specialized KV cache manager (for example, hybrid Mamba -models), Gemma 4 hybrid attention, or, in disaggregated serving, a draft layout +following applies: context parallelism larger than 1 (helix context parallelism +is shared only with a drafter that reads a helix-sharded context), a per-layer +KV-head count on the target or drafter, a specialized KV cache manager that does +not opt in (for example, most hybrid Mamba models), Gemma 4 hybrid attention, or, +in disaggregated serving, a draft layout that cannot be transferred with the target's. When the layers are shared, the worker log reports `target and draft KV layers share one KV cache manager` at startup. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/__init__.py new file mode 100644 index 000000000000..0b253306a8d7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 (NVFP4) / sm_107 / default. + +Imports nothing: ``routing.py`` selects the deployment's ``features`` profile before any module +of the target is imported. +""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/attn_res.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/attn_res.py new file mode 100644 index 000000000000..80ebe988703a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/attn_res.py @@ -0,0 +1,579 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 on Rubin (SM107): attention-residual mixing. + +Before each sublayer norm a token mixes a bank of earlier sublayer outputs (the HF +reference's ``_apply_attn_res``). The functions here pick the kernel for each mix: the +per-token and persistent fused kernels with the RMSNorm fused in, the PDL kernels that also +finish a pending all-reduce (a Lamport push, the row-parallel MoE tail, a deferred MoE +residual), and the FP32 reference outside the kernels' limits. ``_DeferredMoEResidual`` and +``_DeferredRpTail`` carry a MoE layer's residual update to the next mix. +""" + +from __future__ import annotations + +from typing import Any, NamedTuple, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.norm.attn_res_fwd import ( + attn_res_fwd, +) +from tensorrt_llm._torch.cute_dsl_utils import sm_reserve +from tensorrt_llm.logger import logger + +from .features import active as _features +from .layers import LamportPending, RMSNormWeight, rms_norm +from .ops.comm.k3_attn_res_lamport import k3_attn_res_lamport +from .ops.moe.k3_moe_residual_add import k3_moe_residual_add as k3_moe_residual_add_op +from .ops.norm.k3_attn_res_norm import k3_attn_res_norm + +_F = _features() + + +# --------------------------------------------------------------------------- +# attn_res: per-token snapshot mixing (HF `_apply_attn_res`). +# --------------------------------------------------------------------------- + +# Fused attention-residual kernels (every Kimi K3 recipe runs the "split" topology): +# * up to _FUSED_ATTN_RES_MAX_TOKENS tokens the per-token C++ kernels (one CTA for N <= 4, one 8-CTA +# cluster per token for N >= 5): trtllm::attn_res_rmsnorm_fwd / attn_res_add_rmsnorm_fwd; +# * above it the persistent kernel (one CTA per SM looping over tokens), +# trtllm::attn_res_add_rmsnorm_persistent_fwd, where it implements the shape (H == 7168 and +# 2 <= N <= 9); +# * elsewhere trtllm::attn_res_fwd and the RMSNorm, or the FP32 reference past that op's contract. +# The K3 PDL kernels (TRTLLM_K3_ATTN_RES_PDL) take precedence up to their own token limit. +_FUSED_ATTN_RES_MAX_TOKENS = 32 + + +def _use_persistent_attn_res(M: int, H: int, N: int) -> bool: + """Whether the persistent kernel serves this mix: H == 7168, 2 <= N <= 9 and more than + ``_FUSED_ATTN_RES_MAX_TOKENS`` tokens (its grid is sized by the SM count, not the token + count).""" + return H == 7168 and 2 <= N <= 9 and M > _FUSED_ATTN_RES_MAX_TOKENS + + +# Token limit of trtllm::attn_res_fwd (cpp/tensorrt_llm/thop/attnResOp.cpp). +_ATTN_RES_OP_MAX_TOKENS = 16384 + + +def _apply_attn_res_fused( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: RMSNormWeight +) -> Optional[torch.Tensor]: + """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. + + Returns ``None`` when the call falls outside the fused kernel's + contract (dtype/shape/arch) so the caller can use the exact fp32 reference + instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` + layout. Candidate order matches the reference: snapshots first, the + running prefix sum last. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + if K + 1 > 12 or not (4096 <= H <= 8192 and H % 1024 == 0): + return None + attn_res_op = attn_res_fwd + proj_w = proj.weight.reshape(-1).to(torch.bfloat16).contiguous() + norm_w = norm.weight.to(torch.bfloat16).contiguous() + + def run(prefix, block): + t = prefix.shape[0] + output, _rsigma, _probs, _logits = attn_res_op( + prefix.reshape(t, 1, H).contiguous(), + block.reshape(K, t, 1, H).contiguous(), + proj_w, + norm_w, + float(norm.eps), + ) + return output.reshape(t, H) + + if M <= _ATTN_RES_OP_MAX_TOKENS: + return run(prefix_sum, block_residual) + # The op takes at most _ATTN_RES_OP_MAX_TOKENS tokens per call and is + # per-token, so longer prefill chunks run in slices; the fp32 reference + # below would otherwise cost seconds per iteration at 24k+ tokens. + out = torch.empty_like(prefix_sum) + for lo in range(0, M, _ATTN_RES_OP_MAX_TOKENS): + hi = min(M, lo + _ATTN_RES_OP_MAX_TOKENS) + out[lo:hi] = run(prefix_sum[lo:hi], block_residual[:, lo:hi]) + return out + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether a fused kernel served the mix, once per shape (a shape gate that rejects + the call is otherwise silent). Emitted at debug level: it is once per distinct shape, not once + per process, so it is a diagnostic rather than a summary. + """ + logger.debug_once( + f"Kimi K3 attn-res fusion [{site}]: " + f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", + key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", + ) + + +# TRTLLM_K3_ATTN_RES_PDL: up to TRTLLM_K3_ATTN_RES_PDL_MAX_M tokens, the fused attn_res(+add) +# +RMSNorm runs as the k3_attn_res_norm entry (kernels/k3_attn_res.py): +# a CTA cluster per token that loads the weights and snapshot rows and reduces their statistics +# before griddepcontrol.wait, i.e. under the kernel in front of it. It takes precedence over the +# persistent topology at those token counts. The updated prefix sum is bit-identical; the +# normalized output can differ from the C++ per-token kernels by one BF16 rounding step in a small +# fraction of the elements. +# It also leaves the MoE residual update residual + bf16(routed + shared) of the BF16-add path +# (kimi_k3_moe_residual) to the next layer's pre-attention mix, which computes it bit-identically +# (k3_attn_res_moe_add_rmsnorm) behind the latent-up GEMM. +_K3_ATTN_RES_PDL = _F.attn_res_pdl +_K3_ATTN_RES_PDL_MAX_M = _F.attn_res_pdl_max_m + +# Where the mix after attention (A_mid) adds the attention output to the running prefix sum, the +# attention block pushes its o_proj output into the o_proj all-reduce's MNNVL Lamport workspace +# (layers) and leaves the call unfinished; A_mid finishes it (k3_attn_res ``lamport_finish``: +# one-shot transport up to the C++ one-shot size, two-shot above, where each row's owner reduces +# and multicasts it) and advances the workspace: no finish kernel or C++ all-reduce runs in +# between. Where o_proj's GEMM has no push epilogue (the MXFP8 GEMM on a prequantized KDA output) +# a copy kernel pushes its output. The sums have the order and rounding of the finish kernel, so +# A_mid produces the same bits. +# TRTLLM_K3_ULL_OPROJ_SANDWICH=1: A_mid finishes the pushed o_proj all-reduce before its +# griddepcontrol.wait (k3_attn_res ``early_poll``): it becomes resident when +# the o_proj push GEMM passes its own wait, polls every rank's slots while o_proj runs, and only +# its stores wait for the o_proj grid, so the o_proj -> A_mid boundary leaves the critical path +# (the post-attention collective sandwich as two PDL grids). Same bits. +_K3_ULL_OPROJ_SANDWICH = _F.ull_oproj_sandwich +# TRTLLM_K3_ULL_TAIL_SANDWICH=1: the next layer's pre-attention mix +# finishes the row-parallel MoE tail's pushed all-reduce the same way (the pre-attention sandwich: +# tail GEMM + push, then the early-polling mix). Same bits. +_K3_ULL_TAIL_SANDWICH = _F.ull_tail_sandwich + + +class _DeferredMoEResidual(NamedTuple): + """A MoE layer's ``residual + bf16(routed + shared)`` that the next layer fuses or materializes.""" + + residual: torch.Tensor + routed: torch.Tensor + shared: torch.Tensor + + def materialize(self) -> torch.Tensor: + return k3_moe_residual_add_op(self.routed, self.shared, self.residual) + + +class _DeferredRpTail(NamedTuple): + """A MoE layer's ``bf16(residual + all-reduce of the ranks' row-parallel tail partials)``: + ``result`` is the pushed all-reduce, left for the next layer's pre-attention mix to finish + (``LamportPending``).""" + + residual: torch.Tensor + result: Any + + def materialize(self) -> torch.Tensor: + # A BF16 add rounds the FP32 sum to nearest, as the fused mix does. + return self.residual + self.result.materialize().view(self.residual.shape) + + +def _defer_moe_residual(num_tokens: int) -> bool: + """Whether a MoE layer leaves its BF16 residual update to the next layer.""" + return _K3_ATTN_RES_PDL and num_tokens <= _K3_ATTN_RES_PDL_MAX_M + + +def _materialize_prefix_sum(hidden_states): + if isinstance(hidden_states, (_DeferredMoEResidual, _DeferredRpTail)): + return hidden_states.materialize() + return hidden_states + + +def _k3_attn_res_pdl( + prefix_sum: torch.Tensor, + addend: Optional[torch.Tensor], + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, +): + """``k3_attn_res_norm`` (with ``addend`` as the one addend) when enabled and applicable, + else ``None``.""" + if not _K3_ATTN_RES_PDL or prefix_sum.shape[0] > _K3_ATTN_RES_PDL_MAX_M: + return None + from .kernels import k3_attn_res + + weights = (proj.weight.reshape(-1), norm.weight, output_norm.weight) + addends = () if addend is None else (addend,) + if not k3_attn_res.can_use_k3_attn_res(prefix_sum, addends, block_residual, *weights): + return None + eps = float(norm.eps) + out_eps = float(output_norm.eps) + return k3_attn_res_norm(prefix_sum, addends, block_residual, *weights, eps, out_eps) + + +def _k3_attn_res_lamport_applies( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, +) -> bool: + """Whether A_mid on ``prefix_sum`` can read its addend from a Lamport workspace: the + vectorized PDL attn_res kernel serves the ADD_PREFIX mix.""" + if not _K3_ATTN_RES_PDL or prefix_sum.shape[0] > _K3_ATTN_RES_PDL_MAX_M: + return False + from .kernels import k3_attn_res + + weights = (proj.weight.reshape(-1), norm.weight, output_norm.weight) + return k3_attn_res.can_use_k3_attn_res(prefix_sum, (prefix_sum,), block_residual, *weights) + + +def _k3_attn_res_pdl_lamport( + prefix_sum: torch.Tensor, + addend: LamportPending, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, + entry_trigger: bool = False, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """``(prefix_sum + all-reduce result, normed mix)``: the mix finishes the pushed all-reduce and + reads the result from the Lamport workspace (``k3_attn_res_lamport``), or + ``None`` if that does not apply. ``entry_trigger``: release the dependent at kernel entry (the + caller checked that the next launch is the ULL front's head).""" + if tuple(addend.shape) != tuple(prefix_sum.shape) or not _k3_attn_res_lamport_applies( + prefix_sum, block_residual, proj, norm, output_norm + ): + return None + M, H = prefix_sum.shape + _note_attn_res_fusion( + "add+attn_res+norm/pdl/lamport_finish" + ("/entry" if entry_trigger else ""), + True, + M, + H, + int(block_residual.shape[0]) + 1, + ) + ws = addend.workspace + return k3_attn_res_lamport( + prefix_sum, + ws.uc_local, + ws.flags.data_ptr(), + block_residual, + proj.weight.reshape(-1), + norm.weight, + output_norm.weight, + float(norm.eps), + float(output_norm.eps), + *addend.finish_args, + False, + _K3_ULL_OPROJ_SANDWICH, + entry_trigger, + ) + + +def _k3_attn_res_pdl_rp_tail( + deferred: _DeferredRpTail, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, + mxfp8: bool = False, + entry_trigger: bool = False, + prefix_out: Optional[torch.Tensor] = None, +) -> Optional[Tuple[torch.Tensor, ...]]: + """``(prefix_sum, normed mix)`` of a deferred row-parallel MoE tail: the mix finishes the + pushed all-reduce, reads it from the Lamport workspace and adds it to the residual with a + round-to-nearest BF16 add; ``None`` if the + vectorized PDL kernel does not serve the shape. With ``mxfp8``, ``(prefix_sum, normed mix, xq, + xs)``. ``entry_trigger`` (with ``mxfp8``): the kernel releases its dependent at entry (a KDA + QKVG projection that stages its weights, as in ``_k3_attn_res_pdl_moe_add``). ``prefix_out`` + (``TRTLLM_K3_ULL_HYGIENE``, a contiguous buffer of the residual's shape outside + ``block_residual``): the kernel stores ``prefix_sum`` there (the same bits).""" + residual, result = deferred + if tuple(result.shape) != tuple(residual.shape): + return None + from .kernels import k3_attn_res + + weights = (proj.weight.reshape(-1), norm.weight, output_norm.weight) + if not k3_attn_res.can_use_k3_attn_res(residual, (residual,), block_residual, *weights): + return None + M, H = residual.shape + _note_attn_res_fusion( + "rp_tail+attn_res+norm/pdl/lamport_finish", + True, + M, + H, + int(block_residual.shape[0]) + 1, + ) + ws = result.workspace + # TRTLLM_K3_ULL2_COLL_CP_TAIL: the call carries the column-parallel tail's all-gather rows. + finish_ag = result.finish_ag + return k3_attn_res_lamport( + residual, + ws.uc_local, + ws.flags.data_ptr(), + block_residual, + *weights, + float(norm.eps), + float(output_norm.eps), + *result.finish_args, + True, + _K3_ULL_TAIL_SANDWICH or finish_ag, + entry_trigger and mxfp8, + finish_ag, + mxfp8=mxfp8, + updated=prefix_out, + ) + + +def _k3_attn_res_pdl_moe_add( + deferred: _DeferredMoEResidual, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, + mxfp8: bool = False, + entry_trigger: bool = False, + mx_atom: bool = False, +) -> Optional[Tuple[torch.Tensor, ...]]: + """``(prefix_sum, normed mix)`` of a deferred MoE residual update, or ``None``; with + ``mxfp8``, ``(prefix_sum, normed mix, xq, xs)`` (the MXFP8 form of the mix). + ``entry_trigger`` (with ``mxfp8``): the kernel releases its dependent at entry + (``k3_attn_res_norm``), for a KDA QKVG projection that stages its weights. + ``mx_atom`` (with ``mxfp8``): ``xs`` in the mid-M GEMM's layout.""" + residual, routed, shared = deferred + if residual.shape[0] > _K3_ATTN_RES_PDL_MAX_M: + return None + from .kernels import k3_attn_res + + weights = (proj.weight.reshape(-1), norm.weight, output_norm.weight) + if not k3_attn_res.can_use_k3_attn_res(residual, (routed, shared), block_residual, *weights): + return None + M, H = residual.shape + _note_attn_res_fusion("moe_add+attn_res+norm/pdl", True, M, H, int(block_residual.shape[0]) + 1) + if mxfp8: + return k3_attn_res_norm( + residual, + (routed, shared), + block_residual, + *weights, + float(norm.eps), + float(output_norm.eps), + mxfp8=True, + entry_trigger=entry_trigger, + mx_atom=mx_atom, + ) + return k3_attn_res_norm( + residual, + (routed, shared), + block_residual, + *weights, + float(norm.eps), + float(output_norm.eps), + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, +) -> Optional[torch.Tensor]: + """Fuse attention-residual mixing with its immediately following norm.""" + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # The fused path is taken for M <= _FUSED_ATTN_RES_MAX_TOKENS, H == 7168 and + # N <= 12, which is the measured window; larger token counts have not been + # measured and fall back to the unfused add + attn_res_fwd + RMSNorm path. + pdl_out = _k3_attn_res_pdl(prefix_sum, None, block_residual, proj, norm, output_norm) + if pdl_out is not None: + _note_attn_res_fusion("attn_res+norm/pdl", True, M, H, N) + return pdl_out + if _use_persistent_attn_res(M, H, N): + _, output = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd( + prefix_sum.reshape(M, 1, H).contiguous(), + None, + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + float(output_norm.eps), + sm_reserve(), + ) + _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) + return output.reshape(M, H) + + if M > _FUSED_ATTN_RES_MAX_TOKENS or H != 7168 or N > 12: + _note_attn_res_fusion("attn_res+norm", False, M, H, N) + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = torch.ops.trtllm.attn_res_rmsnorm_fwd( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + float(output_norm.eps), + ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) + return output.reshape(M, H) + + +def _apply_attn_res_add_rmsnorm_fused( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm. + + The production residual add produces a BF16 tensor that remains live across + the following MLP. The kernel therefore returns that materialized, + BF16-rounded prefix sum alongside the normalized attention-residual output, + while avoiding a separate add launch and a re-read of the intermediate by + attention-residual selection. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or addend.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not addend.is_cuda + or not block_residual.is_cuda + or prefix_sum.shape != addend.shape + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # Same measured window as _apply_attn_res_rmsnorm_fused above. + pdl_out = _k3_attn_res_pdl(prefix_sum, addend, block_residual, proj, norm, output_norm) + if pdl_out is not None: + _note_attn_res_fusion("add+attn_res+norm/pdl", True, M, H, N) + return pdl_out + if _use_persistent_attn_res(M, H, N): + updated_prefix_sum, output = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd( + prefix_sum.reshape(M, 1, H).contiguous(), + addend.reshape(M, 1, H).contiguous(), + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + float(output_norm.eps), + sm_reserve(), + ) + _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + if M > _FUSED_ATTN_RES_MAX_TOKENS or H != 7168 or N > 12: + _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + addend_kernel = addend.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + updated_prefix_sum, output = torch.ops.trtllm.attn_res_add_rmsnorm_fwd( + layer_kernel, + addend_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + float(output_norm.eps), + ) + _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + +def _apply_attn_res( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: RMSNormWeight +) -> torch.Tensor: + """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). + + prefix_sum: ``[num_tokens, hidden_size]`` + block_residual: ``[num_snapshots, num_tokens, hidden_size]`` + + Inputs fitting the fused kernel's contract dispatch directly to the in-tree + ``trtllm::attn_res_fwd`` op. Only the fallback boundary restores the HF + ``[M, K, H]`` layout. + """ + fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) + if fused is not None: + return fused + # The fp32 reference below is orders of magnitude slower at prefill + # sizes; say so once per shape rather than fall back silently. + _note_attn_res_fusion( + "attn_res", + False, + prefix_sum.shape[0], + prefix_sum.shape[1], + int(block_residual.shape[0]) + 1, + ) + block_residual_hf = block_residual.transpose(0, 1) + v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.eps) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +def _apply_attn_res_and_rmsnorm( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm.""" + fused = _apply_attn_res_rmsnorm_fused(prefix_sum, block_residual, proj, norm, output_norm) + if fused is not None: + return fused + mixed = _apply_attn_res(prefix_sum, block_residual, proj, norm) + return rms_norm(mixed, output_norm.weight, output_norm.eps) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: RMSNormWeight, + output_norm: RMSNormWeight, + entry_trigger: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm. + + ``addend`` may be a ``LamportPending``: the attention block's pushed all-reduce, which the PDL + attn_res kernel finishes where it serves the mix and the finish kernel turns into a tensor + otherwise. ``entry_trigger``: see :func:`_k3_attn_res_pdl_lamport`.""" + if isinstance(addend, LamportPending): + fused = _k3_attn_res_pdl_lamport( + prefix_sum, addend, block_residual, proj, norm, output_norm, entry_trigger + ) + if fused is not None: + return fused + addend = addend.materialize() + fused = _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, proj, norm, output_norm + ) + if fused is not None: + return fused + updated_prefix_sum = prefix_sum + addend + return updated_prefix_sum, _apply_attn_res_and_rmsnorm( + updated_prefix_sum, block_residual, proj, norm, output_norm + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/dflash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/dflash.py new file mode 100644 index 000000000000..049a0445cb73 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/dflash.py @@ -0,0 +1,4284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DFlash speculative-decoding worker, Kimi K3 Rubin private copy. + +The 3-way merge of ``tensorrt_llm/_torch/speculative/dflash.py`` (main) with Jonas +Li's ``user/jonasl/k3-rubin-agentic-opt`` (189fe67b81), worker side only. +``DFlashSpecMetadata`` is main's (the engine creates it, and its ``prepare`` +assigns the slots of generation requests); Jonas's input-commit staging is not +ported. The drafter's context KV lives in the target KV cache manager's +appended draft layers when the engine unifies them (``draft_layer_ids``, +reached through ``get_draft_kv_cache_view``): the worker decodes their pages +from ``kv_cache_block_offsets`` in its own forward, and a generation request +that arrives without a prefill here (disaggregated serving) starts its context +at its first position. + +Jonas's selection: target and greedy draft tokens come from the private +``K3VocabParallelHead`` (MNNVL pair exchange, one-launch Markov chains, +temperature-only sampling) wherever the target LM head is a plain vocab shard; +main's head is not used. His switches, read where he read them; unset, the +rest of the worker is main's: + +* ``TLLM_SPEC_TAIL_FUSION=1`` (+ ``TLLM_DSPARK_TAIL_V2``): fused strict + acceptance, fused selection + acceptance, the parallel block-table decode, the + context latents on a side stream and ``_prepare_gen_drafter_inputs_fused`` in + place of main's ``_prepare_paged_latent_gen_inputs``. +* ``TLLM_DSPARK_DRAFT_GEMM``, ``TRTLLM_K3_DRAFT3_GLUE``: drafter LM-head / glue + variants inside those paths; the drafter embedding there runs as the pushed + MNNVL lookup where it applies. + +Optional attention dependencies used by DFlash's private context cache. + +Per-step glue of the MLA DSpark drafter off its critical path (``draft3_glue``). + +Same values: the split-KV plan of the block decode on a side stream (joined before the first +decode), the shape-only scratch indices built once per shape, no copy of the first layer's +residual, no unused dtype conversion of the slots, the drafter LM head reading a contiguous row +range of the block output in place of a gather. + +Pinned GEMM kernels for the DSpark MLA drafter's BF16 projections. + +The drafter's projections run at 7 to 32 token rows. The generic low-M BF16 +dispatcher picks a kernel per shape from an isolated timing race, which does +not see what the decode graph sees: cold weights (the drafter streams ~80 MB +per layer) and a PDL head start from the predecessor. Under those conditions +(measured on GB300 as predecessor RMSNorm + GEMM in a CUDA graph with PDL, +weights rotated beyond L2) the low-M split-K tensor-core kernel at the tactics +below beats the dispatcher's in-graph choices -- the SIMT direct kernel for +q_b runs 6.8 us in the decode graph against 2.2 us here, and the fused +[q_a | kv_a] projection, created after the dispatcher labelled the modules, +falls back to cuBLAS (8-10 us against 6.6 us). + +``TLLM_DSPARK_DRAFT_GEMM=1`` selects them. They change the drafter's GEMM +accumulation order (split-K partial sums), i.e. draft-token rounding only; the +target's verification keeps the output distribution unchanged. + +All-reduce sites of the DSpark MLA drafter on pushed MNNVL all-reduces. + +Where the drafter's all-reduce runs fused with the residual add and RMSNorm over MNNVL: + +* o_proj and down_proj run as the low-M BF16 split-K GEMM whose epilogue stores the partial + straight into the Lamport slots of the all-reduce (``dspark_draft_gemm_push``), and one site + kernel finishes the all-reduce, adds the residual and applies the following RMSNorm + (``dspark_draft_ar.DraftArNormKernel``). The site kernel is launched behind the GEMM with PDL and + polls while the GEMM runs; its outputs are bitwise the MNNVL fused all-reduce + (``RESIDUAL_RMS_NORM``) of the GEMM output. The GEMM is the split-K kernel at the tactics below + instead of the module's GEMM, i.e. a different FP32 accumulation order of the partial (drafter + rounding only). +* the embedding of the draft block: the owner rank of each token multicasts its table row and every + rank polls it (``DraftEmbedKernel``), bitwise the vocab-sharded lookup + all-reduce. + +A site that does not apply (more than 64 rows, no single-load launch shape, no MNNVL workspace view +during capture) returns None and the caller keeps its path. The Lamport slots come from the +package's pushed all-reduce (``layers``) on the drafter's MNNVL workspace. + +Vocab-parallel token selection for a vocab-sharded LM head. + +A column-parallel LM head leaves every TP rank with a contiguous vocab shard of +the logits. Gathering the full logits to pick one token per row moves +``rows * vocab`` values; this module instead reduces each rank's shard to one +``(value, global index)`` pair per row, exchanges only those pairs, and takes +the winner. The result equals ``torch.argmax`` over the gathered logits bit for +bit: ``NaN`` beats every number, and equal maxima resolve to the lowest global +index (rank shards are contiguous and ordered, so the lowest rank's lowest +column wins). + +Three consumers: + +* :func:`vocab_parallel_argmax` -- greedy target sampling. +* :func:`vocab_parallel_sample` -- temperature-only sampling through the + Gumbel-max trick: ``argmax(logits / T + G)`` with ``G`` standard Gumbel noise + is an exact draw from ``softmax(logits / T)``. The noise of every vocab entry + comes from a counter-based Philox stream keyed by the row's seed, so all + ranks agree on the winner. The distribution matches the gathered sampler; the + individual draws do not (different random stream). +* :func:`vocab_parallel_markov_tokens` -- the DSpark vanilla-Markov greedy + block chain on shard logits: step ``i`` adds ``markov_w2 @ markov_w1[prev_i]`` + to block position ``i`` and takes the global argmax, which becomes + ``prev_{i+1}`` and draft token ``i``. + +The pair exchange is an all-reduce (sum) of a zero-padded ``[rows, 2 * tp]`` +float32 buffer in which each rank fills only its own two columns (``x + 0`` is +exact), over the MNNVL all-reduce when available and NCCL otherwise. Every +launch has a shape fixed by ``(rows, shard width)`` and nothing synchronizes +with the host, so all paths are CUDA-graph capturable. + +Non-greedy target rows are sampled vocab-parallel only when the deployment's +``advanced_sampling_mode`` is ``no_topk_no_topp``. The draft Markov chain of up +to 16 requests and 16 block positions over the MNNVL (or single-rank) exchange +runs as one launch: the CuTe DSL kernel +(:meth:`K3VocabParallelHead.markov_tokens_v3`) on SM100-family devices where it +applies, else :meth:`K3VocabParallelHead.markov_tokens_chain`; any other chain +takes :meth:`K3VocabParallelHead.markov_tokens`. + +An MNNVL all-reduce with its own Lamport workspace. + +The MNNVL all-reduces of one mapping (``layers.TpAllReduce``) share a single +workspace (buffers plus the flags that sequence them), so two of them must not be in +flight at once -- they would interleave in the same Lamport buffers. A reduction that +runs on a side stream concurrently with the main stream's all-reduces needs its own +workspace; this class allocates one the same way +``get_or_scale_allreduce_mnnvl_workspace`` does and calls the same fused kernel, so its +results are those of the shared-workspace all-reduce for the same inputs. +""" + +import functools +import math +import os +import weakref +from collections import deque +from dataclasses import dataclass +from functools import lru_cache +from typing import TYPE_CHECKING, Callable, Optional, Tuple + +import torch +import torch.nn.functional as F +import triton +from torch import nn + +from tensorrt_llm._device_copy import DeviceCopyBuffer +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAReplayState, + KimiK3HybridCacheManager, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID +from tensorrt_llm._torch.pyexecutor.resource_manager import BaseResourceManager +from tensorrt_llm._torch.speculative import DFlashSpecMetadata +from tensorrt_llm._torch.speculative.accept_stats import maybe_create_recorder +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + SpecWorkerBase, +) +from tensorrt_llm._utils import mpi_comm, prefer_pinned +from tensorrt_llm.functional import AllReduceFusionOp, AllReduceParams, AllReduceStrategy +from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +from .features import active as _features +from .kernels.spec.spec_tail_kernels import ( + BLOCK_TABLE_CHUNK, + TAIL_V2_MAX_GENS, + DraftBlockPlacement, + spec_tail_fusion_enabled, + tail_v2_enabled, + vocab_shard_embedding_applies, +) +from .kernels.spec.vocab_parallel_kernels import ( + _CHAIN_MAX_REQS, + _CHAIN_MAX_STEPS, + _LOCAL_BLOCK_V, + _MARKOV_BLOCK_V, + _PairLayout, +) +from .layers import TpAllReduce, mnnvl_serves, mnnvl_workspace_bytes +from .ops.spec.k3_dflash_ctx_positions import k3_dflash_ctx_positions +from .ops.spec.k3_dflash_gen_prepare import k3_dflash_gen_prepare +from .ops.spec.k3_dflash_noise_embedding import k3_dflash_noise_embedding +from .ops.spec.k3_dflash_paged_latent_prepare import k3_dflash_paged_latent_prepare +from .ops.spec.k3_draft_block_table import k3_draft_block_table +from .ops.spec.k3_fused_select_accept import k3_fused_select_accept +from .ops.spec.k3_fused_strict_accept import k3_fused_strict_accept +from .ops.spec.k3_markov_add_local_argmax import k3_markov_add_local_argmax +from .ops.spec.k3_markov_chain import k3_markov_chain +from .ops.spec.k3_markov_reduce_embed import k3_markov_reduce_embed +from .ops.spec.k3_next_new_tokens import k3_next_new_tokens +from .ops.spec.k3_store_context_latents import k3_store_context_latents +from .ops.spec.k3_vocab_local_select import k3_vocab_local_select +from .ops.spec.k3_vocab_reduce_pairs import k3_vocab_reduce_pairs +from .ops.spec.k3_vocab_shard_embedding import k3_vocab_shard_embedding + +_NUM_LAMPORT_BUFFERS = 3 + + +class PrivateMNNVLAllReduce: + """The MNNVL all-reduce over a private workspace (for a side stream). + + Construction is collective over the mapping's TP group. + """ + + def __init__(self, mapping: Mapping, dtype: torch.dtype, max_tokens: int, hidden: int): + from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer + + self.mapping = mapping + self.dtype = dtype + elem = torch.tensor([], dtype=dtype).element_size() + need = mnnvl_workspace_bytes(max_tokens, hidden, mapping.tp_size, dtype) + self.buffer_size_bytes = math.ceil(max(need, 1) / (8 * 1024 * 1024)) * (8 * 1024 * 1024) + self.max_tokens = max_tokens + self.hidden = hidden + self._elem = elem + workspace_bytes = _NUM_LAMPORT_BUFFERS * self.buffer_size_bytes + comm = mpi_comm().Split( + int(mapping.pp_rank * mapping.cp_size + mapping.cp_rank), mapping.tp_rank + ) + force_mn = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" + self._handle = McastGPUBuffer( + workspace_bytes, + mapping.tp_size, + mapping.tp_rank, + mapping.local_rank, + force_mn or mapping.is_multi_node(), + comm.py2f(), + ) + buffer = self._handle.get_uc_buffer( + mapping.tp_rank, (workspace_bytes // torch.float32.itemsize,), torch.float32, 0 + ) + buffer.fill_(-0.0) + torch.cuda.synchronize() + comm.Barrier() + self._comm = comm + self._buffer_base = buffer.view(dtype).view(_NUM_LAMPORT_BUFFERS, -1) + # [cur idx, dirty idx, bytes per buffer, dirty num stages, numBytesToClear[4], access count] + self._buffer_flags = torch.tensor( + [0, 2, self.buffer_size_bytes, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=torch.device("cuda", mapping.local_rank), + ) + + def fits(self, num_tokens: int, hidden: int) -> bool: + need = mnnvl_workspace_bytes(num_tokens, hidden, self.mapping.tp_size, self.dtype) + return need <= self.buffer_size_bytes + + def rms_norm(self, partial: torch.Tensor, gamma: torch.Tensor, eps: float) -> torch.Tensor: + """``RMSNorm(all_reduce(partial))``, as the MNNVL all-reduce with ``RMS_NORM``.""" + assert partial.dtype == self.dtype and self.fits( + partial.numel() // partial.shape[-1], partial.shape[-1] + ) + outputs = torch.ops.trtllm.mnnvl_fusion_allreduce( + partial, + gamma, + None, + eps, + self._buffer_base, + self._buffer_flags, + True, + None, + int(AllReduceFusionOp.RMS_NORM), + ) + return outputs[0] + + +# Programs the local reductions aim to launch (rows x column splits). +_TARGET_PROGRAMS = 256 + + +class _PairExchange: + """Exchanges ``(value, global index)`` float32 pairs per row across TP. + + Over MNNVL each rank owns ``splits`` consecutive pair slots of a zero-padded + ``[rows, width]`` buffer that an all-reduce sums; over NCCL each rank + publishes one pair per row and the pairs are all-gathered. + """ + + def __init__(self, mapping: Mapping): + self.mapping = mapping + self.tp_size = mapping.tp_size + self.tp_rank = mapping.tp_rank + self._mnnvl = None + if self.tp_size > 1 and mnnvl_serves(mapping, torch.float32, AllReduceStrategy.AUTO): + try: + self._mnnvl = TpAllReduce(mapping, torch.float32) + except (RuntimeError, ValueError) as e: + logger.info(f"Vocab-parallel pair exchange: MNNVL unavailable ({e}); using NCCL.") + self.backend = "mnnvl" if self._mnnvl is not None else "nccl" + + def chain_buffers(self, device): + """``(multicast view, local view, set floats)`` of the one-launch Markov + chain's exchange: three sets of ``[steps, requests, ranks, 2]`` float32 + pairs, -0.0 where nothing was written. Allocated collectively on first + use (every rank reaches it at the same point of the warmup), outside + graph capture. A single rank uses one local buffer for both views.""" + if getattr(self, "_chain", None) is not None: + return self._chain + assert not torch.cuda.is_current_stream_capturing() + set_floats = _CHAIN_MAX_STEPS * _CHAIN_MAX_REQS * max(self.tp_size, 1) * 2 + numel = 3 * set_floats + if self.tp_size == 1: + buf = torch.full((numel,), -0.0, dtype=torch.float32, device=device) + self._chain = (buf, buf, set_floats) + return self._chain + from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer + + assert self.backend == "mnnvl" + # The communicator of the exchange's MNNVL workspace (allocated with it). + comm = get_or_scale_allreduce_mnnvl_workspace(self.mapping, torch.float32)["comm"] + force_mn = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" + handle = McastGPUBuffer( + numel * 4, + self.tp_size, + self.tp_rank, + self.mapping.local_rank, + force_mn or self.mapping.is_multi_node(), + comm.py2f(), + ) + local = handle.get_uc_buffer(self.tp_rank, [numel], torch.float32, 0) + mcast = handle.get_mc_buffer([numel], torch.float32, 0) + local.fill_(-0.0) + torch.cuda.synchronize() + comm.Barrier() + self._chain_handle = handle + self._chain = (mcast, local, set_floats) + return self._chain + + def markov_v3_buffers(self, device, entries: int): + """``(local address, multicast address, state)`` of the CuTe DSL Markov + chain's key exchange for ``entries`` (ranks x CTAs per rank): its + rotating sets, zero-filled, and this rank's int32 launch state. Allocated + collectively on first use, outside graph capture. A single rank uses one + local buffer for both addresses.""" + cached = getattr(self, "_markov_v3", None) + if cached is not None and cached[0] == entries: + return cached[1:4] + assert cached is None, "the Markov exchange is sized once per process" + assert not torch.cuda.is_current_stream_capturing() + from .kernels.dspark_markov_greedy import NUM_SETS, STATE_WORDS, set_bytes + + words = NUM_SETS * set_bytes(entries) // 4 + state = torch.zeros((STATE_WORDS,), dtype=torch.int32, device=device) + if self.tp_size == 1: + buf = torch.zeros((words,), dtype=torch.int32, device=device) + self._markov_v3 = (entries, buf.data_ptr(), buf.data_ptr(), state, buf) + return self._markov_v3[1:4] + from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer + + assert self.backend == "mnnvl" + # The communicator of the exchange's MNNVL workspace (allocated with it). + comm = get_or_scale_allreduce_mnnvl_workspace(self.mapping, torch.float32)["comm"] + force_mn = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" + handle = McastGPUBuffer( + words * 4, + self.tp_size, + self.tp_rank, + self.mapping.local_rank, + force_mn or self.mapping.is_multi_node(), + comm.py2f(), + ) + local = handle.get_uc_buffer(self.tp_rank, [words], torch.int32, 0) + mcast = handle.get_mc_buffer([words], torch.int32, 0) + local.zero_() + torch.cuda.synchronize() + comm.Barrier() + self._markov_v3 = (entries, local.data_ptr(), mcast.data_ptr(), state, handle) + return self._markov_v3[1:4] + + def layout(self, rows: int, n_cols: int, device) -> _PairLayout: + # Split each row over enough programs to fill the GPU when rows are few; + # the per-split pairs travel as extra slots of the same exchange. + max_splits = triton.cdiv(n_cols, _LOCAL_BLOCK_V) + splits = 1 + if self.backend == "mnnvl" and self.tp_size > 1: + splits = max(1, min(max_splits, triton.cdiv(_TARGET_PROGRAMS, max(rows, 1)))) + chunk = triton.cdiv(triton.cdiv(n_cols, splits), _LOCAL_BLOCK_V) * _LOCAL_BLOCK_V + splits = max(1, triton.cdiv(n_cols, chunk)) + if self.tp_size == 1 or self.backend == "nccl": + buf = torch.empty((rows, 2 * splits), dtype=torch.float32, device=device) + return _PairLayout(buf, 2 * splits, 0, 0, 0, 0, splits, chunk) + width = triton.next_power_of_2(2 * self.tp_size * splits) + buf = torch.empty((rows, width), dtype=torch.float32, device=device) + lo = 2 * self.tp_rank * splits + return _PairLayout( + buf, width, self.tp_rank * splits, lo, lo + 2 * splits, width, splits, chunk + ) + + def exchange(self, layout: _PairLayout) -> tuple[torch.Tensor, int, int, int]: + """Every rank's pairs: (tensor, row stride, slot stride, number of slots).""" + buf = layout.buffer + rows = buf.shape[0] + if self.tp_size == 1: + return buf, layout.stride_out, 2, layout.splits + if self.backend == "mnnvl": + out = self._mnnvl(buf, all_reduce_params=AllReduceParams()) + return out, layout.stride_out, 2, self.tp_size * layout.splits + from tensorrt_llm._torch.distributed.ops import allgather + + assert layout.splits == 1 + out = allgather(buf, self.mapping, dim=0) + return out, 2, 2 * rows, self.tp_size + + +class K3VocabParallelHead: + """Vocab-parallel token selection over one column-parallel LM head. + + Args: + lm_head: an ``LMHead`` (or ``Linear``) in COLUMN tensor-parallel mode + whose weight holds this rank's contiguous vocab shard. + vocab_size: the full vocabulary size. + """ + + def __init__(self, lm_head, vocab_size: int): + self.lm_head = lm_head + self.mapping = lm_head.mapping + self.vocab_size = int(vocab_size) + self.shard_width = int(lm_head.out_features) + self.col_base = self.mapping.tp_rank * self.shard_width + self.exchange = _PairExchange(self.mapping) + self._markov_w2_shard = None + self._markov_w2_src_ptr = None + logger.info( + f"Vocab-parallel LM-head token selection: tp={self.mapping.tp_size}, " + f"shard={self.shard_width}, vocab={self.vocab_size}, exchange={self.exchange.backend}" + ) + + @staticmethod + def supported(lm_head, vocab_size: int) -> bool: + """Whether ``lm_head`` is a plain vocab-sharded head this class handles.""" + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = getattr(lm_head, "mapping", None) + if mapping is None or mapping.tp_size <= 1 or mapping.enable_attention_dp: + return False + if getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN: + return False + if not getattr(lm_head, "gather_output", False): + return False + if getattr(lm_head, "padding_size", 0) != 0 or getattr(lm_head, "has_any_quant", False): + return False + # The sampler draws Philox noise for 4-column groups of the global vocab. + shard = int(lm_head.out_features) + return shard % 4 == 0 and shard * mapping.tp_size == int(vocab_size) and vocab_size < 2**24 + + def local_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """This rank's ``[rows, shard]`` logits in the head's output dtype.""" + return self.lm_head.apply_linear(hidden_states, self.lm_head.bias) + + def shard_logits(self, hidden_states: torch.Tensor) -> "K3VocabShardedLogits": + """This rank's logits wrapped for a spec worker (see ``K3VocabShardedLogits``).""" + return K3VocabShardedLogits(self, self.local_logits(hidden_states)) + + def gather_full_logits(self, local: torch.Tensor) -> torch.Tensor: + """Full float32 logits, identical to the gathered head + ``.float()``.""" + from tensorrt_llm._torch.distributed import allgather + + return allgather(local, self.mapping, sizes=self.lm_head.gather_output_sizes).float() + + def _launch_local(self, local: torch.Tensor, sampling: Optional[tuple]) -> _PairLayout: + rows, n_cols = local.shape + layout = self.exchange.layout(rows, n_cols, local.device) + k3_vocab_local_select(local, self.col_base, layout, sampling) + return layout + + def _exchange_and_reduce(self, layout: _PairLayout, out: torch.Tensor): + pairs, stride_row, stride_slot, n_slots = self.exchange.exchange(layout) + return k3_vocab_reduce_pairs(pairs, stride_row, stride_slot, n_slots, out) + + def selection_pairs(self, local: torch.Tensor, spec_metadata) -> tuple: + """Every rank's ``(value, global index)`` pairs of the target selection, + not yet reduced: ``(pairs, row stride, slot stride, slots)``. + + The same local kernel and exchange as :meth:`K3VocabShardedLogits.select_tokens` + (argmax for an all-greedy batch, else the Gumbel-max draw), for a caller + that reduces the pairs inside a fused consumer kernel. + """ + assert local.dim() == 2 and local.stride(-1) == 1 + if spec_metadata.is_all_greedy_sample: + layout = self._launch_local(local, None) + else: + from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import ( + GREEDY_TEMPERATURE_THRESHOLD, + ) + + rows = local.shape[0] + layout = self._launch_local( + local, + ( + spec_metadata.temperatures[:rows], + spec_metadata.seeds[:rows], + spec_metadata.offsets[:rows], + float(GREEDY_TEMPERATURE_THRESHOLD), + ), + ) + return self.exchange.exchange(layout) + + def argmax(self, local: torch.Tensor, out_dtype=torch.int64) -> torch.Tensor: + """Global argmax over the vocab of shard logits ``[rows, shard]``.""" + assert local.dim() == 2 and local.stride(-1) == 1 + layout = self._launch_local(local, None) + out = torch.empty((local.shape[0],), dtype=out_dtype, device=local.device) + return self._exchange_and_reduce(layout, out) + + def sample( + self, + local: torch.Tensor, + temperatures: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + greedy_threshold: float, + out_dtype=torch.int64, + ) -> torch.Tensor: + """Temperature-only sampling (Gumbel-max); rows at or below the greedy + temperature threshold take the plain argmax.""" + assert local.dim() == 2 and local.stride(-1) == 1 + layout = self._launch_local(local, (temperatures, seeds, offsets, float(greedy_threshold))) + out = torch.empty((local.shape[0],), dtype=out_dtype, device=local.device) + return self._exchange_and_reduce(layout, out) + + def markov_tokens_chain( + self, + local_block: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + ) -> torch.Tensor: + """:meth:`markov_tokens` with the whole chain in one launch. + + Every step resolves prev, gathers its ``markov_w1`` row, computes this + rank's ``markov_w2`` bias columns, adds the step's shard logits and + publishes the rank's (value, global index) pair; the per-rank pairs + travel through a multicast buffer that every program polls (see + ``_markov_chain_kernel``) instead of an all-reduce launch. The bias is + an FP32 dot of the BF16 row and weights rounded to BF16, like the + ``F.linear`` of :meth:`markov_tokens`, but its summation order differs, + so a draft token can differ where two vocab entries tie to within that + rounding. Needs ``num_reqs <= 16`` and ``block <= 16`` (the caller falls + back otherwise) and the MNNVL (or single-rank) exchange. + """ + num_reqs, block_size, shard = local_block.shape + device = local_block.device + tokens = torch.empty((num_reqs, block_size), dtype=torch.int32, device=device) + if num_reqs == 0 or block_size == 0: + return tokens + assert num_reqs <= _CHAIN_MAX_REQS and block_size <= _CHAIN_MAX_STEPS + w2_shard = self.markov_w2_shard(markov_w2) + assert ( + w2_shard.stride(-1) == 1 and markov_w1.is_contiguous() and local_block.stride(-1) == 1 + ) + programs = triton.cdiv(shard, _MARKOV_BLOCK_V) + # Every program polls for the others' pairs: all of them must be resident. + assert programs <= torch.cuda.get_device_properties(device).multi_processor_count + mcast, local, set_floats = self.exchange.chain_buffers(device) + state = getattr(self, "_chain_state", None) + if state is None or state[0] != (programs, str(device)): + partial = torch.empty( + (_CHAIN_MAX_STEPS * programs * _CHAIN_MAX_REQS * 2,), + dtype=torch.float32, + device=device, + ) + tickets = torch.zeros((_CHAIN_MAX_STEPS,), dtype=torch.int32, device=device) + launch = torch.zeros((1,), dtype=torch.int32, device=device) + self._chain_state = ((programs, str(device)), partial, tickets, launch) + _, partial, tickets, launch = self._chain_state + tp = max(self.exchange.tp_size, 1) + return k3_markov_chain( + local_block, + first_prev_tokens, + markov_w1, + w2_shard, + tokens, + partial, + tickets, + launch, + mcast, + local, + set_floats=set_floats, + vocab=self.vocab_size, + col_base=self.col_base, + tp=tp, + tp_rank=self.mapping.tp_rank if tp > 1 else 0, + ) + + def markov_v3_supported( + self, + local_block: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: Optional[torch.Tensor], + markov_w2: Optional[torch.Tensor], + ) -> bool: + """Whether :meth:`markov_tokens_v3` serves these operands: BF16 logits + and Markov weights (rank a multiple of 16, at most 256), at most 16 + requests and 16 block positions, an SM100-family device with a CTA per + 256 shard rows resident at once, and the MNNVL (or single-rank) + exchange.""" + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + if not IS_CUTLASS_DSL_AVAILABLE or markov_w1 is None or markov_w2 is None: + return False + if self.exchange.tp_size > 1 and self.exchange.backend != "mnnvl": + return False + num_reqs, block_size, shard = local_block.shape + rank_dim = markov_w1.shape[1] + device = local_block.device + if not (1 <= num_reqs <= 16 and 1 <= block_size <= 16) or shard != self.shard_width: + return False + if ( + local_block.dtype != torch.bfloat16 + or markov_w1.dtype != torch.bfloat16 + or markov_w2.dtype != torch.bfloat16 + or local_block.stride(-1) != 1 + or not markov_w1.is_contiguous() + or not markov_w2.is_contiguous() + or markov_w2.shape[1] != rank_dim + or rank_dim % 16 + or rank_dim > 256 + or first_prev_tokens.dtype not in (torch.int32, torch.int64) + or first_prev_tokens.stride(0) != 1 + or markov_w2.data_ptr() % 16 + ): + return False + from .kernels.dspark_markov_greedy import ctas_for, smem_bytes + + props = torch.cuda.get_device_properties(device) + return ( + props.major == 10 + and ctas_for(shard) <= props.multi_processor_count + and smem_bytes(rank_dim) <= props.shared_memory_per_block_optin + ) + + def markov_tokens_v3( + self, + local_block: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + ) -> torch.Tensor: + """:meth:`markov_tokens_chain` as one CuTe DSL launch without the + intra-GPU reduction. + + Each CTA keeps its 256 ``markov_w2`` shard rows in shared memory, adds + the BF16-rounded bias of every request to the step's shard logits in + FP32 (the arithmetic of the other chains) and stores its own + (value, global index) key into every rank's exchange; every CTA polls + all keys of the step (``dspark_markov_greedy``). Same tokens as + :meth:`markov_tokens_chain` wherever both compute the same BF16 bias. + The caller checks :meth:`markov_v3_supported`. + """ + from .kernels.dspark_markov_greedy import ctas_for + from .ops.spec.k3_dspark_markov_greedy import k3_dspark_markov_greedy + + num_reqs, block_size, shard = local_block.shape + device = local_block.device + tokens = torch.empty((num_reqs, block_size), dtype=torch.int32, device=device) + tp = max(self.exchange.tp_size, 1) + box_uc, box_mc, state = self.exchange.markov_v3_buffers(device, tp * ctas_for(shard)) + k3_dspark_markov_greedy( + local_block, + first_prev_tokens, + markov_w1, + self.markov_w2_shard(markov_w2), + tokens, + vocab=self.vocab_size, + tp=tp, + rank=self.mapping.tp_rank if tp > 1 else 0, + state=state, + box_uc=box_uc, + box_mc=box_mc, + multicast=tp > 1, + ) + return tokens + + def markov_w2_shard(self, markov_w2: torch.Tensor) -> torch.Tensor: + """This rank's rows of ``markov_w2`` (cached; a contiguous row slice).""" + if self._markov_w2_src_ptr != markov_w2.data_ptr(): + self._markov_w2_shard = markov_w2[self.col_base : self.col_base + self.shard_width] + self._markov_w2_src_ptr = markov_w2.data_ptr() + return self._markov_w2_shard + + def markov_tokens( + self, + local_block: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: Optional[torch.Tensor], + markov_w2: Optional[torch.Tensor], + ) -> torch.Tensor: + """Greedy DSpark block draft tokens from shard block logits. + + Args: + local_block: ``[B, K, shard]`` shard logits of the block positions. + first_prev_tokens: ``[B]`` anchor token ids (step 0's prev). + markov_w1, markov_w2: ``[vocab, rank]`` Markov head weights, or + ``None`` for a drafter without a Markov head (plain argmax). + Returns: + ``[B, K]`` int32 draft tokens, bit-identical to the per-position + argmax of ``dspark_markov_chain_logits`` on the gathered logits. + """ + num_reqs, block_size, shard = local_block.shape + device = local_block.device + if markov_w1 is None: + flat = local_block.reshape(num_reqs * block_size, shard) + return self.argmax(flat, out_dtype=torch.int32).view(num_reqs, block_size) + tokens = torch.empty((num_reqs, block_size), dtype=torch.int32, device=device) + if num_reqs == 0 or block_size == 0: + return tokens + rank = markov_w1.shape[1] + w2_shard = self.markov_w2_shard(markov_w2) + emb = torch.empty((num_reqs, rank), dtype=markov_w1.dtype, device=device) + pairs, stride_row, stride_slot, n_slots = first_prev_tokens, 0, 0, 1 + for step in range(block_size + 1): + k3_markov_reduce_embed( + pairs, + stride_row, + stride_slot, + n_slots, + first_prev_tokens, + markov_w1, + emb, + tokens, + self.vocab_size, + step, + ) + if step == block_size: + break + # Same F.linear as the reference chain, on this rank's w2 rows. + bias = F.linear(emb, w2_shard) + step_logits = local_block[:, step] + layout = self.exchange.layout(num_reqs, shard, device) + k3_markov_add_local_argmax(step_logits, bias, self.col_base, layout) + pairs, stride_row, stride_slot, n_slots = self.exchange.exchange(layout) + return tokens + + +class K3VocabShardedLogits: + """This rank's vocab shard of the target logits, handed to a spec worker in + place of the gathered float32 logits. + + The worker either selects tokens from it vocab-parallel + (:meth:`select_tokens`) or, for a batch the vocab-parallel samplers do not + cover, materializes the gathered logits with :meth:`gather`. + """ + + __slots__ = ("head", "local") + + def __init__(self, head: K3VocabParallelHead, local: torch.Tensor): + self.head = head + self.local = local + + @property + def device(self) -> torch.device: + return self.local.device + + def gather(self) -> torch.Tensor: + """The gathered ``[rows, vocab]`` float32 logits.""" + return self.head.gather_full_logits(self.local) + + def select_pairs(self, spec_metadata) -> tuple: + """The unreduced exchanged pairs of :meth:`select_tokens` (see + :meth:`K3VocabParallelHead.selection_pairs`).""" + return self.head.selection_pairs(self.local, spec_metadata) + + def select_tokens(self, spec_metadata) -> torch.Tensor: + """One target token per row: the argmax for an all-greedy batch, else a + temperature-only draw (the caller guarantees the sampling mode).""" + if spec_metadata.is_all_greedy_sample: + return self.head.argmax(self.local) + from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import GREEDY_TEMPERATURE_THRESHOLD + + rows = self.local.shape[0] + return self.head.sample( + self.local, + spec_metadata.temperatures[:rows], + spec_metadata.seeds[:rows], + spec_metadata.offsets[:rows], + GREEDY_TEMPERATURE_THRESHOLD, + ) + + +@lru_cache(maxsize=1) +def get_dflash_paged_append() -> Callable[..., None]: + """Load flashinfer's paged K/V append: a plain scatter into an HND page pool.""" + if not IS_FLASHINFER_AVAILABLE: + raise RuntimeError( + "DFlash paged context cache requires flashinfer, which is not installed." + ) + import flashinfer + + return flashinfer.page.append_paged_kv_cache + + +# FP8 E4M3 storage of an MLA drafter's context latent: one per-tensor scale of +# 1.0 for the latent and for the query that reads it -- the scheme of the +# target MLA's FP8 KV cache without calibrated kv scales (quant_scale_kv and +# quant_scale_q default to 1, bmm1_scale = softmax_scale * dq_q * dq_kv, +# bmm2_scale = dq_kv). Values saturate at the format's largest finite value, as +# the target's FP8 cache writes do, instead of overflowing to NaN. +DRAFT_FP8_MAX = float(torch.finfo(torch.float8_e4m3fn).max) +DRAFT_FP8_KV_SCALE = 1.0 + + +def to_draft_fp8(x: torch.Tensor) -> torch.Tensor: + """``x`` as saturated FP8 E4M3 at the drafter's unit per-tensor scale.""" + return x.clamp(-DRAFT_FP8_MAX, DRAFT_FP8_MAX).to(torch.float8_e4m3fn) + + +@dataclass(frozen=True) +class HelixDraftShard: + """One rank's share of a drafter's paged context KV under helix. + + Global page ``g`` of a request belongs to CP rank ``g % cp_size``, which + stores its pages back to back: local page ``g // cp_size``. Every helper + maps global token positions or lengths into that rank-local layout. + """ + + cp_size: int + cp_rank: int + page_size: int + + def local_slots(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Rank-local slot of each global position, and whether this rank owns it. + + An unowned position maps to the slot this rank keeps for the token one + or more pages away, never one of the ``page_size`` tokens around it. + """ + page = torch.div(positions, self.page_size, rounding_mode="floor") + slots = torch.div(page, self.cp_size, rounding_mode="floor") * self.page_size + return slots + positions % self.page_size, page % self.cp_size == self.cp_rank + + def local_lengths(self, lengths: torch.Tensor) -> torch.Tensor: + """Number of owned positions below each global length.""" + span = self.page_size * self.cp_size + rem = lengths % span - self.cp_rank * self.page_size + return torch.div(lengths, span, rounding_mode="floor") * self.page_size + rem.clamp( + 0, self.page_size + ) + + +def combine_helix_partials( + out: torch.Tensor, lse: torch.Tensor, empty: Optional[torch.Tensor] = None +) -> torch.Tensor: + """Merge per-rank partial attention over disjoint KV shards. + + Args: + out: ``[cp, N, H, D]`` partial outputs, one per CP rank. + lse: ``[cp, N, H]`` base-2 log-sum-exp of the matching scores, as + TRTLLM-Gen returns it. + empty: optional ``[cp, N]`` bool, True where a rank held no keys for + a row; its partial is ignored whatever the kernel wrote. + Returns: + ``[N, H, D]`` in ``out``'s dtype. + """ + lse = lse.float() + if empty is not None: + lse = lse.masked_fill(empty.unsqueeze(-1), float("-inf")) + peak = lse.amax(dim=0, keepdim=True) + weight = torch.exp2(lse - peak.clamp_min(torch.finfo(torch.float32).min)) + weight = weight.nan_to_num(nan=0.0) + total = weight.sum(dim=0).clamp_min(torch.finfo(torch.float32).tiny) + merged = (weight.unsqueeze(-1) * out.float().nan_to_num(0.0)).sum(dim=0) + return (merged / total.unsqueeze(-1)).to(out.dtype) + + +def draft3_glue_enabled() -> bool: + """Whether the drafter forward takes its per-step glue off the critical path.""" + return _features().draft3_glue + + +# role -> ((max rows, kernel, split-K tactic (mma_m, mma_n, split_k, ab_stages)), ...) +_PINS = { + # [q_a | kv_a]: 2112 x 7168 + "q_kv_a": ( + (8, "splitk", (64, 16, 4, 8)), + (16, "splitk", (64, 16, 4, 9)), + (32, "splitk", (64, 16, 2, 11)), + ), + # q_b: 768 x 1536 per rank + "q_b": ((16, "splitk", (64, 8, 4, 10)), (32, "splitk", (64, 16, 3, 10))), + # gate_up: 1792 x 7168 per rank + "gate_up": ((16, "splitk", (64, 16, 4, 8)), (32, "splitk", (64, 32, 4, 8))), + # down: 7168 x 896 per rank + "down": ((32, "cublas", None),), + # context latents of all layers: 2880 x 7168 + "ctx_latent": ( + (8, "splitk", (128, 8, 4, 6)), + (16, "splitk", (64, 16, 2, 10)), + (32, "splitk", (64, 32, 2, 9)), + ), + # the draft block's LM head shard: 10240 x 7168 + "lm_head": ((16, "splitk", (128, 16, 1, 6)), (32, "splitk", (128, 32, 1, 5))), +} + + +# Rubin (SM107), the helix16 drafter's per-rank shapes, keyed by (N, K); same entry format. Takes +# precedence over ``_PINS`` for these shapes on SM107. Measured as CUDA graphs of 12 back-to-back +# GEMMs with PDL over > L2 of rotating weights, us per call, cuBLAS -> pinned split-K at 7 / 14 / +# 28 rows. +_PINS_SM107 = { + # [q_a | kv_a] (replicated; scripts/draftfuse/bench_gemm.py, job 2706844): 6.46 / 6.48 / 7.17 + # -> 3.95 / 4.07 / 4.95. + (2112, 7168): ( + (8, "splitk", (64, 8, 4, 12)), + (16, "splitk", (64, 16, 4, 10)), + (32, "splitk", (64, 32, 4, 8)), + ), + # context latents of all layers (replicated; job 2706844): 6.79 / 7.45 / 7.15 -> 4.61 / 4.71 / + # 5.38. + (2880, 7168): ( + (8, "splitk", (64, 8, 2, 12)), + (16, "splitk", (64, 16, 4, 10)), + (32, "splitk", (64, 32, 4, 8)), + ), + # The MLP and LM head shards of the TP16 group (scripts/ull_shell/bench_draft_gemm.py + # BG_TP=16, hecate job 834542, HBM locked): gate_up 6.02 / 6.11 / 5.72 -> 3.66 / 3.86 / 4.54; + # down 4.22 / 3.31 / 3.37 -> 2.50 / 2.48 / 2.60; LM head 15.65 / 14.56 / 16.09 -> 14.12 / + # 14.14 / 14.45. + # gate_up: 2 x 896 + (1792, 7168): ( + (8, "splitk", (64, 8, 4, 12)), + (16, "splitk", (64, 16, 4, 10)), + (32, "splitk", (64, 32, 4, 8)), + ), + # down + (7168, 896): ( + (8, "splitk", (64, 8, 1, 12)), + (16, "splitk", (64, 16, 1, 11)), + (32, "splitk", (64, 32, 1, 9)), + ), + # the draft block's LM head shard + (10240, 7168): ( + (8, "splitk", (128, 8, 2, 6)), + (16, "splitk", (64, 16, 1, 11)), + (32, "splitk", (64, 32, 1, 9)), + ), + # q_b of a helix drafter (latent pages sharded over the group, every head on every rank): + # 64 heads x 192. cuBLAS -> pinned at 8 / 16 / 32 rows: 4.75 / 4.79 / 6.65 -> 4.26 / 4.34 / + # 4.49 us (scripts l0ull2drafter/bench_draft_gemm.py BG_TP=1, hecate job 838031, HBM locked). + (12288, 1536): ( + (8, "splitk", (64, 8, 1, 12)), + (16, "splitk", (64, 16, 1, 11)), + (32, "splitk", (64, 32, 1, 9)), + ), +} + + +def _pins(role: str, n: int, k: int, device: torch.device): + if (n, k) in _PINS_SM107 and _sm107(device.index): + return _PINS_SM107[(n, k)] + return _PINS.get(role) + + +@functools.cache +def _sm107(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) == (10, 7) + + +def draft_gemm_enabled() -> bool: + """Whether ``TLLM_DSPARK_DRAFT_GEMM=1`` pins the drafter's GEMM kernels.""" + return _features().dspark_draft_gemm + + +def draft_gemm( + x: torch.Tensor, weight: torch.Tensor, role: str, strided: bool = False +) -> Optional[torch.Tensor]: + """``x @ weight.T`` with the kernel pinned for ``role`` at this row count. + + Returns None (the caller keeps its module path) when the switch is off, the + row count is above the pinned range, or the operands do not fit the kernel. + ``strided``: a 2-D ``x`` may be a row-major view with a 16-byte-aligned row + stride (e.g. the q_a columns of the [q_a | kv_a] rows), read in place. + """ + if not draft_gemm_enabled() or weight.dim() != 2 or not x.is_cuda: + return None + k = x.shape[-1] + m = x.numel() // k if k else 0 + pins = _pins(role, weight.shape[0], k, x.device) + pin = next((p for p in pins or () if 0 < m <= p[0]), None) + if ( + pin is None + or x.dtype != torch.bfloat16 + or weight.dtype != torch.bfloat16 + or weight.dim() != 2 + or weight.shape[1] != k + or not weight.is_contiguous() + or not x.is_cuda + ): + return None + _, kind, tactic = pin + if kind == "cublas": + return F.linear(x, weight) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + from .kernels.dspark_draft_gemm_push import ensure_cute_runtime + from .kernels.low_m_bf16_splitk import SplitKTactic + from .ops.spec.k3_splitk_dense import k3_splitk_dense + + # detach(): the kernel's DLPack export refuses tensors that require grad, + # which module parameters do even under inference mode. + strided = strided and x.dim() == 2 and x.stride(1) == 1 and x.stride(0) % 8 == 0 + x2d = x.detach() if strided else x.detach().reshape(m, k) + weight = weight.detach() + if ( + not (strided or x2d.is_contiguous()) + or k % 128 + or x2d.data_ptr() % 32 + or weight.data_ptr() % 32 + ): + return None + out = torch.empty((m, weight.shape[0]), dtype=x.dtype, device=x.device) + ensure_cute_runtime() + k3_splitk_dense( + x2d, + weight.t(), + None, + out, + get_env_enable_pdl(), + SplitKTactic(*tactic), + strided_a=strided and not x2d.is_contiguous(), + ) + return out.view(*x.shape[:-1], weight.shape[0]) + + +# Split-K tactics (mma_m, mma_n, split_k, ab_stages) of the push GEMM per (role, max rows); the +# first entry whose row bound covers the call applies. Rubin, TP8 drafter shapes: o_proj +# 7168 x 1024, down_proj 7168 x 1792. TP4 sweep, chain of 10 sites (GEMM + site) with cold +# weights, us per site at 7 / 14 / 28 rows: split 1 at mma_n = the row tile is best for both roles +# (o_proj 6.92 / 7.43 / 10.03, down 7.12 / 7.21 / 10.09; split 2: 7.60 / 7.90 / 11.16; cuBLAS + +# C++ fused all-reduce: 9.76 / 8.83 / 10.86) (scripts/draftfuse/bench_ar.py, job 2706843). +_PUSH_TACTICS = { + "o_proj": ( + (8, (64, 8, 1, 0)), + (16, (64, 16, 1, 0)), + (32, (64, 32, 1, 0)), + (64, (64, 64, 1, 0)), + ), + "down": ((8, (64, 8, 1, 0)), (16, (64, 16, 1, 0)), (32, (64, 32, 1, 0)), (64, (64, 64, 1, 0))), +} + + +def push_tactic(role: str, m: int, n: int, k: int, table=None): + """The ``SplitKTactic`` of the push GEMM for ``role`` at ``m`` rows, or None. A table entry + with zero stages takes the deepest pipeline that fits.""" + from .kernels.dspark_draft_gemm_push import validate_push_tactic + from .kernels.low_m_bf16_splitk import SplitKTactic + + entry = next((t for rows, t in (table or _PUSH_TACTICS).get(role, ()) if m <= rows), None) + if entry is None: + return None + mma_m, mma_n, split_k, stages = entry + for ab_stages in [stages] if stages else range(12, 0, -1): + tactic = SplitKTactic(mma_m, mma_n, split_k, ab_stages) + try: + validate_push_tactic(tactic, m, n, k) + except ValueError: + continue + return tactic + return None + + +class DraftArSites: + """The pushed all-reduce sites of one drafter, on its o_proj all-reduce (a + ``layers.TpAllReduce`` whose MNNVL workspace every drafter all-reduce shares).""" + + def __init__(self, all_reduce): + self.all_reduce = all_reduce + + def linear_ar_norm( + self, + x: torch.Tensor, + weight: torch.Tensor, + role: str, + residual: torch.Tensor, + norm, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """``norm(all_reduce(x @ weight^T) + residual)`` and the new residual, or None if the site + does not apply (the caller then runs its own path).""" + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + from .kernels.dspark_draft_ar import ar_norm_config + from .layers import lamport_push_target + from .ops.spec.k3_draft_ar_norm import k3_draft_ar_norm + from .ops.spec.k3_draft_gemm_push import k3_draft_gemm_push + + if ( + x.dim() != 2 + or x.dtype != torch.bfloat16 + or weight.dtype != torch.bfloat16 + or not x.is_contiguous() + or not weight.is_contiguous() + or not residual.is_contiguous() + or residual.dtype != torch.bfloat16 + or getattr(norm, "weight", None) is None + or x.data_ptr() % 32 + or weight.data_ptr() % 32 + ): + return None + m, k = x.shape + n = weight.shape[0] + if residual.shape != (m, n) or m == 0: + return None + tactic = push_tactic(role, m, n, k) + if tactic is None or not ar_norm_config(m, n, x.device)[0]: + return None + target = lamport_push_target(self.all_reduce, m, n) + if target is None: + return None + pdl = get_env_enable_pdl() + k3_draft_gemm_push(x.detach(), weight.detach(), target.output, target, tactic, pdl=pdl) + out = torch.empty_like(residual) + residual_out = torch.empty_like(residual) + k3_draft_ar_norm( + target.workspace, + target.one_shot, + target.cpp_one_shot, + residual, + norm.weight.detach(), + norm.variance_epsilon, + out, + residual_out, + use_pdl=pdl, + ) + return out, residual_out + + +def draft_embed_site(embedding, ids: torch.Tensor) -> Optional[torch.Tensor]: + """``embedding(ids)`` of a vocab-parallel ``layers.EmbeddingTable`` whose all-reduce + runs over MNNVL, as one owner-multicast kernel; None where it does not apply.""" + weight = getattr(embedding, "weight", None) + all_reduce = getattr(embedding, "all_reduce", None) + mnnvl = getattr(all_reduce, "mnnvl_allreduce", None) + if weight is None or mnnvl is None or not getattr(embedding, "vocab_parallel", False): + return None + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + from .kernels.dspark_draft_ar import ELTS, EMBED_THREADS + from .layers import mnnvl_lamport_push_target + from .ops.spec.k3_draft_embed import k3_draft_embed + + if ( + weight.dtype != torch.bfloat16 + or not weight.is_contiguous() + or ids.dim() != 1 + or ids.dtype not in (torch.int32, torch.int64) + or not ids.is_contiguous() + or ids.shape[0] == 0 + or weight.shape[1] % (ELTS * EMBED_THREADS) + ): + return None + # The table's shards: ceil(vocab / tp) rows per rank (this rank's padded to that), the last + # one possibly shorter. + vocab = embedding.num_embeddings + slice_width = weight.shape[0] + if slice_width != -(-vocab // mnnvl.mapping.tp_size) or embedding.vocab_start % slice_width: + return None + m, hidden = ids.shape[0], weight.shape[1] + target = mnnvl_lamport_push_target(mnnvl, m, hidden, one_shot=True) + if target is None: + return None + k3_draft_embed( + target.workspace, + ids, + weight.detach(), + slice_width, + vocab, + target.output, + use_pdl=get_env_enable_pdl(), + ) + return target.output + + +# Backends that keep the draft's private context K/V in a paged HND pool. +# They share the pool layout, page table and paged append op +_PAGED_ATTENTION_BACKENDS = ("TRTLLM", "FA4") + +# Eager iterations with context requests in which a unified draft cache's host +# block counts are checked against the decoded table (_check_host_block_counts). +_HOST_BLOCK_COUNT_CHECKS = 8 + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig + + +def dflash_context_dtype(draft_model: nn.Module) -> torch.dtype: + """Use projection activations for the context cache, not packed weight storage.""" + projection = getattr(draft_model, "fc", None) + if projection is None: + return torch.bfloat16 + return getattr(projection, "dtype", projection.weight.dtype) + + +def last_accepted_hidden(projected_rows: torch.Tensor, num_accepted: torch.Tensor) -> torch.Tensor: + """Select the target row that predicted each newly committed anchor. + + Rows are [requests, verification tokens, hidden], after the target projection + and hidden normalization: hidden_norm(fc(target_hidden)). Accepted counts + include the target's bonus token, so its predicting row is at count minus one. + """ + indices = (num_accepted - 1).clamp_min(0).long() + return projected_rows.gather( + 1, indices[:, None, None].expand(-1, 1, projected_rows.shape[-1]) + ).squeeze(1) + + +def compute_dflash_ctx_buffer_bytes( + max_batch_size: int, + max_ctx_len: int, + block_size: int, + num_attn_layers: int, + num_kv_heads_per_rank: int, + head_dim: int, + dtype_bytes: int = 2, + attention_backend: str = "VANILLA", + page_size: int = 32, + kv_factor: int = 2, + paged: Optional[bool] = None, +) -> int: + """Per-GPU bytes ``_lazy_init_ctx_buffers`` will allocate for the DFlash + pooled-context K/V buffers. + + Mirrors the allocation exactly: ``max_batch_size + 1`` slots (one scratch + slot for padding/warmup dummies), capacity ``max_ctx_len + block_size`` + per slot, ``kv_factor`` halves for every drafter attention layer (K and V, + so 2; the MLA drafter keeps one latent and no V, so 1). The dense layout + holds ``kv_factor`` contiguous ``[slots, L, capacity, Hkv, D]`` tensors; + the paged layout holds one ``[L, pages, kv_factor, Hkv, page, D]`` tensor + whose per-slot capacity is rounded up to whole pages. + + ``paged`` selects the layout. At runtime the caller passes the resolved + ``_ctx_paged`` flag (the backend's kernels are paged, or the drafter opts + into pages itself). When ``None`` -- the config-time estimate, which has + no drafter to ask -- it falls back to whether ``attention_backend`` is a + paged backend. + """ + num_slots = max_batch_size + 1 + capacity = max_ctx_len + block_size + if paged is None: + paged = attention_backend in _PAGED_ATTENTION_BACKENDS + if paged: + pages_per_slot = (capacity + page_size - 1) // page_size + capacity = pages_per_slot * page_size + per_slot_elems = kv_factor * num_attn_layers * capacity * num_kv_heads_per_rank * head_dim + return num_slots * per_slot_elems * dtype_bytes + + +def dflash_draft_slot_ids( + num_gens: int, + block_size: int, + num_draft_tokens: int, + shift_label: bool, + device="cuda", +) -> torch.Tensor: + """Flat indices into the [num_gens * block_size] drafter block outputs + whose hidden states produce the K draft-token logits per request. + + Plain DFlash (K2.7 convention, shift_label off): the hidden state at + mask slot j (j = 1..K) predicts draft token j; slot 0 (which holds the + anchor/bonus token) is unused. + + DSpark shift_label convention (DeepSpec: labels for a block anchored at + position p are input_ids[p+1 .. p+block_size], so the hidden state at + block slot j predicts the token at position p+1+j): slots 0..K-1 are + used, and slot 0 — the anchor token slot — predicts the first draft + token. + """ + request_bases = torch.arange(num_gens, dtype=torch.long, device=device) * block_size + offsets = torch.arange(num_draft_tokens, dtype=torch.long, device=device) + first_slot = 0 if shift_label else 1 + return (request_bases.unsqueeze(1) + first_slot + offsets.unsqueeze(0)).flatten() + + +def dflash_noise_block_embedding( + embed_tokens: nn.Module, + bonus: torch.Tensor, + mask_token_id: int, + block_size: int, + trained_mask_embedding: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """The drafter's input block per request, [len(bonus), block_size, hidden]: slot 0 the bonus token's embedding, + every other slot the mask token's. + + The mask row is the drafter's own trained one when it has one (``trained_mask_embedding``); otherwise the shared + lookup's. Both rows go through ``embed_tokens.forward`` (NOT ``.weight[...]``) so TP-sharded vocabs mask out + ranks that don't own the token id and all-reduce, the mask row included, so every rank makes the same call. + + ``trained_mask_embedding`` ([hidden]) is used as given: the drafter casts it to its own dtype and device at load. + """ + num_gens = bonus.shape[0] + mask_tok = torch.full((1,), int(mask_token_id), dtype=torch.long, device=bonus.device) + combined_embed = embed_tokens(torch.cat([bonus, mask_tok], dim=0)) + embed_bonus = combined_embed[:num_gens] + embed_mask = combined_embed[num_gens] + if trained_mask_embedding is not None: + embed_mask = trained_mask_embedding + noise_embed_2d = embed_mask.expand(num_gens, block_size, -1).clone() + noise_embed_2d[:, 0, :] = embed_bonus + return noise_embed_2d + + +def dflash_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: + """Positions a drafter that indexes absolute positions must be able to encode. + + The forward reads two position sets past the accumulated context, both + starting from a ctx_len clamped to ``max_ctx``: + + * block decode ``ctx_len + num_accepted + j``, ``j in [0, block_size)`` + * context KV ``ctx_len + [0, max_draft_len]`` + + so the largest index is ``max_ctx + max(block_size, max_draft_len + 1) - 1`` + and a table needs one entry more than that. Derived from the runtime + ``max_ctx`` rather than from model_config: py_executor_creator raises the + engine's max_seq_len past the configured value and never writes it back, so + any config-derived reconstruction drifts from what ctx_len is clamped to. + """ + return int(max_ctx) + max(int(block_size), int(max_draft_len) + 1) + + +def dflash_allocated_ctx_limit( + block_counts: torch.Tensor, page_size: int, block_size: int +) -> torch.Tensor: + """Context length a request may advertise, given the pages it holds. + + Leaves ``block_size`` positions of room: the MLA path writes the block's own + latents at ``ctx_len .. ctx_len + block_size``, so stopping at exactly the + allocation puts the first of them on an unallocated page, whose table entry + ``_refresh_ctx_block_tables`` clamped to 0 -- another request's block. A + silent cross-request write, not a fault. + + Per-request only on the ``python`` V2 backend, which propagates + ``BAD_PAGE_INDEX``. The default ``cpp`` one maps it to 0 + (kvCacheManagerV2Utils.cu), every row counts full, and this degenerates to + the block-table-width bound upstream already applied. Do not read a + per-request guarantee into it there. + """ + return (block_counts * page_size - block_size).clamp(min=0) + + +@dataclass +class K3VocabParallelSelection: + """The tokens a DFlash worker selects from this rank's vocab shard of the + target LM head instead of from the gathered logits. + + It exists when the target LM head is a plain vocab-sharded, gathered TP + head (:meth:`K3VocabParallelHead.supported`) and the deployment has + neither occurrence penalties nor a guided decoder, which both edit the full + target logits. :attr:`head` is the private ``K3VocabParallelHead`` (MNNVL + pair exchange), with Jonas's selection semantics: the sampling decision is + made per batch (:meth:`K3DFlashWorker._vocab_parallel_target_selection_applies`) + and the greedy draft tokens come from + :meth:`K3DFlashWorker._draft_vocab_parallel_head`. + """ + + head: K3VocabParallelHead + + +class K3DFlashWorker(SpecWorkerBase): + """ + Worker for DFlash speculative decoding. + + DFlash uses the draft model with mask tokens to predict multiple + draft tokens in parallel. The DFlash draft model uses cross-attention + to target model hidden states captured from specific layers. + + The target features are projected through fc + hidden_norm and fed + to the draft model as K/V context in cross-attention. The context + accumulates across steps - each step adds newly accepted tokens' + projected hidden states to a per-request buffer, giving the draft + model the full history of target features. + + Reference: https://arxiv.org/pdf/2602.06036 + """ + + def __init__( + self, + spec_config: "DFlashDecodingConfig", + mapping: Mapping, + use_separate_draft_kv_cache: bool = False, + ): + super().__init__(use_separate_draft_kv_cache) + self.spec_config = spec_config + self.mapping = mapping + self._resolved_mask_token_id = None + self._resolved_block_size = None + self._compute_block_size = None + + # Pre-allocated per-slot context K/V buffers, built lazily on the + # first forward. Fixed-size so slot-indexed reads/writes are CUDA + # graph compatible. + self._ctx_buf_inited = False + self._ctx_len = None + # Host shadows of _ctx_len and of each request's prompt progress. + self._ctx_len_host = None + self._req_ctx_pos = {} + # True when the drafter's context lives in the target manager's draft + # layers (the engine unified them); see _lazy_init_ctx_buffers. + self._ctx_unified_cache = False + # The manager layer ids of the drafter's context, and the K-buffer + # offset that manager's per-layer page tables carry. + self._ctx_layer_ids = None + self._ctx_block_offset = 0 + self._ctx_kv_factor = 2 + # Iterations left in which the host block counts are checked against + # the device-decoded table (see _check_host_block_counts). + self._draft_table_checks_left = 0 + # False once a check found those counts disagreeing with the table. + self._ctx_unified_host_counts = True + # Snapshot for rolling back in-place _ctx_len updates when a forward + # fails (or after warmup). See _ensure_spec_dec_state_restored. + self._saved_ctx_len = None + self._saved_ctx_len_host = None + self._saved_req_ctx_pos = None + self._ctx_len_restore_pending = False + # Deferred kv_lens_cuda rewind state (see _prepare_kv_for_draft_forward, + # _apply_kv_rewind_after_draft, _ensure_spec_dec_state_restored). + self._kv_rewind_pending = False + self._kv_rewind_amount = None + self._kv_rewind_nc = None + self._kv_rewind_bs = None + # TLLM_DSPARK_TAIL_V2: the rewind rides in the next-new-tokens kernel + # (see _fused_select_accept, _fused_kv_rewind_operands). + self._fused_kv_rewind_armed = False + self._batch_to_slot_buffer = None + self._max_ctx = 0 + self._ctx_k_buf = None # [max_batch+1, L, max_ctx+block, nkv, hd] + self._ctx_v_buf = None + # Set by _lazy_init_ctx_buffers; every context write switches on the + # cache layout rather than on which kernel reads it. + self._ctx_paged = False + # [L, pages, K/V, nkv, page, hd] when privately allocated, or the draft + # KV cache manager's per-layer pool views when bound to it. + self._ctx_kv_buf = None + # [L] int64 base addresses of the manager's per-layer pool views. + self._ctx_kv_ptrs = None + # Whether a generation step's drafter inputs can be prepared by + # _prepare_paged_latent_gen_inputs; decided with the pool binding. + self._paged_latent_prepare = False + # The manager _ctx_kv_buf is borrowed from, so a swap can be detected. + self._ctx_kv_manager = None + # Static per-slot page table for the private arena. When bound to the + # manager it is replaced by _ctx_block_tables, refreshed per iteration. + self._ctx_page_table = None + self._ctx_block_tables = None # [max_batch+1, max_blocks_per_seq] int32 + self._ctx_block_indptr = None + # Pages the manager handed this request, per batch row. Saturates at + # the table width on the default backend -- dflash_allocated_ctx_limit. + self._ctx_block_counts = None + # Host copy of this iteration's per-row counts, read from the manager's + # caches (unified draft cache); None when only the device holds them. + self._ctx_block_counts_host: list[int] | None = None + self._ctx_block_partial = None # parallel decode's per-chunk block counts + self._ctx_pool_idx = 0 + self._ctx_block_divisor = 1 # encoded offset -> raw pool block index + self._ctx_kv_indptr = None + self._ctx_kv_last_page_len = None + self._ctx_page_size = 32 + self._ctx_pages_per_slot = 0 + self._ctx_paged_append = None + # True only when the drafter reads the draft KV cache manager's pool + # AND that manager joins the target's reuse protocol, which is what + # makes a reused prefix's drafter K/V still addressable. Set in + # _lazy_init_ctx_buffers; read by _store_prefill_context. + self._ctx_reuse_addressable = False + self._dflash_attention_backend = spec_config.attention_backend + # Set by bind_target_lm_head. + self._vocab_parallel: Optional[K3VocabParallelSelection] = None + # The private K3VocabParallelHead per LM head (keyed by id), built on + # first use and shared by the target and draft selections of one head. + self._vocab_parallel_heads = {} + # Set by set_target_model and set_draft_model. + self._target_model_ref: Optional[weakref.ref] = None + self._draft_model_ref: Optional[weakref.ref] = None + # The drafter's helix mapping when its context pages are CP-sharded. + self._helix_mapping = None + # The side-stream context latents' state (see _launch_ctx_side). + self._ctx_side = None + self._ctx_side_ar = None + self._ctx_side_stream = None + + # Slot management (Python, updated in prepare() and eager mode) + self._req_to_slot = {} # request_id -> slot index + self._free_slots = deque() # available slot indices + self._dummy_slot = None # for cudagraph padding or warmup dummy requests + self._graph_dummy_id_floor = 1 << 62 + + # Opt-in acceptance-statistics recorder (None unless + # TLLM_DFLASH_ACCEPT_STATS_DIR is set; see accept_stats.py). Only + # consulted behind `is not None` checks — zero overhead when off. + self._accept_stats = maybe_create_recorder( + spec_config.max_draft_len, getattr(mapping, "rank", 0) or 0 + ) + if self._accept_stats is not None: + logger.info( + f"DFlash: acceptance-statistics recording enabled -> {self._accept_stats.path}" + ) + + # type(self), not a literal: K3DSparkWorker does not override + # __init__, so a hardcoded name reports the base class and the log + # cannot evidence which worker a run actually used. + logger.info( + f"{type(self).__name__} initialized with " + f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" + ) + + @property + def max_draft_len(self) -> int: + return self.spec_config.max_draft_len + + @property + def _draft_tokens_per_req(self) -> int: + """Target-forward tokens per gen request: K drafts + 1 bonus. + + Previously this was 2K (K+1 accepted + K-1 mask fillers), but the + fillers contribute nothing beyond the K+1 prefix the acceptance + check and context store consume, so carrying them through every + target layer is wasted work that dominates step time at large batch. + """ + return self.max_draft_len + 1 + + def _validate_draft_attention_backend(self, draft_model) -> None: + draft_model_dflash_attention_backend = getattr( + draft_model, "dflash_attention_backend", None + ) + if draft_model_dflash_attention_backend is None: + raise ValueError("DFlash draft model is missing dflash_attention_backend.") + if self._dflash_attention_backend == "AUTO": + # The draft model resolved AUTO against its own family and this + # build; adopt that, because _lazy_init_ctx_buffers keys paging and + # the trtllm-gen shape checks off the concrete value. + self._dflash_attention_backend = draft_model_dflash_attention_backend + return + if draft_model_dflash_attention_backend != self._dflash_attention_backend: + raise ValueError( + "DFlash worker and draft model attention backends must match; " + f"worker={self._dflash_attention_backend}, " + f"model={draft_model_dflash_attention_backend}." + ) + + def set_draft_model(self, draft_model) -> None: + super().set_draft_model(draft_model) + self._validate_draft_attention_backend(draft_model) + # Non-owning: the draft model is a sibling module of this worker. + self._draft_model_ref = weakref.ref(draft_model) + self._helix_mapping = getattr(draft_model, "_helix_mapping", None) + # The DFlash 2 selector indexes full-vocab codebooks by draft-token + # id, so a d2t-remapped draft vocab would score the wrong tokens. + if self._d2t is not None and getattr(draft_model, "is_dflash2", False): + raise NotImplementedError( + "DFlash 2 candidate selection requires a shared draft/target vocab " + "(d2t vocab mapping is not supported)." + ) + lilicorr = getattr(draft_model, "lilicorr", None) + if lilicorr is not None: + if self._d2t is not None: + raise NotImplementedError("LiLiCorr requires a shared draft/target vocabulary") + if self.max_draft_len >= lilicorr.block_size: + raise ValueError( + f"LiLiCorr checkpoint was trained with block_size={lilicorr.block_size}. " + f"Set max_draft_len to at most {lilicorr.block_size - 1}; " + f"got {self.max_draft_len}." + ) + if not self.spec_config.target_layer_ids: + raise ValueError( + "LiLiCorr requires nonempty target_layer_ids to capture target hidden states." + ) + + def set_target_model(self, target_model: nn.Module) -> None: + # Non-owning: the target model owns this worker. + self._target_model_ref = weakref.ref(target_model) + + def bind_target_lm_head(self, lm_head: nn.Module, logits_processor: nn.Module) -> None: + """Decide the :class:`K3VocabParallelSelection` over the target LM head. + + The shard logits are ``lm_head(hidden_states)``, which are the target's + logits only when its logits processor is the plain + :class:`LogitsProcessor`; a custom one (a final norm or head transform + before the LM head) keeps the gathered logits. Binding the same head + again keeps its selection. + + The head is the private ``K3VocabParallelHead`` (Jonas's MNNVL pair + exchange), under main's eligibility gates (no guided decoder, no + occurrence penalties, the plain logits processor) plus its own + :meth:`K3VocabParallelHead.supported`. + """ + if ( + self.guided_decoder is not None + or self.spec_config.enable_penalty + or type(logits_processor) is not LogitsProcessor + ): + self._vocab_parallel = None + return + head = self._vocab_parallel_head(lm_head) + if head is None: + self._vocab_parallel = None + return + if self._vocab_parallel is None or self._vocab_parallel.head is not head: + # Jonas's head decides temperature-only sampling per batch + # (_vocab_parallel_target_selection_applies), not here. + self._vocab_parallel = K3VocabParallelSelection(head=head) + + def bind_draft_lm_head(self, draft_model) -> None: + """Post-load hook for the draft side of the selection. + + The greedy draft tokens are chosen per step by + :meth:`_draft_vocab_parallel_head`, so nothing is bound here. + """ + del draft_model + + def setup_aliases(self) -> None: + """Post-load hook: decide the vocab-parallel selection, the target side + and then the draft side. + + The model loader runs it after loading: the target model's logits + processor is final (models install a custom one after the base + constructor), and the drafter has taken over the target modules it + shares (``load_weights_from_target_model``). + """ + target_model = None if self._target_model_ref is None else self._target_model_ref() + if target_model is not None: + self.bind_target_lm_head(target_model.lm_head, target_model.logits_processor) + if self._draft_model_ref is not None: + self.bind_draft_lm_head(self._draft_model_ref()) + + def post_load_weights(self) -> None: + self.setup_aliases() + + def set_guided_decoder(self, guided_decoder) -> bool: + # The guided decoder masks the full target logits. + self._vocab_parallel = None + return super().set_guided_decoder(guided_decoder) + + @property + def target_vocab_parallel_head(self) -> Optional[K3VocabParallelHead]: + return None if self._vocab_parallel is None else self._vocab_parallel.head + + def target_logits( + self, + hidden_states: torch.Tensor, + lm_head, + logits_processor, + attn_metadata, + spec_metadata, + draft_model, + ): + """This rank's vocab shard (``K3VocabShardedLogits``) once the target + head is bound for vocab-parallel selection, the gathered logits + otherwise.""" + head = self.target_vocab_parallel_head + if head is not None: + return head.shard_logits(hidden_states) + return super().target_logits( + hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) + + def _helix_shard(self) -> Optional[HelixDraftShard]: + """This rank's share of the context pages when the drafter is helix-sharded.""" + helix = self._helix_mapping + if helix is None: + return None + return HelixDraftShard(helix.cp_size, helix.cp_rank, self._ctx_page_size) + + def _vocab_parallel_head(self, lm_head) -> Optional[K3VocabParallelHead]: + """The cached private ``K3VocabParallelHead`` of ``lm_head``, or None when + the head is not a plain vocab shard.""" + if lm_head is None: + return None + key = id(lm_head) + if key not in self._vocab_parallel_heads: + vocab = int(getattr(lm_head, "num_embeddings", 0) or 0) + supported = vocab > 0 and K3VocabParallelHead.supported(lm_head, vocab) + self._vocab_parallel_heads[key] = ( + K3VocabParallelHead(lm_head, vocab) if supported else None + ) + return self._vocab_parallel_heads[key] + + def _vocab_parallel_target_selection_applies(self, spec_metadata) -> bool: + """Whether this batch's target tokens can come from shard logits. + + Greedy batches always; temperature-only batches only under the + ``no_topk_no_topp`` sampling mode and without rejection sampling, which + needs the full target distribution. + """ + if spec_metadata.is_all_greedy_sample: + return True + if self._vocab_parallel is None: + return False + if spec_metadata.use_rejection_sampling: + return False + return spec_metadata.advanced_sampling_mode == AdvancedSamplingMode.NO_TOPK_NO_TOPP + + def _draft_vocab_parallel_head(self, draft_model, spec_metadata): + """The private head greedy draft tokens are selected vocab-parallel with, + or None. + + Besides Jonas's conditions, a drafter with a LiLiCorr or DFlash 2 path + selector keeps main's gathered path, which runs those selectors. + """ + if ( + spec_metadata.wants_advanced_draft_sampling + or self.guided_decoder is not None + or self._d2t is not None + or getattr(draft_model, "lilicorr", None) is not None + or getattr(draft_model, "has_candidate_selector", False) + ): + return None + return self._vocab_parallel_head(getattr(draft_model, "lm_head", None)) + + def _vocab_parallel_markov_weights(self, draft_model): + """``(markov_w1, markov_w2)`` of the intra-block Markov head the draft + tokens chain through, or ``(None, None)`` for the plain per-position + argmax. Plain DFlash has no Markov head.""" + del draft_model + return None, None + + def _k3_vocab_parallel_draft_tokens( + self, + draft_head: K3VocabParallelHead, + draft_model, + hidden_states_out: torch.Tensor, + gen_gather_ids: torch.Tensor, + num_gens: int, + num_draft_tokens: int, + inputs: dict, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Greedy draft tokens straight from this rank's vocab shard of the + block logits through the private head, bit-identical to refining the + gathered logits and taking the per-position argmax. + + Returns ``([num_gens, K]`` int32 draft tokens, the ``[num_gens * K, + hidden]`` block hidden states they were computed from). + """ + K = num_draft_tokens + markov_w1, markov_w2 = self._vocab_parallel_markov_weights(draft_model) + # TRTLLM_K3_DRAFT3_GLUE: a contiguous row range of the block output + # (one request) is read in place instead of gathered. + rows = ( + self._block_gather_range(draft_model, num_gens, K, hidden_states_out.shape[0]) + if draft3_glue_enabled() + else None + ) + block_hidden = ( + hidden_states_out[rows[0] : rows[1]] + if rows is not None + else hidden_states_out[gen_gather_ids] + ) + local_block = ( + draft_gemm(block_hidden, draft_head.lm_head.weight, "lm_head") + if getattr(draft_head.lm_head, "bias", None) is None + else None + ) + if local_block is None: + local_block = draft_head.local_logits(block_hidden) + markov = draft_head.markov_tokens + if ( + markov_w1 is not None + and num_gens <= TAIL_V2_MAX_GENS + and (draft_head.exchange.backend == "mnnvl" or draft_head.exchange.tp_size == 1) + # The one-launch chain holds at most 16 requests and 16 block positions. + and num_gens <= 16 + and K <= 16 + ): + markov = draft_head.markov_tokens_chain + local_block = local_block.view(num_gens, K, -1) + if draft_head.markov_v3_supported( + local_block, inputs["first_prev_tokens"], markov_w1, markov_w2 + ): + markov = draft_head.markov_tokens_v3 + gen_draft_tokens = markov( + local_block, + inputs["first_prev_tokens"], + markov_w1, + markov_w2, + ) + return gen_draft_tokens, block_hidden + + def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype, kv_factor=2): + """Fail with the arithmetic before allocating the drafter context arena. + + The arena is dense in max_seq_len and lands *after* the KV cache manager + has already carved its pool, so free_gpu_memory_fraction does not bound + it. Sized off an unpinned max_seq_len it reaches hundreds of GiB, and the + bare torch.zeros that follows reports only the failed allocation, naming + neither the buffer nor the knob that controls it. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + # Reuse the allocation-sizing helper so this runtime check and the + # config-time estimate cannot drift. num_slots and capacity are already + # resolved here, so feed them directly (max_batch_size = num_slots - 1, + # block_size = 0, dtype_bytes = itemsize); kv_factor and the resolved + # paged flag carry the drafter-specific layout. + arena = compute_dflash_ctx_buffer_bytes( + num_slots - 1, + capacity, + 0, + L, + nkv, + hd, + dtype_bytes=itemsize, + attention_backend=self._dflash_attention_backend, + page_size=self._ctx_page_size, + kv_factor=kv_factor, + paged=self._ctx_paged, + ) + + # mem_get_info() excludes the caching allocator's reserved-but-unused + # segments, which torch.zeros() below can serve from. + free = torch.cuda.mem_get_info()[0] + ( + torch.cuda.memory_reserved() - torch.cuda.memory_allocated() + ) + gib = 1 << 30 + logger.info( + f"DFlash: drafter ctx arena needs {arena / gib:.2f} GiB " + f"(slots={num_slots}, layers={L}, capacity={capacity}, " + f"kv_heads={nkv}, head_dim={hd}), {free / gib:.2f} GiB free" + ) + if arena >= free: + # Warn, not raise: this estimate has under-reported real + # headroom, and torch.zeros below is the authority. + logger.warning( + f"DFlash drafter context arena needs {arena / gib:.2f} GiB but only " + f"{free / gib:.2f} GiB is free. It is dense in sequence length: " + f"slots={num_slots} x layers={L} x capacity={capacity} x " + f"kv_heads={nkv} x head_dim={hd} x {itemsize}B x {kv_factor} halves. " + f"capacity comes from max_seq_len (+{self._resolved_block_size} " + f"block slack), so pass an explicit max_seq_len matching the " + f"sequences you actually serve instead of letting it fall back to " + f"the drafter's max_position_embeddings; lowering max_batch_size " + f"scales it down linearly too." + ) + + def _managed_ctx_pool( + self, + draft_kv_cache_manager: BaseResourceManager | None, + L: int, + nkv: int, + hd: int, + dtype: torch.dtype, + kv_factor: int = 2, + ) -> list[torch.Tensor] | None: + """Per-layer views of the draft KV cache manager's pool, or None. + + The drafter's layers are the manager's ``draft_layer_ids`` when it + appends them to the target's (the unified draft cache), else its first + ``L`` layers. ``get_buffers(..., "HND")`` hands back the per-layer shape + the private arena uses; ``dtype`` is the pool's (FP8 for a drafter with + ``context_kv_fp8``). + + A list, not one stacked tensor: only the two consumers that index it by + layer need it, and the two managers disagree on what a stack would mean. + V1 slices one page-major allocation, so the layers are strided views of + a shared storage; V2 wraps each layer's pool base address in its own + tensor. Everything downstream indexes ``[layer_idx]`` first anyway. + + Note the page count is not the drafter's to interpret: a V2 view spans + the whole interleaved pool, not this layer's slice, so only shape[1:] is + checked here; the index space is settled in _init_ctx_block_offsets. + A unified cache cannot fall back to the private arena (the manager + allocates, reuses and transfers the drafter's pages), so a mismatch + there raises. + """ + if draft_kv_cache_manager is None: + # No separate draft KV cache. Attention DP alone does NOT disable + # it: _util.py skips that bail for an external drafter, which is + # what DSpark is, so this drafter does get a manager under DP. + return None + draft_layer_ids = tuple(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) + layer_ids = draft_layer_ids if draft_layer_ids else tuple(range(L)) + if len(layer_ids) != L: + reason = ( + f"the KV cache manager exposes {len(layer_ids)} draft layers but the " + f"drafter has {L}" + ) + if self._ctx_unified_cache: + raise ValueError(f"DFlash: {reason}.") + logger.warning(f"DFlash: {reason}; using the private ctx arena.") + return None + self._ctx_layer_ids = layer_ids + layers = [ + draft_kv_cache_manager.get_buffers(layer_id, kv_layout="HND") for layer_id in layer_ids + ] + base = layers[0] + # kv_factor 1 is the MLA drafter's SELFKONLY pool: one latent, no V. + expected = (kv_factor, nkv, base.size(-2), hd) + for i, layer in enumerate(layers): + if tuple(layer.shape[1:]) != expected or layer.dtype != dtype: + reason = ( + f"draft KV pool layer {i} is {tuple(layer.shape)} / {layer.dtype}, " + f"not {('pages',) + expected} / {dtype}" + ) + if self._ctx_unified_cache: + raise ValueError(f"DFlash: the unified draft KV pool does not match: {reason}.") + # Fall back rather than fail: the private arena is what every + # DFlash model shipped with, so an unrecognized pool costs the + # old memory footprint, not the run. + logger.warning(f"DFlash: {reason}; using the private ctx arena.") + return None + logger.info( + f"DFlash: binding the drafter ctx cache to the draft KV cache manager pool: " + f"{L} x {tuple(base.shape)} {base.dtype}, page_size={base.size(-2)}, " + f"layers={list(layer_ids)}, unified={self._ctx_unified_cache}" + ) + return layers + + def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page_size) -> bool: + """Validate the pool layout and resolve its offset decode constants. + + ``kv_cache_block_offsets`` entries are ``pool_block_index * + num_pool_layers * kv_factor`` plus the K/V field index. What that has to + be divided by to index a ``get_buffers`` view depends on the manager -- + see the stride check below; attention_backend/utils.py hardcodes the V1 + divisor and is not a general contract. A V2 manager with per-layer page + tables also offsets every entry by the layer's K buffer within the + physical slot (``_ctx_block_offset``). + """ + mgr = draft_kv_cache_manager + pool_mapping = mgr.kv_cache_pool_mapping + first_layer_id = self._ctx_layer_ids[0] if self._ctx_layer_ids is not None else 0 + layer_offset = mgr.layer_offsets[first_layer_id] + self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) + per_layer_tables = bool(getattr(mgr, "_use_per_layer_page_tables", False)) + num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) + self._ctx_block_offset = 0 + if isinstance(mgr, KVCacheManagerV2): + converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) + layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] + unsupported = None + if mgr.enable_swa_scratch_reuse: + unsupported = "DFlash managed context does not support SWA scratch pages" + for layer_id in self._ctx_layer_ids: + local_id = mgr.layer_offsets[layer_id] + if ( + mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id + or mgr.get_layer_page_index_scale(layer_id) != converter.scale + ): + unsupported = ( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" + ) + break + if unsupported is not None: + # A unified manager was selected for this layout, so the drafter + # cannot move to private storage without breaking allocation, + # reuse and transfer. A separate manager keeps the private arena. + if self._ctx_unified_cache: + raise NotImplementedError(unsupported) + logger.warning(f"{unsupported}; using the private ctx arena.") + return False + if per_layer_tables: + # Attention tables address the physical pool origin, whereas + # get_buffers() starts at this layer's own K buffer. + self._ctx_block_offset = int(converter.layer_offset) + elif num_pool_layers != L: + logger.warning( + f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " + f"the drafter has {L}; using the private ctx arena." + ) + return False + manager_kv_factor = ( + mgr.get_layer_kv_factor(first_layer_id) + if hasattr(mgr, "get_layer_kv_factor") + else mgr.kv_factor + ) + if manager_kv_factor != self._ctx_kv_factor: + if self._ctx_unified_cache: + raise ValueError( + f"DFlash: the unified draft layers have KV factor {manager_kv_factor}, " + f"the drafter {self._ctx_kv_factor}." + ) + logger.warning( + "DFlash: draft KV layer factor does not match the drafter; " + "using the private ctx arena." + ) + return False + # V1 slices a page-major pool, so stride(0) spans every layer; V2 wraps + # the layer base in a dense view over interleaved pages, so stride(0) is + # one layer slot. Guessing wrong is silent: acceptance quietly drops. + layer_block_elems = manager_kv_factor * nkv * page_size * hd + stride0 = pool[0].stride(0) + if stride0 == layer_block_elems: + self._ctx_block_divisor = manager_kv_factor + elif stride0 == num_pool_layers * layer_block_elems: + self._ctx_block_divisor = num_pool_layers * manager_kv_factor + else: + message = ( + f"DFlash: draft KV pool layer stride {stride0} matches neither V1's " + f"page-major slice ({num_pool_layers * layer_block_elems}) nor V2's " + f"interleaved view ({layer_block_elems})" + ) + if self._ctx_unified_cache: + raise ValueError(f"{message}.") + logger.warning( + f"{message}; using the private ctx arena rather than guessing block indices." + ) + return False + logger.info( + f"DFlash: ctx block offsets pool_idx={self._ctx_pool_idx}, " + f"divisor={self._ctx_block_divisor}, block_offset={self._ctx_block_offset}, " + f"view_pages={pool[0].size(0)}" + ) + return True + + def _init_ctx_block_tables(self, num_slots: int, max_blocks: int) -> None: + """Size the per-iteration block table (and the parallel decode's scratch).""" + self._ctx_block_tables = torch.zeros( + (num_slots, max_blocks), dtype=torch.int32, device="cuda" + ) + self._ctx_block_indptr = torch.arange( + 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" + ) + self._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") + # Per-(request, column chunk) allocated-block counts of the parallel decode. + self._ctx_block_partial = torch.zeros( + num_slots * ((max_blocks + BLOCK_TABLE_CHUNK - 1) // BLOCK_TABLE_CHUNK), + dtype=torch.int32, + device="cuda", + ) + logger.info(f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}") + + @staticmethod + def _check_block_within_allocation( + num_ctx: torch.Tensor, allocated_ctx_limit: torch.Tensor, block_size: int + ) -> None: + """Raise when a request's drafter block would leave its allocated pages. + + The shared manager sizes the request, so a shortfall there would have the + block's scratch write land on the last context position, or on another + request's page. Asynchronous, so it adds no host sync and is safe to + capture in a CUDA graph; on the device it traps instead of raising. + """ + torch._assert_async( + (num_ctx <= allocated_ctx_limit).all(), + f"DFlash: the drafter's block of {block_size} positions falls outside the " + "pages the shared KV cache manager allocated for the request", + ) + + def _refresh_ctx_block_tables( + self, attn_metadata, num_seqs: int, request_ids: list[int] | None = None + ) -> bool: + """Decode this iteration's draft block table into the persistent buffer. + + In place and fixed-shape so a captured graph keeps reading the same + tensor. Returns False when there is no table to fill, i.e. every path + that stayed on the private arena. + + The unified draft cache's pages are rows of the target metadata's + ``kv_cache_block_offsets`` (the draft layers' pool), a separate draft + manager's are ``draft_kv_cache_block_offsets``. For the unified cache the + per-request block counts are also read on the host from the manager's + caches, so the prefill store need not wait for the device. + """ + self._ctx_block_counts_host = None + if self._ctx_block_tables is None or num_seqs <= 0: + return False + offsets_name = ( + "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + src = getattr(attn_metadata, offsets_name, None) + if src is None: + # The table is bound but unfillable. Leaving it at zeros would send + # every request's every layer to pool block 0 -- silently. + raise RuntimeError( + "DFlash: bound to the draft KV pool but " + f"{type(attn_metadata).__name__} carries no " + f"{offsets_name} this iteration." + ) + # Field 0 of the K/V pair: both name the same block for a kv_factor=2 + # pool, and the drafter addresses blocks, not planes. The block count + # is taken before the clamp erases the placeholders; it is informative + # only on the python V2 backend, the cpp one already mapped + # BAD_PAGE_INDEX to 0. The kernels remove the layer offset in int64 + # and leave the metadata's table unchanged. + if tail_v2_enabled(): + # TLLM_DSPARK_TAIL_V2: the same decode with a row's columns spread + # over programs. + k3_draft_block_table( + src[self._ctx_pool_idx, :num_seqs, 0], + self._ctx_block_divisor, + self._ctx_block_tables[:num_seqs], + self._ctx_block_counts[:num_seqs], + self._ctx_block_partial, + block_offset=self._ctx_block_offset, + ) + else: + k3_draft_block_table( + src[self._ctx_pool_idx, :num_seqs, 0], + self._ctx_block_divisor, + self._ctx_block_tables, + self._ctx_block_counts, + block_offset=self._ctx_block_offset, + ) + # Only the context rows' counts are read on the host (the prefill store). + num_contexts = min(int(attn_metadata.num_contexts), num_seqs) + cache_map = getattr(self._ctx_kv_manager, "kv_cache_map", None) + if ( + self._ctx_unified_cache + and self._ctx_unified_host_counts + and cache_map is not None + and request_ids is not None + and num_contexts > 0 + ): + counts = [ + int(cache_map[request_id].num_blocks) if request_id in cache_map else 0 + for request_id in request_ids[:num_contexts] + ] + if self._draft_table_checks_left > 0 and not torch.cuda.is_current_stream_capturing(): + self._draft_table_checks_left -= 1 + if not self._check_host_block_counts(counts): + return True + self._ctx_block_counts_host = counts + return True + + def _check_host_block_counts(self, counts: list[int]) -> bool: + """Compare the manager's block counts of the context rows with the decoded table's. + + Run on the first eager iterations that carry context requests (it reads + the device counts back). The decoded count saturates at the table width + where the manager writes unallocated entries as page 0 and the layer + offset is 0, which bounds nothing, so that is accepted for any host + count within the width. On a disagreement the host counts are dropped + for good and the prefill store reads the device counts instead. + """ + width = self._ctx_block_tables.size(1) + device_counts = self._ctx_block_counts[: len(counts)].tolist() + if all(d == h or (d == width and h <= width) for d, h in zip(device_counts, counts)): + return True + logger.error( + "DFlash: the unified draft cache's host block counts " + f"{counts[:8]} disagree with the decoded table's {device_counts[:8]}; " + "the prefill store reads the device counts from now on." + ) + self._draft_table_checks_left = 0 + self._ctx_unified_host_counts = False + return False + + def _check_unified_helix(self, draft_kv_cache_manager) -> None: + """Require the drafter's helix sharding to be the target manager's. + + A unified cache's draft pages follow the target's helix ledger (global + page g on CP rank g % cp_size, as its local page g // cp_size), which is + what the drafter's own sharding (``_helix_shard``) assumes of its pages. + """ + mapping = draft_kv_cache_manager.mapping + target = (mapping.cp_size, mapping.cp_rank) if mapping.has_cp_helix() else None + helix = self._helix_mapping + drafter = None if helix is None else (helix.cp_size, helix.cp_rank) + if target != drafter: + raise ValueError( + "DFlash: the unified draft KV pages are helix-sharded as " + f"(cp_size, cp_rank)={target}, but the drafter reads them as {drafter}." + ) + + def _lazy_init_ctx_buffers( + self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None + ): + self._validate_draft_attention_backend(draft_model) + # The target manager also holding the drafter's layers (the unified draft + # cache): their pages are rows of the target metadata's block offsets. + self._ctx_unified_cache = ( + draft_kv_cache_manager is not None + and draft_kv_cache_manager is getattr(attn_metadata, "kv_cache_manager", None) + ) + offsets_name = ( + "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + # Only TrtllmAttentionMetadata builds the block offsets. Without them the + # table stays all-zero and every request would read and write pool + # block 0 -- silently, at acceptance 1.0. + if draft_kv_cache_manager is not None and not hasattr(attn_metadata, offsets_name): + if self._ctx_unified_cache: + raise NotImplementedError( + f"DFlash unified KV requires {offsets_name} in attention metadata" + ) + logger.warning( + f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; " + "using the private ctx arena." + ) + draft_kv_cache_manager = None + + if self._ctx_buf_inited: + if self._ctx_kv_manager is draft_kv_cache_manager: + return + # KV cache estimation swaps a probe manager for the real one. + # Holding the old pool's tensors across that writes freed memory. + logger.info("DFlash: draft KV cache manager replaced, rebinding the ctx cache.") + self._ctx_buf_inited = False + + # Recorded even when the buffers stay private, so that the estimation + # pass falling back to the arena still rebinds once the real pool lands. + self._ctx_kv_manager = draft_kv_cache_manager + self._ctx_block_tables = None + self._ctx_block_counts = None + self._ctx_block_partial = None + self._ctx_block_counts_host = None + self._ctx_layer_ids = None + self._ctx_block_offset = 0 + self._ctx_reuse_addressable = False + self._paged_latent_prepare = False + max_batch = spec_metadata.max_num_requests + + # ctx_len is 1:1 with the target's positions, so max_seq_len bounds it. + # max_position_embeddings does not: K3's drafter advertises 1048576, + # which sizes the arena at hundreds of GiB. + config = getattr(draft_model, "config", None) + max_pos = getattr(config, "max_position_embeddings", None) if config else None + # Take the served length from the KV cache manager, which every backend + # carries on the metadata base class. A `max_seq_len` attribute is only + # declared by some metadata classes (TrtllmAttentionMetadata has one -- + # a wrapper returning this same manager value; the base class and + # FlashInferAttentionMetadata do not), so reading the attribute alone + # lets max_position_embeddings win the min() below and sizes the arena + # for the advertised positions instead of the served length. The + # attribute stays as the fallback for no-cache attention, where it is + # set by hand and there is no manager to ask. + cache_manager = getattr(attn_metadata, "kv_cache_manager", None) + runtime_max = getattr(cache_manager, "max_seq_len", None) + if runtime_max is None: + runtime_max = getattr(attn_metadata, "max_seq_len", None) + candidates = [c for c in (runtime_max, max_pos) if c is not None] + self._max_ctx = min(candidates) if candidates else 8192 + # py_executor_creator inflates max_seq_len by at most this much for spec + # dec (one term is overlap-only), so compare against what is served. + from tensorrt_llm._torch.speculative.utils import get_num_extra_kv_tokens + + spec_slack = 2 * (self.spec_config.tokens_per_gen_step - 1) + get_num_extra_kv_tokens( + self.spec_config + ) + if runtime_max is not None and self._max_ctx < runtime_max - spec_slack: + # Warn, not raise: the per-request skip in _store_prefill_context + # keeps an oversized request from taking the executor down. + logger.warning( + f"DFlash drafter covers {self._max_ctx} positions but the engine serves " + f"up to {runtime_max - spec_slack}. max_position_embeddings={max_pos} is " + "the binding constraint, so requests past it will draft nothing." + ) + + dtype = dflash_context_dtype(draft_model) + # An FP8 drafter context (context_kv_dtype) is held in an FP8 pool. + pool_dtype = torch.float8_e4m3fn if getattr(draft_model, "context_kv_fp8", False) else dtype + + # Reserve slot index max_batch as a scratch slot for padding/unknown + # dummies; real requests only draw slots 0..max_batch-1, so dummy + # writes land here and can't corrupt a real request's context. + self._dummy_slot = max_batch + num_slots = max_batch + 1 + # CUDA-graph padding uses ids counting down from + # CUDA_GRAPH_DUMMY_REQUEST_ID; a floor of one draft block separates them + # from real ids. Imported lazily to avoid an import cycle. + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + + self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len + + self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") + self._ctx_len_host = [0] * num_slots + self._batch_to_slot_buffer = DeviceCopyBuffer( + torch.zeros(max_batch, dtype=torch.long, device="cuda"), + pin_memory=prefer_pinned(), + ) + self._free_slots = deque(range(max_batch)) + self._req_to_slot = {} + self._req_ctx_pos = {} + + # checkpoint's trained block width + self._resolved_block_size = getattr(draft_model, "block_size", None) or ( + self.max_draft_len + 1 + ) + # what the draft forward actually computes + self._compute_block_size = self._draft_block_width(draft_model) + if self._compute_block_size > self._resolved_block_size: + slack = self._compute_block_size - self.max_draft_len + raise ValueError( + f"DFlash checkpoint was trained with block_size={self._resolved_block_size}." + f"Lower max_draft_len to at most {self._resolved_block_size - slack}." + f"Current max_draft_len={self.max_draft_len}." + ) + if self._ctx_unified_cache: + reserved = getattr(draft_kv_cache_manager, "_shared_draft_block_width", 0) + if self._compute_block_size > reserved: + raise ValueError( + f"Drafter block width {self._compute_block_size} exceeds the " + f"{reserved} positions reserved by the shared KV cache manager. " + "The manager must reserve the drafter's resolved block width." + ) + + # +block_size slack lets each backend append the per-iteration noise + # K/V after the accumulated request context. + assert hasattr(draft_model, "_build_fused_kv_buffers"), ( + "DFlash draft model must define _build_fused_kv_buffers." + ) + draft_model._build_fused_kv_buffers() + L = draft_model._num_attn_layers + nkv = draft_model._num_kv_heads + hd = draft_model._head_dim + kv_factor = getattr(draft_model, "_kv_factor", 2) + self._ctx_kv_factor = kv_factor + capacity = self._max_ctx + self._compute_block_size + # Published before any drafter table is built, so a drafter may read it + # instead of its config cap. _max_ctx is attn_metadata.max_seq_len, which + # py_executor_creator raises past model_config.max_seq_len and never + # writes back. _compute_block_size: the block decode's j runs over it. + draft_model._runtime_position_ceiling = dflash_position_ceiling( + self._max_ctx, self._compute_block_size, self.max_draft_len + ) + # A drafter asks for pages on its own (the MLA one does, to get its + # footprint under free_gpu_memory_fraction) while still reading them + # eagerly. + use_paged = self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS or getattr( + draft_model, "_paged_ctx_cache", False + ) + self._ctx_paged = use_paged + if use_paged: + pool = self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, pool_dtype, kv_factor) + # The manager's page size wins when bound to it: its pool is already + # carved, so the drafter adopts the geometry rather than imposing one. + page_size = self._ctx_page_size if pool is None else pool[0].size(-2) + self._ctx_page_size = page_size + # Settle the block table before committing to the pool: it is the + # last thing that can rule the pool out, and falling back after + # taking the pool branch would leave no buffer allocated at all. + if pool is not None and not self._init_ctx_block_offsets( + draft_kv_cache_manager, pool, L, nkv, hd, page_size + ): + pool = None + if pool is None: + if self._ctx_unified_cache: + raise ValueError("DFlash could not bind the unified manager's draft KV layers") + if getattr(draft_model, "context_kv_fp8", False): + raise ValueError( + "speculative_config.context_kv_dtype='fp8' needs the drafter's context " + "KV in an FP8 draft KV cache manager pool." + ) + # Per-slot reservation is an arena-only notion; the pool path + # takes its pages from the manager one iteration at a time. + self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size + total_pages = num_slots * self._ctx_pages_per_slot + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) + # HND layout consumed by both FlashInfer's paged append and the + # TRTLLM-Gen launcher: [L, pages, K/V, Hkv, page, D]. + self._ctx_kv_buf = torch.zeros( + (L, total_pages, kv_factor, nkv, page_size, hd), + dtype=dtype, + device="cuda", + ) + self._ctx_page_table = torch.arange( + total_pages, dtype=torch.int32, device="cuda" + ).view(num_slots, self._ctx_pages_per_slot) + self._ctx_kv_indptr = torch.arange( + 0, + total_pages + 1, + self._ctx_pages_per_slot, + dtype=torch.int32, + device="cuda", + ) + else: + # Nothing reserved: pages come from the manager's per-request + # block table one iteration at a time, so the footprint follows + # the sequences served rather than max_batch x max_seq_len. + if self._ctx_unified_cache: + self._check_unified_helix(draft_kv_cache_manager) + self._draft_table_checks_left = _HOST_BLOCK_COUNT_CHECKS + self._ctx_unified_host_counts = True + self._ctx_kv_buf = pool + self._ctx_kv_ptrs = torch.tensor( + [layer.data_ptr() for layer in pool], dtype=torch.int64, device="cuda" + ) + self._init_ctx_block_tables(num_slots, draft_kv_cache_manager.max_blocks_per_seq) + # The fused preparation stores single-latent context entries + # projected from the captured target features through fc and + # hidden_norm (which every loaded DFlash drafter owns), as the + # eager path does; LiLiCorr drafters keep the eager path, and so + # does a helix-sharded drafter (rank-local pages), which only + # _prepare_gen_drafter_inputs_fused handles. + self._paged_latent_prepare = ( + kv_factor == 1 + and not draft_model._uses_lilicorr + and self._helix_mapping is None + ) + # Only a manager that publishes and matches its own blocks can + # hand back a reused prefix's drafter K/V; the private arena and + # an unpaired draft pool both start every request at 0. + self._ctx_reuse_addressable = self._ctx_unified_cache or bool( + getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) + ) + self._ctx_kv_last_page_len = torch.full( + (num_slots,), page_size, dtype=torch.int32, device="cuda" + ) + else: # dense private arena (VANILLA, unpaged) + if self._ctx_unified_cache: + raise ValueError( + "DFlash: the drafter reads a dense context arena, not the unified " + "manager's draft KV pages." + ) + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) + kv_shape = (num_slots, L, capacity, nkv, hd) + self._ctx_k_buf = torch.zeros(kv_shape, dtype=dtype, device="cuda") + # An MLA drafter stores one latent per token; leaving _ctx_v_buf as + # None is what tells the write paths below there is no V half. + self._ctx_v_buf = ( + torch.zeros(kv_shape, dtype=dtype, device="cuda") if kv_factor == 2 else None + ) + self._ctx_buf_inited = True + + if not self._ctx_reuse_addressable and getattr( + draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False + ): + # Every fallback above (shape mismatch, stride mismatch, pool layer + # count) logs its own reason and then costs only memory -- except + # this one, where the scheduler goes on matching prefixes the + # drafter can no longer read. Silent, and it looks exactly like the + # unpaired build in every metric but acceptance length. + logger.warning( + "DFlash: the draft KV cache manager joins block reuse but the drafter " + "fell back to a private ctx arena, so a reused prefix's drafter K/V " + "stays unreachable and acceptance length drops as if unpaired." + ) + + logger.info( + f"DFlash: allocated ctx buffers: max_batch={max_batch}, " + f"max_ctx={self._max_ctx}, dtype={dtype}, " + f"dflash_attention_backend={self._dflash_attention_backend}, " + f"reuse_addressable={self._ctx_reuse_addressable}, " + f"unified={self._ctx_unified_cache}" + ) + + def _ctx_paged_index_args(self): + """(kv_indices, kv_indptr) for the append op, and what ``rows`` means. + + Bound to the manager the page table is per-request and refreshed every + iteration, so rows are batch positions; on the private arena it is the + static per-slot table and rows are slots. Both are dense and uniformly + strided, so kv_indptr is the same arange either way -- only the width + differs. + """ + if self._ctx_block_tables is not None: + table = self._ctx_block_tables + indptr = self._ctx_block_indptr + else: + table = self._ctx_page_table + indptr = self._ctx_kv_indptr + return table.flatten().contiguous(), indptr + + def _assign_slot( + self, + req_id: int, + reset: bool = False, + updates: Optional[dict] = None, + bases: Optional[dict] = None, + ): + """Get (or refresh) this request's context slot; None when none is free. + + ``reset`` recycles an existing slot, which is how a reused request id is + told apart from a continuation. Pass ``updates`` to defer the length + resets into a caller's batch instead of writing _ctx_len per request; + either way the host mirror stays in step with the device tensor. + ``bases`` is the base-page batch of main's ``DFlashSpecMetadata.prepare()``; + this worker's slots always read from position 0, so it is left untouched. + """ + + def clear(slot: int) -> None: + if updates is None: + self._write_ctx_len({slot: 0}) + else: + updates[slot] = 0 + + if reset and req_id in self._req_to_slot: + old_slot = self._req_to_slot.pop(req_id) + clear(old_slot) + self._free_slots.append(old_slot) + if req_id not in self._req_to_slot: + if not self._free_slots: + return None + slot = self._free_slots.popleft() + self._req_to_slot[req_id] = slot + clear(slot) + return self._req_to_slot[req_id] + + def _get_ctx_paged_append(self) -> Callable[..., None]: + """flashinfer paged append, shared by the TRTLLM and FA4 backends.""" + if self._ctx_paged_append is None: + self._ctx_paged_append = get_dflash_paged_append() + return self._ctx_paged_append + + def _store_context_kv_paged( + self, + k: torch.Tensor, + v: torch.Tensor, + rows: torch.Tensor, + positions: torch.Tensor, + owned: Optional[torch.Tensor] = None, + ) -> None: + """Write K/V at rank-local ``positions``. + + ``owned`` (helix): False entries rewrite their slot's current value, + which keeps a fixed-size, CUDA-graph-safe store. + """ + fp8_cache = self._ctx_kv_buf[0].dtype == torch.float8_e4m3fn + helix = self._helix_mapping is not None + if v is None or self._dflash_attention_backend == "VANILLA" or fp8_cache or helix: + # Single-latent (MLA) cache: one vector per token per layer, so + # there is no K/V interleave for append_paged_kv_cache to do and the + # write is a plain scatter into the page the block table names. + # VANILLA (a unified cache's pages), FP8 and helix-sharded caches + # write the same way. + table = ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ) + page_size = self._ctx_page_size + pages = table[rows.long(), positions.long() // page_size].long() + offsets = positions.long() % page_size + keep = None if owned is None else owned.view(-1, *([1] * (k.dim() - 2))) + for layer_idx in range(k.size(1)): + pool = self._ctx_kv_buf[layer_idx] + # Byte views for FP8: index_put support varies across builds. + view = pool.view(torch.uint8) if fp8_cache else pool + for plane, new in ((0, k), (1, v)): + if new is None: + continue + if fp8_cache and v is None: + # The MLA latent saturates (to_draft_fp8), as its block + # decode quantizes the query it is read with. + new = to_draft_fp8(new[:, layer_idx]) + else: + new = new[:, layer_idx].to(pool.dtype) + if fp8_cache: + new = new.view(torch.uint8) + if keep is not None: + new = torch.where(keep, new, view[pages, plane, :, offsets]) + view[pages, plane, :, offsets] = new + return + append_paged_kv_cache = self._get_ctx_paged_append() + + # Convert at most once and reuse for every layer. Calling ``.to()`` in + # the op call would allocate converted tensors once per layer when the + # inputs are not already int32. + rows_i32 = rows.to(torch.int32) + positions_i32 = positions.to(torch.int32) + kv_indices, kv_indptr = self._ctx_paged_index_args() + for layer_idx in range(k.size(1)): + pool = self._ctx_kv_buf[layer_idx] + append_paged_kv_cache( + append_key=k[:, layer_idx].to(pool.dtype).contiguous(), + append_value=v[:, layer_idx].to(pool.dtype).contiguous(), + batch_indices=rows_i32, + positions=positions_i32, + paged_kv_cache=pool, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=self._ctx_kv_last_page_len, + kv_layout="HND", + ) + + def _prepare_attn_metadata_for_dflash(self, attn_metadata, spec_metadata): + """Save attn_metadata fields that DFlash modifies during forward.""" + is_capturing = torch.cuda.is_current_stream_capturing() + + if spec_metadata.is_cuda_graph and not is_capturing: + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda", "kv_lens_cuda") + else: + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") + + def _prepare_kv_for_draft_forward( + self, + attn_metadata, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + ): + """Adjust kv_lens_cuda so the draft model sees correct RoPE positions.""" + if hasattr(attn_metadata, "kv_lens_cuda"): + self._kv_rewind_amount = 1 - num_accepted_tokens[num_contexts:batch_size] + self._kv_rewind_nc = num_contexts + self._kv_rewind_bs = batch_size + + if batch_size > num_contexts: + attn_metadata.kv_lens_cuda[num_contexts:batch_size] += 1 + self._kv_rewind_pending = True + + attn_metadata.update_for_spec_dec() + + def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): + """Apply the deferred kv_lens rewind after the draft forward.""" + self._kv_rewind_pending = False + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if is_warmup: + # kv_lens_cuda was saved by prepare_for_spec_dec in this mode and + # is restored wholesale, so no rewind is needed. + return + + if self._kv_rewind_amount is not None and hasattr(attn_metadata, "kv_lens_cuda"): + nc = self._kv_rewind_nc + bs = self._kv_rewind_bs + attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount + attn_metadata.kv_lens_cuda[nc:bs].clamp_(min=0) + + def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> None: + """Add this step's accepted tokens to the context lengths of ``slots``, on the device. + + Padding, warmup and attention-DP dummy rows share the dummy slot, which no request owns + or frees. Emptying it right after the add keeps it at 0 between steps: those rows attend + no context, and their next context K / V starts at column 0, inside the padding + request's own pages (past them, its page-table row maps page 0, another request's). + The store is to a fixed slot, so a captured step replays it. + """ + self._ctx_len[slots] += num_accepted + self._ctx_len.clamp_(max=self._max_ctx) + self._ctx_len[self._dummy_slot].zero_() + + def _rebuild_batch_to_slot(self, request_ids: list[int]) -> None: + mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] + self._batch_to_slot_buffer.update(torch.tensor(mapping, dtype=torch.long)) + + def _seed_generation_slot(self, slot: int, position: int, updates: dict, bases: dict) -> None: + """Record the context length of a slot newly assigned to a generation row whose first + token sits at absolute ``position`` (main's ``DFlashSpecMetadata.prepare()`` calls this + for a disaggregated generation request, or one that lost its slot). + + On the unified cache the draft layers travel with the target's prompt KV, so the slot + holds drafter K/V for every position below ``position``. The private arena holds none + for such a request, so its slot starts empty. Either way the slot reads from position 0, + so ``bases`` is left untouched. + """ + updates[slot] = min(position, self._max_ctx) if self._ctx_unified_cache else 0 + + def _write_ctx_len( + self, updates: dict[int, int], base_updates: Optional[dict[int, int]] = None + ) -> None: + """Apply a slot-to-length mapping to _ctx_len in one async scatter. + + ``base_updates`` comes from main's ``DFlashSpecMetadata.prepare()``; this worker has no + base pages (its slots always read from position 0), so it is ignored. + """ + if not updates: + return + for slot, value in updates.items(): # Update host side mirror + self._ctx_len_host[slot] = value + + slots, values = zip(*updates.items()) + pinned = prefer_pinned() + slots_pinned = torch.tensor(slots, dtype=torch.long, pin_memory=pinned) + vals_pinned = torch.tensor(values, dtype=torch.long, pin_memory=pinned) + self._ctx_len.index_copy_( + 0, + slots_pinned.to("cuda", non_blocking=True), + vals_pinned.to("cuda", non_blocking=True), + ) + + def _restore_ctx_len_host(self) -> None: + """Roll the host shadows back to their pre-forward snapshot.""" + if self._saved_ctx_len_host is not None: + self._ctx_len_host = list(self._saved_ctx_len_host) + if self._saved_req_ctx_pos is not None: + self._req_ctx_pos = dict(self._saved_req_ctx_pos) + + def _store_prefill_context( + self, + draft_model, + spec_metadata: "DFlashSpecMetadata", + attn_metadata, + position_ids: torch.Tensor, + total_target_tokens: int, + ): + """Capture prefill hidden states and store as initial accumulated context. + + During prefill (context requests), the target model processes all prompt + tokens. We project their captured hidden states through fc + hidden_norm + and store them per-request so the draft model can use the full prompt + as cross-attention context on subsequent gen steps. + """ + if not hasattr(draft_model, "fc") or not hasattr(draft_model, "hidden_norm"): + return + + num_ctx_tokens = attn_metadata.num_ctx_tokens + if num_ctx_tokens == 0: + return + + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + if captured_hs is None: + return + # Project context tokens through fc + hidden_norm + ctx_proj = draft_model.project_target_hidden(captured_hs[:num_ctx_tokens]) + + # Split by request and store/append accumulated context. + # Context requests may arrive in chunks (chunked prefill), so we must + # APPEND successive chunks for the same request rather than overwriting. + # A chunk is only an append if it starts where the last one ended + + # Scalar (non-rope) token indices on the host. The device tensor holds + # the same values, but reading it costs a D2H sync per request. + host_pos = spec_metadata.host_position_ids + # Slot -> value, uploaded once after the loop: assigning a Python scalar + # through a CUDA index copies from pageable memory and waits for the + # stream, i.e. for the whole target forward of this step. A failure + # mid-loop needs no partial upload: the forward's snapshot restores both. + ctx_len_updates = {} + offset = 0 + num_contexts = attn_metadata.num_contexts + shard = self._helix_shard() + # A context write addresses pages through row i of the block table, so + # it must fit that row's allocation as well as the arena's length. The + # unified draft cache read this iteration's counts from the manager on + # the host; reading them back from the device would wait for the whole + # target forward of this step, so only a separate draft manager's table + # needs one tolist() for the loop. Under helix each local page stands + # for cp_size global ones. + cp = 1 if shard is None else shard.cp_size + page_span = cp * self._ctx_page_size + host_counts = self._ctx_block_counts_host + if self._ctx_block_tables is None: + ctx_alloc = None + elif host_counts is not None and len(host_counts) >= num_contexts: + ctx_alloc = [count * page_span for count in host_counts[:num_contexts]] + else: + ctx_alloc = (self._ctx_block_counts[:num_contexts] * page_span).tolist() + for i in range(num_contexts): + req_id = spec_metadata.request_ids[i] + slen = int(attn_metadata._seq_lens[i]) + chunk_proj = ctx_proj[offset : offset + slen].detach() + chunk_pos = position_ids[offset : offset + slen].long().detach() + + if slen == 0: + first_pos = 0 + elif host_pos is not None: + first_pos = int(host_pos[offset]) + else: + first_pos = chunk_pos[0].item() + + # A chunk continues what we already stored only if it starts exactly + # where the previous one ended. Everything else -- a fresh request, a + # request id reused after completion, a prefill restarted after + # preemption -- has to start from a clean slot. + reset = self._req_ctx_pos.get(req_id) != first_pos + if self._assign_slot(req_id, reset=reset, updates=ctx_len_updates) is None: + logger.warning("DFlash: no free slots, skipping context store") + self._req_ctx_pos.pop(req_id, None) + offset += slen + continue + + slot = self._req_to_slot[req_id] + cur = ctx_len_updates.get(slot, self._ctx_len_host[slot]) + cap = self._max_ctx if ctx_alloc is None else min(self._max_ctx, ctx_alloc[i]) + if cur == 0 and first_pos > 0 and self._ctx_reuse_addressable: + # Reuse let the target skip [0, first_pos), so only the tail + # reaches the capture buffer. The prefix K/V are still in the + # paired pool: precompute_context_kv is per-token in (hidden, p). + cur = first_pos + if cur + slen > cap: + # Request-level, like the no-free-slots path above: truncating + # would silently draft from a stale prefix, but killing the + # forward would take every other in-flight request with it. + logger.warning( + f"DFlash: ctx overflow on slot {slot} " + f"({cur} + {slen} > {cap}); skipping its context " + "store, so this request drafts nothing." + ) + self._req_to_slot.pop(req_id, None) + self._req_ctx_pos.pop(req_id, None) + ctx_len_updates[slot] = 0 + self._free_slots.append(slot) + offset += slen + continue + end = cur + slen + actual = end - cur + # Anchor to the true position rather than accumulating, and advance + # even when the slot is full so the next chunk still reads as a + # continuation. + self._req_ctx_pos[req_id] = first_pos + slen + if actual > 0: + cache_dtype = ( + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + ) + if cache_dtype == torch.float8_e4m3fn: + # The K/V projection runs in the drafter's weight dtype; + # the store converts its outputs to the cache dtype. + cache_dtype = chunk_proj.dtype + chunk_proj_cast = chunk_proj[:actual].to(cache_dtype) + ctx_len_updates[slot] = end + # Precompute post-norm/post-RoPE K,V for this prefill chunk + # so decode iters can read without re-projecting. + # + # RoPE positions must be slot-local (cur..end), NOT the + # target-forward positions in chunk_pos: gen-time appends and + # draft queries number positions from _ctx_len, so the stored + # keys must share that numbering. Without KV-block reuse the + # two coincide, but on a prefix-cache hit chunk_pos starts at + # the reused-token count, which would offset every stored key + # from the numbering the queries use (RoPE attention depends + # only on position differences, so the uniform slot-local + # renumbering is exact). + local_pos = torch.arange(cur, end, dtype=torch.long, device="cuda") + chunk_k, chunk_v = draft_model.precompute_context_kv(chunk_proj_cast, local_pos) + # chunk_k/v: [actual, L, nkv, hd] → [L, actual, nkv, hd] + if self._ctx_paged: + # Manager block tables are keyed by batch position, the + # private arena's by slot. See _ctx_paged_index_args. + row = i if self._ctx_block_tables is not None else slot + positions, store_k, store_v = local_pos, chunk_k, chunk_v + if shard is not None: + # Eager prefill: keep only this rank's pages. + positions, owned = shard.local_slots(positions) + positions, store_k = positions[owned], chunk_k[owned] + store_v = None if chunk_v is None else chunk_v[owned] + self._store_context_kv_paged( + store_k, + store_v, + torch.full((positions.numel(),), row, dtype=torch.long, device="cuda"), + positions, + ) + else: # VANILLA DFlash backend (FlashAttention) + self._ctx_k_buf[slot, :, cur:end] = chunk_k.permute(1, 0, 2, 3) + if chunk_v is not None: + self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) + offset += slen + + self._write_ctx_len(ctx_len_updates) + + def _check_no_real_generation( + self, spec_metadata, num_contexts: int, batch_size: int, kv_cache_manager + ): + """Reject a generation request on a context-KV-only drafter. + + Only dummies reach a context server's generation rows: warmup and + attention-DP padding requests (flagged by the target KV cache manager) + and CUDA-graph padding. A real one means the server is not context-only. + """ + is_dummy = getattr(kv_cache_manager, "_request_id_to_is_dummy", {}) + real = [ + request_id + for request_id in spec_metadata.request_ids[num_contexts:batch_size] + if request_id != ATTENTION_DP_DUMMY_REQUEST_ID + and request_id < self._graph_dummy_id_floor + and not is_dummy.get(request_id, False) + ] + if real: + raise RuntimeError( + f"Generation requests {real[:4]} reached a DSpark drafter built with " + "speculative_config.context_kv_only, which never drafts; set it only " + "on disaggregated context servers." + ) + + def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): + # Restore first (in warmup mode kv_lens_cuda was saved and comes back + # wholesale), then apply any pending rewind for the other modes so a + # failed draft forward does not leave kv_lens_cuda incremented. + super()._ensure_spec_dec_state_restored(attn_metadata, spec_metadata) + if ( + getattr(self, "_kv_rewind_pending", False) + and attn_metadata is not None + and spec_metadata is not None + ): + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + if ( + getattr(self, "_ctx_len_restore_pending", False) + and self._ctx_len is not None + and self._saved_ctx_len is not None + ): + # A failed forward must not keep this iteration's in-place + # _ctx_len updates: roll back to the pre-forward snapshot. + self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_host() + self._ctx_len_restore_pending = False + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + logits_vocab_shard = False + K = spec_metadata.runtime_draft_len + if isinstance(logits, K3VocabShardedLogits): + if K > 0 and self._vocab_parallel_target_selection_applies(spec_metadata): + # The target tokens come from the shard: this step returns + # the shard, which the engine gathers only for a logits + # post-processor. + raw_logits = logits.local + logits_vocab_shard = True + else: + logits = raw_logits = logits.gather() + + if K == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + # Lazy init buffers and attach worker reference for prepare() + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + previous_ctx_len = self._ctx_len + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) + spec_metadata._dflash_worker = self + if ( + self._ctx_unified_cache + and self._ctx_len is not previous_ctx_len + and spec_metadata.request_ids + ): + # New buffers clear request slots; seed them before generation. + spec_metadata.prepare() + # Before any store: prefill and decode both address pages through it. + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + self._refresh_ctx_block_tables(attn_metadata, batch_size, spec_metadata.request_ids) + + # Save context lengths so both warmup and a failed forward can roll + # back the in-place _ctx_len updates made during drafting. + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # Never allocate the snapshot while capturing a CUDA graph: the + # clone would live in the graph memory pool, and its replay-time + # writes could alias blocks reused by later captures. Rollback is + # only meaningful for eager/warmup forwards anyway; a failure + # during capture aborts the graph itself, and captured ops do not + # mutate _ctx_len until replay. + self._saved_ctx_len = self._ctx_len.clone() + self._saved_ctx_len_host = list(self._ctx_len_host) + self._saved_req_ctx_pos = dict(self._req_ctx_pos) + self._ctx_len_restore_pending = True + + self._ctx_side = self._launch_ctx_side( + draft_model, spec_metadata, attn_metadata, input_ids.shape[0] + ) + + self._execute_guided_decoder_if_present(logits) + + fused_accept = self._fused_select_accept(logits, attn_metadata, spec_metadata) + if fused_accept is not None: + accepted_tokens, num_accepted_tokens, bookkeeping_done = fused_accept + else: + bookkeeping_done = False + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + + # Opt-in acceptance recording (env-gated; eager-mode measurement + # runs only). Skipped for CUDA-graph batches (capture/replay/warmup + # use synthetic requests and forbid the host sync). + if ( + self._accept_stats is not None + and num_gens > 0 + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + ): + self._accept_stats.on_accept( + spec_metadata.request_ids[num_contexts:batch_size], + num_accepted_tokens[num_contexts:batch_size].tolist(), + ) + + # Update GDN/Mamba recurrent states to the accepted token's state. + if ( + not bookkeeping_done + and num_gens > 0 + and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager) + ): + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + if not bookkeeping_done: + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) + + # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. + # The draft model uses standard 1D RoPE, so only scalar positions are needed. + if position_ids.ndim == 3: + position_ids = position_ids[0, 0] + else: + position_ids = position_ids.squeeze(0) + + # Get total tokens processed by target model (for hidden state extraction) + total_target_tokens = input_ids.shape[0] + + # Capture prefill (context) hidden states for future gen steps. + # This gives the draft model the full prompt context, not just gen tokens. + if num_contexts > 0: + self._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens + ) + # Rebuild batch_to_slot after prefill assigns new slots + if self._ctx_buf_inited and spec_metadata.request_ids: + self._rebuild_batch_to_slot(spec_metadata.request_ids) + + # A context-KV-only drafter (disaggregated context server) stores the + # prompt's drafter context above and never drafts. Main has no + # speculative_config.context_kv_only yet, so a drafter reports False. + drafts = not getattr(draft_model, "context_kv_only", False) + if not drafts and num_gens > 0: + self._check_no_real_generation( + spec_metadata, num_contexts, batch_size, attn_metadata.kv_cache_manager + ) + inputs = ( + self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=total_target_tokens, + ) + if drafts + else None + ) + + vocab_size = None + if num_gens > 0 and drafts: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + # The fused preparations place the block in the pages: main's + # as a DraftBlockPlacement, TLLM_SPEC_TAIL_FUSION's as the MLA + # block fixup operands. + fixup_kwargs = ( + {"block_fixup": inputs["mla_block_fixup"]} + if inputs.get("mla_block_fixup") is not None + else {} + ) + hidden_states_out = draft_model.dflash_forward( + noise_embedding=inputs["noise_embedding"], + query_positions=inputs["query_positions"], + num_ctx_per_req=inputs["num_ctx_per_req"], + ctx_k_cache=inputs["ctx_k_cache"], + ctx_v_cache=inputs["ctx_v_cache"], + ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], + ctx_kv_cache=inputs["ctx_kv_cache"], + ctx_page_table=inputs["ctx_page_table"], + block_placement=inputs["block_placement"], + **fixup_kwargs, + ) + + # Gather K logits per gen request from the block outputs. + # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. + # Which block slots carry them is a drafter-family convention, + # resolved through _draft_slot_ids. + block_size = self._compute_block_size + gen_gather_ids = self._block_gather_ids( + draft_model, num_gens, block_size, K, hidden_states_out.shape[0] + ) + + # Jonas's draft-token selection through the private head. + draft_head = self._draft_vocab_parallel_head(draft_model, spec_metadata) + if draft_head is not None: + gen_draft_tokens, gen_hidden_states = self._k3_vocab_parallel_draft_tokens( + draft_head, + draft_model, + hidden_states_out, + gen_gather_ids, + num_gens, + K, + inputs, + ) + vocab_size = draft_head.vocab_size + else: + gen_hidden_states = hidden_states_out[gen_gather_ids] + gen_logits = draft_model.logits_processor( + gen_hidden_states, draft_model.lm_head, attn_metadata, True + ) + + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + + gen_logits = self._refine_block_logits( + draft_model, gen_logits, inputs, spec_metadata + ) + + # DFlash 2: replace the independent per-position picks with + # one coherent path through the block (absent for plain + # DFlash). + if getattr(draft_model, "lilicorr", None) is not None: + gen_logits = self._apply_lilicorr( + draft_model, + gen_logits, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["lilicorr_anchor"], + spec_metadata, + ) + vocab_size = gen_logits.shape[-1] + elif getattr(draft_model, "has_candidate_selector", False): + gen_logits = self._apply_dflash2_selector( + draft_model, + gen_logits, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + spec_metadata, + ) + vocab_size = gen_logits.shape[-1] + + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + ) + + # Opt-in confidence-calibration recording (env-gated). The + # confidence values come from an optional provider on the + # recorder (None here -> no calibration rows collected; the + # DSpark confidence-scheduled verification MR supplies the + # real provider, see accept_stats.py). Guarded off for + # CUDA-graph batches and d2t vocab-mapped drafters. + if ( + self._accept_stats is not None + and self._accept_stats.confidence_provider is not None + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + and self._d2t is None + ): + self._accept_stats.record_draft_confidence( + spec_metadata.request_ids[num_contexts:batch_size], + draft_model, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + gen_draft_tokens, + ) + + else: + # Warmup and attention-DP padding rows of a context-KV-only drafter + # get the same zero placeholder as context requests. + gen_draft_tokens = torch.zeros((num_gens, K), dtype=torch.int32, device="cuda") + + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a one-hot placeholder so + # they are a legal distribution when they become gen requests next iter. + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, vocab_size) + + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + else: + next_draft_tokens = gen_draft_tokens + + self._restore_attn_metadata_from_spec_dec(attn_metadata) + next_new_tokens = None + fused_rewind = self._fused_kv_rewind_operands(attn_metadata, spec_metadata, batch_size) + if ( + fused_rewind is not None + and next_draft_tokens.shape[0] == batch_size + and next_draft_tokens.is_contiguous() + and accepted_tokens.stride(-1) == 1 + ): + # The rewind rides in the next-new-tokens kernel. + self._kv_rewind_pending = False + next_new_tokens = k3_next_new_tokens( + accepted_tokens[:batch_size], + next_draft_tokens, + num_accepted_tokens[:batch_size], + *fused_rewind, + ) + else: + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + self._rollback_guided_decoder_after_verify(num_accepted_tokens) + + if next_new_tokens is None: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + # Restore context lengths after warmup; real runs keep the updates. + if is_warmup: + self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_host() + self._ctx_len_restore_pending = False + + return { + # The gathered float32 target logits, or this rank's vocab shard + # when the target tokens were selected from it. + "logits": raw_logits, + "logits_vocab_shard": logits_vocab_shard, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def _launch_ctx_side(self, draft_model, spec_metadata, attn_metadata, total_target_tokens): + """``TLLM_DSPARK_TAIL_V2``: start the context-latent chain on a side stream. + + The step's context latents (TP-sharded fc + all-reduce with hidden_norm, + the all-layer latent GEMM, RMSNorm + RoPE at ``ctx_len[slot] + j``) only + need the captured target hidden states, so they run concurrently with + the target-token selection and acceptance. The all-reduce uses a private + MNNVL workspace (the main stream's all-reduces run meanwhile). Same + kernels and inputs as the main-stream path, so the same values. + + Returns ``(latents [L, N, D], positions-read event, done event)``, or None + when the path does not apply (then the main stream computes them). + """ + if ( + not tail_v2_enabled() + or draft_model is None + or getattr(draft_model, "context_kv_only", False) + ): + return None + num_gens = attn_metadata.num_seqs + if attn_metadata.num_contexts != 0 or num_gens == 0 or num_gens > TAIL_V2_MAX_GENS: + return None + if not self._fused_gen_drafter_inputs_applies(draft_model, spec_metadata): + return None + # A TP-sharded drafter fc publishes its shard as (weight, begin, end, + # all_reduce); without one the main stream computes the latents. + shard = getattr(draft_model, "_fc_shard", None) + layer_major = getattr(draft_model, "precompute_context_kv_layer_major", None) + if not shard or layer_major is None: + return None + weight, begin, end, all_reduce = shard + if all_reduce.mnnvl_allreduce is None or draft_model.hidden_norm.eps is None: + return None + width = self._draft_tokens_per_req + tokens = num_gens * width + hidden = weight.shape[0] + ar = self._ctx_side_ar + if ar is None or not ar.fits(tokens, hidden): + if torch.cuda.is_current_stream_capturing(): + return None + from .dflash import PrivateMNNVLAllReduce + + # Collective: every rank reaches this at the same warmup step. + ar = PrivateMNNVLAllReduce(all_reduce.mapping, weight.dtype, max(tokens, 256), hidden) + self._ctx_side_ar = ar + self._ctx_side_stream = torch.cuda.Stream() + main = torch.cuda.current_stream() + side = self._ctx_side_stream + fork = torch.cuda.Event() + fork.record(main) + side.wait_event(fork) + positions_read = torch.cuda.Event() + done = torch.cuda.Event() + with torch.cuda.stream(side): + positions = k3_dflash_ctx_positions( + self._ctx_len, self._batch_to_slot_buffer.tensor[:num_gens], width + ) + positions_read.record(side) + captured = spec_metadata.get_hidden_states(total_target_tokens)[:tokens] + captured = captured.to(draft_model.fc.weight.dtype) + partial = F.linear(captured[:, begin:end], weight) + projected = ar.rms_norm( + partial, draft_model.hidden_norm.weight, draft_model.hidden_norm.eps + ) + projected = projected.to(self._ctx_latent_dtype(draft_model)) + latent = layer_major(projected, positions) + done.record(side) + if latent is None: + main.wait_event(done) + return None + # Consumed on the main stream after it waits on ``done``. + latent.record_stream(main) + return latent, positions_read, done + + def _fused_select_accept(self, logits, attn_metadata, spec_metadata): + """``TLLM_DSPARK_TAIL_V2``: target selection, strict acceptance and the + drafter's KV-length prep in one kernel. + + Applies to vocab-parallel target selection through the private head + with strict, uniform-width acceptance. + It writes what ``sample_and_accept_draft_tokens`` and + ``_prepare_kv_for_draft_forward`` write, after + ``_prepare_attn_metadata_for_dflash`` (moved ahead: it only saves + metadata tensors the acceptance does not touch). Returns ``(accepted, + num_accepted, bookkeeping_done)``, or None to take the unfused path. + + With a mamba-hybrid target the kernel also writes what + ``update_mamba_states`` writes when that is only the KDA replay record + (:meth:`_kda_replay_operands`). Otherwise it fuses only the selection + and the acceptance (``bookkeeping_done`` False), and + ``update_mamba_states`` and the KV-length prep run as on the unfused + path, in the same order. + """ + self._fused_kv_rewind_armed = False + if not tail_v2_enabled() or not isinstance(logits, K3VocabShardedLogits): + return None + if logits.head.exchange.backend != "mnnvl" and logits.head.exchange.tp_size > 1: + return None + # Main's ragged verification packs per-request windows; the kernel + # reads a uniform K + 1 tokens per generation request. + if spec_metadata.is_ragged_verify: + return None + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + if num_gens > TAIL_V2_MAX_GENS: + return None + if num_gens > 0 and self._can_use_rejection_sampling(spec_metadata): + return None + draft_tokens = self._reshape_draft_tokens_for_accept(spec_metadata, num_gens, logits.device) + pairs = logits.select_pairs(spec_metadata) + force = self._force_accept_kernel_args(draft_tokens.shape[-1], spec_metadata, logits.device) + kda_replay = None + if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): + fusable, kda_replay = self._kda_replay_operands(attn_metadata) + if not fusable: + accepted, num_accepted, _ = k3_fused_select_accept( + pairs, draft_tokens, num_contexts, batch_size, force + ) + return accepted, num_accepted, False + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + kv_lens = getattr(attn_metadata, "kv_lens_cuda", None) + accepted, num_accepted, rewind = k3_fused_select_accept( + pairs, + draft_tokens, + num_contexts, + batch_size, + force, + kda_replay=kda_replay, + kv_lens=kv_lens, + ) + if kv_lens is not None: + self._kv_rewind_amount = rewind + self._kv_rewind_nc = num_contexts + self._kv_rewind_bs = batch_size + self._kv_rewind_pending = True + attn_metadata.update_for_spec_dec() + self._fused_kv_rewind_armed = kv_lens is not None + return accepted, num_accepted, True + + @staticmethod + def _kda_replay_operands(attn_metadata): + """``(fusable, kda_replay)`` for :func:`fused_select_accept` with a + mamba-hybrid target. + + Fusable when ``update_mamba_states(attn_metadata, num_accepted, + state_indices=mamba_metadata.state_indices)`` is + ``KimiK3HybridCacheManager``'s and writes nothing but + ``KDAReplayState.record_acceptance``. ``kda_replay`` is then + ``(prev_num_accepted_tokens, state_indices, dummy_mask)``, or None when + that call records nothing (a context-only KDA manager, or no local KDA + layers). + """ + mgr = attn_metadata.kv_cache_manager + if ( + not isinstance(mgr, KimiK3HybridCacheManager) + or type(mgr).update_mamba_states is not KimiK3HybridCacheManager.update_mamba_states + ): + return False, None + if mgr.kda_context_only or mgr.local_num_mamba_layers == 0: + return True, None + replay = mgr._speculative_state + if ( + type(replay) is not KDAReplayState + or replay.prev_num_accepted_tokens is None + or mgr._dummy_request_mask is None + ): + return False, None + state_indices = attn_metadata.mamba_metadata.state_indices + if state_indices is None: + state_indices = mgr.get_state_indices() + dummy_mask = mgr._dummy_request_mask + if state_indices.stride(-1) != 1 or dummy_mask.stride(-1) != 1: + return False, None + return True, (replay.prev_num_accepted_tokens, state_indices, dummy_mask) + + def _fused_kv_rewind_operands(self, attn_metadata, spec_metadata, batch_size: int): + """``(kv_lens, amounts, num_contexts)`` when the rewind after the draft + forward runs inside the next-new-tokens kernel, else None (the + ``_apply_kv_rewind_after_draft`` path, which also covers warmup).""" + if not self._fused_kv_rewind_armed: + return None + self._fused_kv_rewind_armed = False + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if ( + is_warmup + or self._kv_rewind_amount is None + or not hasattr(attn_metadata, "kv_lens_cuda") + or self._kv_rewind_bs != batch_size + or not spec_tail_fusion_enabled() + ): + return None + return attn_metadata.kv_lens_cuda, self._kv_rewind_amount, self._kv_rewind_nc + + def _force_accept_kernel_args( + self, runtime_draft_len: int, spec_metadata, device: torch.device + ) -> Optional[tuple]: + """The ``force`` argument of ``fused_strict_accept`` / ``fused_select_accept`` + that reproduces :meth:`_apply_force_accepted_tokens` (same warmup skip + and draws).""" + if self.force_num_accepted_tokens == 0.0: + return None + if ( + spec_metadata is not None + and spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + ): + return None + int_part = int(self.force_num_accepted_tokens) + frac_part = self.force_num_accepted_tokens - int_part + max_total = runtime_draft_len + 1 + base_total = min(int_part + 1, max_total) + if frac_part > 0.0 and base_total < max_total: + self._ensure_force_accept_rng_state(device) + return ( + base_total, + frac_part, + max_total, + self._force_accept_rng_pool, + self._force_accept_rng_counter, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + ) + return (base_total,) + + def _sample_and_accept_draft_tokens_base( + self, + logits, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + spec_metadata, + ): + """``TLLM_SPEC_TAIL_FUSION=1``: main's strict acceptance and the force + override in one launch (``fused_strict_accept``), the same values. + + Unset, or under main's ragged verification, main's implementation runs. + """ + if not spec_tail_fusion_enabled() or ( + spec_metadata is not None and spec_metadata.is_ragged_verify + ): + return super()._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata + ) + if isinstance(logits, torch.Tensor) and logits.dim() == 1: + logits = logits.unsqueeze(0) + target_tokens = self._sample_tokens_for_batch( + logits, spec_metadata, num_contexts, batch_size + ) + return k3_fused_strict_accept( + target_tokens, + draft_tokens, + num_contexts, + batch_size, + self._force_accept_kernel_args( + draft_tokens.shape[-1], spec_metadata, target_tokens.device + ), + ) + + def _sample_tokens_for_batch( + self, + logits, + spec_metadata, + num_contexts: int, + batch_size: int, + ) -> torch.Tensor: + """Main's sampler, which also takes the private head's shard logits. + + The worker only passes shard logits through when vocab-parallel + selection covers this batch's sampling mode. + """ + if isinstance(logits, K3VocabShardedLogits): + return logits.select_tokens(spec_metadata) + return super()._sample_tokens_for_batch(logits, spec_metadata, num_contexts, batch_size) + + def _block_gather_ids( + self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int, num_rows: int + ) -> torch.Tensor: + """The block-output rows gathered for the K draft logits per request. + + Shape-only, so with ``TLLM_SPEC_TAIL_FUSION=1`` it is built once per + shape and reused instead of being recomputed every step. + """ + key = (num_gens, block_size, num_draft_tokens, num_rows) + if spec_tail_fusion_enabled(): + cache = self.__dict__.setdefault("_block_gather_id_cache", {}) + if key in cache: + return cache[key] + ids = self._draft_slot_ids(draft_model, num_gens, block_size, num_draft_tokens) + # Shields only the last request: at block_size == K with shift_label + # off, slots run 1..K, so every request reads the next one's slot 0 + # and the last overruns. Degrades, never raises. + ids = ids.clamp(max=num_rows - 1) + if spec_tail_fusion_enabled() and not torch.cuda.is_current_stream_capturing(): + self._block_gather_id_cache[key] = ids + return ids + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + # batch_indices_cuda is the identity here (DFlashSpecMetadata.prepare). + if ( + spec_tail_fusion_enabled() + and next_draft_tokens.shape[0] == batch_size + and next_draft_tokens.is_contiguous() + and accepted_tokens.stride(-1) == 1 + ): + return k3_next_new_tokens( + accepted_tokens[:batch_size], next_draft_tokens, num_accepted_tokens[:batch_size] + ) + return super()._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens + ) + + def _draft_first_slot(self, draft_model) -> int: + """The block slot whose hidden state predicts the first draft token + (``_draft_slot_ids`` takes slots ``first .. first + K - 1`` of each block).""" + return 1 + + def _block_gather_range( + self, draft_model, num_gens: int, num_draft_tokens: int, num_rows: int + ) -> Optional[tuple]: + """``(start, stop)`` when the gathered block rows are one contiguous range + (one request), else None.""" + if num_gens != 1: + return None + first = self._draft_first_slot(draft_model) + stop = first + num_draft_tokens + return (first, stop) if stop <= num_rows else None + + def _draft_block_width(self, draft_model) -> int: + """Block slots the draft forward must compute for K draft tokens. + + Plain DFlash reads slots 1..K, so slot 0 is pure overhead and the + forward needs K+1. Families whose slot convention differs override + this alongside :meth:`_draft_slot_ids` — the two must agree, since + the width bounds the slots the gather is allowed to name. + """ + return self.max_draft_len + 1 + + def _draft_slot_ids( + self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int + ) -> torch.Tensor: + """Block-output slots whose hidden states produce the K draft logits. + + Plain DFlash uses the K2.7 convention: mask slots 1..K. Drafter + families with another convention override this — see + :meth:`K3DSparkWorker._draft_slot_ids` for the shift_label + variant, which reads slots 0..K-1 instead. + """ + return dflash_draft_slot_ids(num_gens, block_size, num_draft_tokens, False, device="cuda") + + def _refine_block_logits( + self, + draft_model, + gen_logits: torch.Tensor, + inputs: dict, + spec_metadata, + ) -> torch.Tensor: + """Refine the block logits between the draft forward and sampling. + + Plain DFlash proposes the backbone's logits unchanged. Drafter + families carrying extra heads override this — see + :meth:`K3DSparkWorker._refine_block_logits` for the Markov + intra-block bias. + """ + return gen_logits + + def _apply_dflash2_selector( + self, + draft_model, + gen_logits: torch.Tensor, + block_hidden_states: torch.Tensor, + anchor_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Replace DFlash's independent per-position picks with the DFlash 2 + candidate selector's chosen path. + + Returns full-vocab ``[num_gens, K, vocab]`` logits carrying only each + position's top-k candidates, scored under the predecessor the walk + settled on. Downstream sampling is unchanged and stays lossless: each + row is the distribution its draft token was drawn from. + + The walk always chains through the greedy successor, so under sampling + a draft token can differ from the predecessor the next position was + scored under. That costs acceptance rate, not correctness — same + trade-off as ``_apply_dspark_markov_bias``. + """ + selector = draft_model.candidate_selector + candidate_ids, unary_logits, block_logits = self._dflash2_global_top_k( + gen_logits, + spec_metadata, + top_k=selector.top_k, + full_vocab=selector.predecessor_codebook.shape[0], + ) + return draft_model.select_candidate_path( + block_logits, + candidate_ids, + unary_logits, + block_hidden_states, + anchor_tokens.long(), + ) + + def _apply_lilicorr( + self, + draft_model, + gen_logits: torch.Tensor, + draft_hidden: torch.Tensor, + anchor_hidden: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + if anchor_hidden is None: + raise RuntimeError( + "LiLiCorr requires the projected target row that predicted the anchor" + ) + full_vocab = draft_model.config.vocab_size + candidate_ids, unary_logits, block_logits = self._dflash2_global_top_k( + gen_logits, spec_metadata, draft_model.lilicorr.candidate_topk, full_vocab + ) + # Candidate features use probabilities normalized over the entire vocabulary, + # including the mass outside top-k and on other tensor-parallel ranks. + log_normalizer = torch.logsumexp(gen_logits.float(), dim=-1, keepdim=True) + if gen_logits.shape[-1] != full_vocab: + from tensorrt_llm._torch.distributed.ops import allgather + + log_normalizer = torch.logsumexp( + allgather(log_normalizer, self.mapping, dim=-1), dim=-1, keepdim=True + ) + proposal = draft_model.select_lilicorr_path( + candidate_ids, unary_logits.float() - log_normalizer, draft_hidden, anchor_hidden + ) + block_logits.fill_(float("-inf")) + block_logits.scatter_(-1, candidate_ids, proposal.to(block_logits.dtype)) + return block_logits + + def _dflash2_global_top_k( + self, + gen_logits: torch.Tensor, + spec_metadata, + top_k: int, + full_vocab: int, + ): + """Per-position global top-k candidates, and the tensor to score into. + + The selector indexes full-vocab codebooks, so candidate ids must be + global. A vocab-sharded draft head is reduced by all-gathering only each + rank's local (id, value) pairs — 2 * top_k per position instead of the + whole vocabulary. Every rank then writes identical rows, so downstream + sampling needs no further collective. + + Returns ``(candidate_ids, unary_logits, block_logits)``, where + ``block_logits`` is ``gen_logits`` rewritten in place when it is already + full-vocab (the common single-rank case) and a fresh tensor otherwise. + """ + shard = gen_logits.shape[-1] + if shard == full_vocab: + unary_logits, candidate_ids = torch.topk(gen_logits, top_k, dim=-1) + return candidate_ids, unary_logits, gen_logits + + mapping = self.mapping + if not ( + self._draft_logits_are_sharded(gen_logits, spec_metadata) + and mapping is not None + and shard * mapping.tp_size == full_vocab + ): + raise NotImplementedError( + f"DFlash 2 candidate selection needs the drafter's full vocab " + f"({full_vocab}) or a plain TP column shard of it; got draft " + f"logits of width {shard}." + ) + + from tensorrt_llm._torch.distributed.ops import allgather + + local_values, local_ids = torch.topk(gen_logits, top_k, dim=-1) + local_ids = local_ids + mapping.tp_rank * shard + # Interleaved (id, value) pairs, matching _get_local_max_and_combined. + combined = torch.stack([local_ids.float(), local_values.float()], dim=-1).flatten(-2) + gathered = allgather(combined, mapping, dim=-1) + unary_logits, selected = torch.topk(gathered[..., 1::2], top_k, dim=-1) + candidate_ids = gathered[..., 0::2].gather(-1, selected).long() + block_logits = gen_logits.new_full((*gen_logits.shape[:-1], full_vocab), float("-inf")) + return candidate_ids, unary_logits, block_logits + + def prepare_1st_drafter_inputs( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DFlashSpecMetadata, + draft_model: nn.Module, + total_target_tokens: int = 0, + ): + """Prepare inputs for DFlash's draft forward. + + For gen requests, builds: + - noise_embedding: token embeddings for [accepted + mask] tokens + - query_positions: position IDs for the query tokens + - num_ctx_per_req: per-request context length in the pool + - ctx_k_cache / ctx_v_cache / ctx_cache_batch_idx: slot-indexed + views of the persistent per-layer K/V pool. + """ + num_contexts = attn_metadata.num_contexts + batch_size = attn_metadata.num_seqs + num_gens = batch_size - num_contexts + + # Resolve mask_token_id and block_size once, cache for subsequent calls + if self._resolved_mask_token_id is None: + if ( + hasattr(self.spec_config, "mask_token_id") + and self.spec_config.mask_token_id is not None + ): + self._resolved_mask_token_id = self.spec_config.mask_token_id + elif hasattr(draft_model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.mask_token_id + elif hasattr(draft_model.model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.model.mask_token_id + else: + raise ValueError( + "DFlash requires mask_token_id to be set. Please set it in DFlashDecodingConfig " + "or ensure the draft model config has 'dflash_config.mask_token_id' or 'mask_token_id'." + ) + mask_token_id = self._resolved_mask_token_id + + # Get the embed_tokens layer from the draft model + embed_tokens = draft_model.draft_model_full.model.embed_tokens + hidden_dim = ( + spec_metadata.hidden_size if spec_metadata.hidden_size > 0 else hidden_states.shape[-1] + ) + lilicorr_anchor = None + # The drafter's own trained mask row, if it kept one for the mask id in use. + trained_mask_embedding = ( + getattr(draft_model, "mask_token_embedding", None) + if mask_token_id == getattr(draft_model, "mask_token_id", None) + else None + ) + + # One fused preparation at most: TLLM_SPEC_TAIL_FUSION=1 selects + # Jonas's, unset main's; each falls back to the eager path below. + if num_gens > 0: + if spec_tail_fusion_enabled(): + if self._fused_gen_drafter_inputs_applies(draft_model, spec_metadata): + return self._prepare_gen_drafter_inputs_fused( + accepted_tokens, + num_accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + embed_tokens, + total_target_tokens, + mask_token_id, + trained_mask_embedding, + ) + elif self._fused_gen_prepare_applies(spec_metadata): + return self._prepare_paged_latent_gen_inputs( + accepted_tokens, + num_accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + total_target_tokens, + mask_token_id, + trained_mask_embedding, + ) + + if num_gens > 0: + gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] + gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] + + total_tokens_per_req = self._draft_tokens_per_req # K+1 + K = spec_metadata.runtime_draft_len + + # Get captured multi-layer hidden states from spec_metadata + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + has_target_features = ( + captured_hs is not None + and hasattr(draft_model, "fc") + and hasattr(draft_model, "hidden_norm") + ) + + # Use cached block_size (resolved once on first call) + block_size = self._compute_block_size + query_tokens_per_req = block_size + + # Get slots for gen requests from pre-computed mapping + slots = self._batch_to_slot_buffer.tensor[num_contexts : num_contexts + num_gens] + gen_rows_out = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + + K_plus_1 = K + 1 + + bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) + bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() + + ctx_len_gen = self._ctx_len[slots] + j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") + offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") + + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + + noise_embed_2d = dflash_noise_block_embedding( + embed_tokens, bonus, mask_token_id, query_tokens_per_req, trained_mask_embedding + ) + + # Accumulate new accepted features into context buffers + if has_target_features: + gen_start = attn_metadata.num_ctx_tokens + # Target now processes exactly K+1 tokens per gen req, so the + # captured slice is already the full set we need to project. + gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] + gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) + projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) + if getattr(draft_model, "lilicorr", None) is not None: + # The last accepted output is predicted by the corresponding input + # row; the newly committed token itself has no target row yet. + projected_rows = projected_to_store.reshape(num_gens, K_plus_1, -1) + lilicorr_anchor = last_accepted_hidden(projected_rows, gen_num_accepted) + gen_num_accepted_long = gen_num_accepted.long() + col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) + owned = None + shard = self._helix_shard() + if shard is not None: + # Helix: rank-local slots; the other ranks' positions + # rewrite their slot's current value (see + # _store_context_kv_paged). + assert K + 1 <= shard.page_size, "helix alias slots must lie outside the write" + col_idx, owned = shard.local_slots(col_idx) + owned = owned.reshape(-1) + if self._ctx_block_tables is not None: + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. + gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) + if self._ctx_block_tables is None: + # The private arena. + col_idx = col_idx.clamp(max=self._max_ctx - 1) + + # Fixed-size writes for CUDA graph compatibility: + # Write ALL entries but zero out invalid ones. Invalid + # writes land at clamped column indices beyond valid + # _ctx_len range, so they're harmless. + slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) + col_flat = col_idx.reshape(-1) + proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] + pos_flat = ctx_position_ids.long().reshape(-1) + mask_1d = write_mask.reshape(-1) + + # Fast path: store the pre-projected/pre-RoPE'd K/V. + # dflash_forward reads these directly via cache_batch_idx. + cache_dtype = ( + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + ) + if cache_dtype == torch.float8_e4m3fn: + # Project in the drafter's weight dtype; the store converts. + cache_dtype = proj_flat.dtype + k_new, v_new = draft_model.precompute_context_kv( + proj_flat.to(cache_dtype), pos_flat + ) + mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) + k_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) + slot_long = slot_flat.long() + col_long = col_flat.long() + if self._ctx_paged: + if self._ctx_block_tables is not None: + # Batch positions of the gen requests, matching the + # per-request block table's row order. + rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + else: + rows_long = slot_long + self._store_context_kv_paged(k_new, v_new, rows_long, col_long, owned) + else: # VANILLA DFlash backend (FlashAttention) + self._ctx_k_buf[slot_long, :, col_long] = k_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new + + self._advance_ctx_len(slots, gen_num_accepted_long) + + num_ctx_per_req_t = self._ctx_len[slots] + if self._ctx_block_tables is not None: + # The write above clamps columns to the request's allocation, so + # a context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. + # Under helix local page k holds global page k * cp_size + + # cp_rank and is the alias of every other page in that span, + # so k local pages cover k * cp_size global pages. + shard = self._helix_shard() + cp = 1 if shard is None else shard.cp_size + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts * cp, self._ctx_page_size, block_size + ) + if self._ctx_unified_cache: + self._check_block_within_allocation( + num_ctx_per_req_t, allocated[gen_rows_out], block_size + ) + num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + noise_embedding = noise_embed_2d + query_positions = query_position_ids.long() + + # Update seq_lens for gen requests to K+1 (the number of tokens + # the target forward actually processed). + attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = ( + total_tokens_per_req + ) + attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req + else: + noise_embedding = hidden_states.new_empty(0, 0, hidden_dim) + query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") + num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") + slots = torch.empty(0, dtype=torch.long, device="cuda") + gen_rows_out = slots + bonus = torch.empty(0, dtype=torch.long, device="cuda") + + return { + "lilicorr_anchor": lilicorr_anchor, + "noise_embedding": noise_embedding, + "query_positions": query_positions, + "num_ctx_per_req": num_ctx_per_req_t, + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + # Slots index the private arena; the manager's block table is keyed + # by batch position, so the drafter reads its own rows there. + "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, + # Anchor token per gen request (block slot 0): last accepted + # token. The dspark Markov chain conditions its first step on it. + "first_prev_tokens": bonus, + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ), + "block_placement": None, + } + + def _fused_gen_prepare_applies(self, spec_metadata) -> bool: + """Whether :meth:`_prepare_paged_latent_gen_inputs` covers this step: a + drafter it supports (``_paged_latent_prepare``), captured target + features, drafting from all ``K + 1`` target tokens.""" + return ( + self._paged_latent_prepare + # Shared allocations require the generic path's strict capacity + # check; the fused preparation silently caps the context length. + and not self._ctx_unified_cache + and spec_metadata.runtime_draft_len + 1 == self._draft_tokens_per_req + and spec_metadata.captured_hidden_states is not None + ) + + def _prepare_paged_latent_gen_inputs( + self, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DFlashSpecMetadata, + draft_model: nn.Module, + total_target_tokens: int, + mask_token_id: int, + trained_mask_embedding: Optional[torch.Tensor], + ) -> dict: + """:meth:`prepare_1st_drafter_inputs` for the generation requests of a + step (context requests may precede them) whose drafter context is a + paged single-latent cache. + + Same inputs and side effects (context latents stored, ``_ctx_len`` + advanced, ``seq_lens`` set) with the per-request bookkeeping in one + kernel and the latent store of every layer in another; also places the + block's own positions in the pages for the drafter's block decode. + """ + num_contexts = attn_metadata.num_contexts + num_gens = attn_metadata.num_seqs - num_contexts + width = self._draft_tokens_per_req + block_size = self._compute_block_size + prep = k3_dflash_paged_latent_prepare( + num_accepted_tokens, + accepted_tokens, + self._batch_to_slot_buffer.tensor[num_contexts : num_contexts + num_gens], + self._ctx_len, + self._ctx_block_counts, + self._ctx_block_tables, + attn_metadata._seq_lens_cuda, + num_contexts, + self._dummy_slot, + self._max_ctx, + self._ctx_page_size, + block_size, + width, + ) + attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = width + + noise_embedding = dflash_noise_block_embedding( + draft_model.draft_model_full.model.embed_tokens, + prep.bonus, + mask_token_id, + block_size, + trained_mask_embedding, + ) + gen_start = attn_metadata.num_ctx_tokens + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + gen_hs = captured_hs[gen_start : gen_start + num_gens * width] + projected = draft_model.project_target_hidden(gen_hs.reshape(-1, gen_hs.shape[-1])) + latent, _ = draft_model.precompute_context_kv( + projected.to(self._ctx_kv_buf[0].dtype), prep.ctx_positions + ) + k3_store_context_latents( + latent[:, :, 0], + self._ctx_kv_buf, + self._ctx_kv_ptrs, + prep.store_pages, + prep.store_offsets, + prep.store_mask, + mla=False, + ) + return { + "lilicorr_anchor": None, + "noise_embedding": noise_embedding, + "query_positions": prep.query_positions, + "num_ctx_per_req": prep.num_ctx, + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + "ctx_cache_batch_idx": torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ), + "first_prev_tokens": prep.bonus, + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": self._ctx_block_tables, + "block_placement": DraftBlockPlacement( + prep.blk_pages, prep.blk_offsets, prep.blk_seq_lens, prep.blk_kv_bounds + ), + } + + def _ctx_latent_dtype(self, draft_model) -> torch.dtype: + """Dtype the context latents are computed in: the pool's, or the + projection's activation dtype for an FP8 pool, whose stores quantize them.""" + dtype = self._ctx_kv_buf[0].dtype + return dflash_context_dtype(draft_model) if dtype == torch.float8_e4m3fn else dtype + + def _fused_gen_drafter_inputs_applies(self, draft_model, spec_metadata) -> bool: + """Whether ``_prepare_gen_drafter_inputs_fused`` (``TLLM_SPEC_TAIL_FUSION=1``) + covers this step: a generation batch drafting from a single-latent (MLA) + context that lives in the draft KV manager's paged pool. + + LiLiCorr drafters keep the eager path, which computes their anchor. + """ + return ( + spec_tail_fusion_enabled() + and self._ctx_paged + and self._ctx_block_tables is not None + and getattr(draft_model, "_kv_factor", 2) == 1 + and hasattr(draft_model, "fc") + and hasattr(draft_model, "hidden_norm") + and getattr(draft_model, "lilicorr", None) is None + # The fused latent store writes in the projection's dtype or FP8. + and self._ctx_kv_buf[0].dtype + in (dflash_context_dtype(draft_model), torch.float8_e4m3fn) + and spec_metadata.get_hidden_states(0) is not None + and spec_metadata.runtime_draft_len + 1 == self._draft_tokens_per_req + and self._compute_block_size is not None + ) + + def _ctx_kv_pool_ptrs(self) -> torch.Tensor: + """Device array of the per-layer context-pool base addresses (cached).""" + ptrs = tuple(buf.data_ptr() for buf in self._ctx_kv_buf) + if getattr(self, "_ctx_kv_ptrs_key", None) != ptrs: + self._ctx_kv_ptrs = torch.tensor(ptrs, dtype=torch.int64, device="cuda") + self._ctx_kv_ptrs_key = ptrs + return self._ctx_kv_ptrs + + def _fused_prepare_constants( + self, num_contexts: int, num_gens: int, block_size: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """The block fixup's upper-triangle mask and the gen rows' batch indices. + + Shape-only. With ``TLLM_DSPARK_TAIL_V2=1`` they are built once per shape + outside graph capture and reused, so no kernel recomputes them per step. + """ + key = (num_contexts, num_gens, block_size) + cache = self.__dict__.setdefault("_fused_prepare_constant_cache", {}) + if tail_v2_enabled() and key in cache: + return cache[key] + idx = torch.arange(block_size, device="cuda") + valid = (idx.view(1, block_size) > idx.view(block_size, 1)).view( + 1, block_size, 1, block_size + ) + batch_idx = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + if tail_v2_enabled() and not torch.cuda.is_current_stream_capturing(): + cache[key] = (valid, batch_idx) + return valid, batch_idx + + def _prepare_gen_drafter_inputs_fused( + self, + accepted_tokens, + num_accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + embed_tokens, + total_target_tokens: int, + mask_token_id: int, + trained_mask_embedding: Optional[torch.Tensor] = None, + ) -> dict: + """``prepare_1st_drafter_inputs`` for the paged single-latent context + (``TLLM_SPEC_TAIL_FUSION=1``; unset, main's + :meth:`_prepare_paged_latent_gen_inputs` runs instead). + + Produces the same drafter inputs and the same side effects (context + latents stored, ``_ctx_len`` advanced, ``seq_lens`` set) with the + per-request bookkeeping in one kernel, the noise embedding in one copy + kernel and the latent store in one scatter kernel. Also returns the + block-decode fixup operands, which ``dflash_forward`` would otherwise + rebuild from a gathered copy of the block table. + """ + from .drafter import _MLABlockFixup + + num_contexts = attn_metadata.num_contexts + num_gens = attn_metadata.num_seqs - num_contexts + width = self._draft_tokens_per_req + block_size = self._compute_block_size + slots = self._batch_to_slot_buffer.tensor[num_contexts : num_contexts + num_gens] + side = self._ctx_side + self._ctx_side = None + if side is not None: + # The side stream has read _ctx_len before prep advances it. + torch.cuda.current_stream().wait_event(side[1]) + shard = self._helix_shard() + prep = k3_dflash_gen_prepare( + num_accepted_tokens, + accepted_tokens, + slots, + self._ctx_len, + self._ctx_block_counts, + self._ctx_block_tables, + attn_metadata._seq_lens_cuda, + num_contexts, + num_gens, + self._dummy_slot, + self._max_ctx, + int(mask_token_id), + self._ctx_page_size, + block_size, + width, + cp_size=1 if shard is None else shard.cp_size, + cp_rank=0 if shard is None else shard.cp_rank, + ) + attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = width + if self._ctx_unified_cache: + # The kernel caps the advertised context so the block fits the + # allocation; a shared allocation must not need the cap. + self._check_block_within_allocation(self._ctx_len[slots], prep["num_ctx"], block_size) + + # A TP-sharded vocab masks its rows and all-reduces (embed_tokens.forward + # semantics; one lookup kernel when the embedding is vocab-sharded). + # Owner-rank multicast of each token's row instead of the all-reduce, + # where it applies. + embed_ids = prep["embed_ids"] + combined_embed = draft_embed_site(embed_tokens, embed_ids) + if combined_embed is None and vocab_shard_embedding_applies(embed_tokens, embed_ids): + local_rows = k3_vocab_shard_embedding( + embed_ids, embed_tokens.weight, embed_tokens.vocab_start, embed_tokens.vocab_end + ) + combined_embed = embed_tokens.all_reduce(local_rows) + if combined_embed is None: + combined_embed = embed_tokens(prep["embed_ids"]) + if trained_mask_embedding is not None: + # The drafter's own trained mask row replaces the shared lookup's + # (the last row), as in dflash_noise_block_embedding. + combined_embed = torch.cat( + ( + combined_embed[:num_gens], + trained_mask_embedding.to(combined_embed.dtype).unsqueeze(0), + ) + ) + noise_embedding = k3_dflash_noise_embedding(combined_embed, block_size) + + if side is not None: + # Computed on the side stream from the same captured hidden states + # and the same positions (ctx_len[slot] + j). + torch.cuda.current_stream().wait_event(side[2]) + latent = side[0].transpose(0, 1) + else: + gen_start = attn_metadata.num_ctx_tokens + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + gen_hs = captured_hs[gen_start : gen_start + num_gens * width] + projected = draft_model.project_target_hidden(gen_hs.reshape(-1, gen_hs.shape[-1])) + projected = projected.to(self._ctx_latent_dtype(draft_model)) + # Layer-major latents skip the per-layer staging copies when the + # drafter offers them; the values are the same either way. + layer_major = getattr(draft_model, "precompute_context_kv_layer_major", None) + latent = None if layer_major is None else layer_major(projected, prep["ctx_positions"]) + if latent is not None: + latent = latent.transpose(0, 1) + else: + latent = draft_model.precompute_context_kv(projected, prep["ctx_positions"])[0][ + :, :, 0 + ] + store_args = ( + latent, + self._ctx_kv_buf, + self._ctx_kv_pool_ptrs(), + prep["store_pages"], + prep["store_offsets"], + prep["store_mask"], + ) + side_launch = getattr(draft_model, "draft3_side_launch", None) + if draft3_glue_enabled() and side_launch is not None: + # The store runs beside the drafter's first projections; dflash_forward + # joins it before its first draft-pool write (the store's masked rows + # zero slots the block latents then take). + side_launch( + lambda: k3_store_context_latents(*store_args, mla=True), + keep=(latent, prep["store_pages"], prep["store_offsets"], prep["store_mask"]), + ) + else: + k3_store_context_latents(*store_args, mla=True) + + valid, batch_idx = self._fused_prepare_constants(num_contexts, num_gens, block_size) + block_fixup = _MLABlockFixup( + pages=prep["blk_pages"], + offsets=prep["blk_offsets"], + valid=valid, + page_tables_i32=self._ctx_block_tables[num_contexts : num_contexts + num_gens], + seq_lens_i32=prep["blk_seq_lens"], + kv_bounds_i32=prep["blk_kv_bounds"], + owned=None if shard is None else self._blk_owned_mask(prep["blk_owned"]), + ) + return { + "lilicorr_anchor": None, + "noise_embedding": noise_embedding, + "query_positions": prep["query_positions"], + "num_ctx_per_req": prep["num_ctx"], + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + "ctx_cache_batch_idx": batch_idx, + "first_prev_tokens": prep["bonus"], + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": self._ctx_block_tables, + "block_placement": None, + "mla_block_fixup": block_fixup, + } + + @staticmethod + def _blk_owned_mask(blk_owned: torch.Tensor) -> torch.Tensor: + """The ``[gens, block, 1]`` bool ownership mask of the prep kernel's 0 / 1 int8 rows: + reinterpreted in place under ``TRTLLM_K3_DRAFT3_GLUE`` (same bytes as the conversion), + else converted.""" + if draft3_glue_enabled() and blk_owned.dtype in (torch.int8, torch.uint8): + return blk_owned.view(torch.bool).unsqueeze(-1) + return blk_owned.bool().unsqueeze(-1) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/drafter.py new file mode 100644 index 000000000000..34ebb1d98558 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/drafter.py @@ -0,0 +1,3159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi-K3 (Rubin) standalone MLA DSpark drafter: K3-owned weights and an explicit kernel chain. + +The drafter is ``K3MLADSparkForCausalLM``, a ``DFlashForCausalLM`` (the engine recognises the +speculative model by that type and runs the DFlash worker contract on it). Below that base the +backbone is this package's: the weights live in ``layers.py`` containers +(``Bf16Linear``, ``RMSNormWeight``, ``GatedMlp``, ``EmbeddingTable``, ``TpAllReduce``), loaded by +:meth:`K3MLADSparkForCausalLM.load_weights`, and each step is the kernel chain +:meth:`K3MLADSparkForCausalLM.dflash_forward` spells out: + + noise embedding (worker) vocab-sharded lookup + all-reduce + per layer: + input RMSNorm (+ residual) FlashInfer; under TP the previous layer's + down all-reduce + residual + RMSNorm + [q_a | kv_a] GEMM (or q_a, kv_a) cuBLAS / pinned split-K + q_a RMSNorm, q_b GEMM FlashInfer, cuBLAS / pinned split-K + query RoPE, q_nope @ W_UK -> latent query cute_dsl_dspark_rope_into, bmm_out + block latent RMSNorm + RoPE, store into pool k3_dspark_rmsnorm_rope, scatter + (dspark_tail_v2: both RoPEs and the store k3_dspark_mla_block_rope, bmm_out) + split-KV plan (first layer of a step) the MLA decode's planner kernel + latent attention over context + block cute_dsl_mla_decode_fp16_blackwell + helix: merge of the rank partials _helix_post_process / k3_helix_tail + V up-projection bmm_out (or einsum) + o_proj, all-reduce + residual + RMSNorm pushed split-K site, or cuBLAS + MNNVL + gate_up GEMM, SwiGLU, down GEMM cuBLAS / pinned split-K, silu_and_mul + final RMSNorm (+ the last down all-reduce) + +and the context latents of the target hidden states are fc (+ all-reduce with the RMSNorm fused) +-> one GEMM over the stacked ``kv_a`` weights of every layer -> ``k3_dspark_mla_ctx_norm_rope``. + +Features (``features.py``): ``dspark_draft_gemm`` pins GEMMs, ``draft3_glue`` takes +per-step glue off the critical path, ``spec_tail_fusion`` / ``dspark_tail_v2`` fuse the block +stores and RoPE, ``ull2_draft_a2a`` / ``ull2_draft_store`` the helix page-sharded latent cache. The +TP fusions (q_a / kv_a as one GEMM, the pushed all-reduce sites, the TP-sharded fc) change the +drafter's rounding only; target verification keeps the output lossless. + +No registry entries: main's registry keeps resolving its own drafter classes (the backbone is +registered only while the base ``__init__`` builds it, under a private name); the K3 shell builds +this drafter through :func:`build_k3_dspark_draft`. + +Standalone DSpark worker, Kimi K3 Rubin private copy. + +The 3-way merge of ``DSparkWorker`` in ``tensorrt_llm/_torch/speculative/dspark.py`` +(main) with Jonas Li's ``user/jonasl/k3-rubin-agentic-opt`` (189fe67b81). Only the +standalone worker (DFlash lineage) is here: the embedded DeepSeek-V4 worker +(``DSv4DSparkWorker``), its ``DSparkSpecMetadata`` and helpers stay main's. + +Jonas's additions: the Markov chain through the private vocab-parallel head +(see :meth:`K3DSparkWorker._draft_vocab_parallel_head`), and his fused +full-vocab chain in place of the drafter's ``apply_markov_chain_logits`` where +the gathered logits fit it. + +Resolution of the DSpark drafter head switches from a drafter's pretrained config. + +DSpark vanilla Markov head: the per-step token-bigram logit bias and the intra-block chain. + +The MLA DSpark drafter's helix layer on the target's exchange, with fewer launches. + +Under helix the MLA drafter keeps one latent per token, so its context pages are sharded over the +whole repurposed group (``Mapping.helix_mapping_for_kv_heads(1)``): every rank attends all heads +over its pages and the per-rank partials are merged by a helix exchange. Switches (environment, +default off; each falls back to the current path where it does not apply): + +* ``TRTLLM_K3_ULL2_DRAFT_A2A=1``: the drafter's helix mapping takes the cp_config of the target's + helix mapping (``inherit_cp_config``), so its exchange runs on the target's transport, e.g. the + MNNVL FIFO all-to-all + ``helix_post_process_native`` for ``use_nccl_for_alltoall: false``, + instead of the NCCL all-to-all with its torch reformats. Same values: both combines run the + same arithmetic on the same records. With an MNNVL cp_config, the V up-projection, the + exchange and the combine of a layer are then one launch of the Kimi-K3 helix tail kernel + (``k3_helix_tail``, rows copied, no fold), on the receive buffers of the target's tail when the + mappings match. The V up-projection sums in another FP32 order than the cuBLAS BMM: drafter + rounding only. +* ``TRTLLM_K3_ULL2_DRAFT_STORE=1`` (with ``TLLM_DSPARK_TAIL_V2=1``): under helix too, one kernel + rotates the query's rope half, normalizes and rotates the block latents into a scratch row + buffer, and one masked scatter stores the rows this rank's pages hold + (``scatter_rows_to_pool(rows_mask=owned)``), instead of the separate RoPE launches and a + gather / where / index_put round trip of the pool. Same values. + +Helix CP post-processing: all-to-all exchange and combine of partial attention outputs. + +Holds ``_helix_post_process`` with the zero-local-KV sanitize and the compiled NCCL +reformat helpers around the all-to-all. + +Speculative shell of the Kimi K3 Rubin targets: the private DSpark drafter and worker. + +The K3 text model (``k3_rubin/text_model.py``, ``KimiK3DefaultForCausalLM``) is built by +main's one-engine shell, ``SpecDecOneEngineForCausalLM.__init__``, which builds main's +DSpark drafter through the draft-model registry (``get_draft_model`` -> +``_build_dspark_draft``) and main's ``DSparkWorker`` through ``get_spec_worker``. +:func:`install_k3_dspark` then replaces both with the package's own: the drafter is +``modeling_dspark_mla.K3MLADSparkForCausalLM`` (built by ``build_k3_dspark_draft``, no +registry entry) and the worker ``dspark.K3DSparkWorker`` on ``dflash.K3DFlashWorker``, which +carry Jonas Li's Rubin DSpark fusions (``user/jonasl/k3-rubin-agentic-opt``) merged onto +main's engine contract. + +Main's drafter is constructed and discarded. The model loader constructs models under +``MetaInitMode`` (``pyexecutor/model_loader.py``), so that costs no device memory; only +the loader's non-meta fallback allocates it once, and it is freed with the old modules. + +Only one-engine DSpark with a standalone MLA drafter is replaced. SA, DFlash, a GQA +DSpark drafter and a draft embedded in the target keep main's drafter and worker, so the +target stays identical to main for them. + +Contract with the text model (``KimiK3DefaultForCausalLM._install_private_spec``): call +``install_k3_dspark(self, model_config)`` once, right after +``SpecDecOneEngineForCausalLM.__init__`` and before the helix CP mapping is restored, so +the private drafter sees the mapping main's drafter was built under. It swaps +``draft_model`` / ``spec_worker`` in ``self.epilogue`` too, calls +``set_draft_model`` / ``set_target_model(self)`` on the new worker, leaves ``self.lm_head`` +and ``self.logits_processor`` untouched (main's vocab-parallel selection requires an +``LMHead`` and exactly a ``LogitsProcessor``), and the private worker keeps the attribute +contract main's ``DFlashSpecMetadata.prepare()`` reads (``_ctx_buf_inited``, +``_req_to_slot``, ``_req_ctx_pos``, ``_free_slots``, ``_graph_dummy_id_floor``, +``_dummy_slot``, ``_batch_to_slot_buffer``) and calls (``_assign_slot(rid, updates=, bases=)``, +``_seed_generation_slot(slot, position, updates, bases)``, ``_write_ctx_len(updates, bases)``). +""" + +import contextlib +import inspect +import math +from functools import lru_cache +from typing import Callable, Dict, Iterator, List, NamedTuple, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.bmm_out import bmm_out +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.distributed import HelixAllToAllNative, alltoall_helix +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.pyexecutor.config_utils import is_mla +from tensorrt_llm._torch.speculative.vocab_parallel import markov_chain_logits +from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm.functional import ( + AllReduceFusionOp, + AllReduceParams, + AllReduceStrategy, + RopeEmbeddingUtils, +) +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import CpType, Mapping + +from .dflash import ( + DraftArSites, + HelixDraftShard, + K3DFlashWorker, + dflash_draft_slot_ids, + draft3_glue_enabled, + draft_gemm, +) +from .features import active as _features +from .kernels.spec.dspark_markov_fused import fused_markov_supported +from .kernels.spec.spec_tail_kernels import ( + DraftBlockPlacement, + spec_tail_fusion_enabled, + tail_v2_enabled, +) +from .layers import ( + Bf16Linear, + EmbeddingTable, + GatedMlp, + RMSNormWeight, + TensorParallelMode, + TpAllReduce, + bf16_linear, + gated_activation, + subgroup_mapping, +) +from .ops.spec.k3_fused_markov_chain_logits import k3_fused_markov_chain_logits +from .ops.spec.k3_scatter_rows_to_pool import k3_scatter_rows_to_pool + +# Published DSpark drafters spell the head switches four different ways, and a +# reader that knows only one of them degrades silently: the heads are skipped, +# their weights are dropped, and nothing raises. Resolution order matches +# ``TorchLlmArgs.validate_speculative_config`` so the model and the user-visible +# spec_config can never disagree about whether a head is on. +_DSPARK_HEAD_KEY_ALIASES = { + # RadixArk/Kimi-K3-DSpark ships ``enable_confidence_head`` top-level. + "use_confidence_head": ("use_confidence_head", "enable_confidence_head"), +} + + +def resolve_dspark_head_config(pretrained_config, key): + """Resolve one DSpark head switch across every spelling in the wild. + + Looks in ``dspark_config``, then ``dflash_config``, then the top level as + ``dspark_`` and as ````, for each accepted alias of ``key``. + Returns ``None`` when the drafter declares the switch nowhere. + """ + dspark_cfg = getattr(pretrained_config, "dspark_config", None) or {} + dflash_cfg = getattr(pretrained_config, "dflash_config", None) or {} + for name in _DSPARK_HEAD_KEY_ALIASES.get(key, (key,)): + for value in ( + dspark_cfg.get(name), + dflash_cfg.get(name), + getattr(pretrained_config, f"dspark_{name}", None), + getattr(pretrained_config, name, None), + ): + if value is not None: + return value + return None + + +def get_masked_input_and_mask( + input_: torch.Tensor, + vocab_start_index: int, + vocab_end_index: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + vocab_mask = (input_ >= vocab_start_index) & (input_ < vocab_end_index) + valid_offset = vocab_start_index * vocab_mask + input_ = vocab_mask * (input_ - valid_offset) + return input_, ~vocab_mask.unsqueeze(-1) + + +def markov_prev_embeddings(prev_tokens: torch.Tensor, markov_w1: torch.Tensor) -> torch.Tensor: + """``markov_w1[prev_tokens]`` with out-of-vocab anchors masked to zero. + + The anchor is the last accepted token, which on the one-model rejection path + comes from ``trtllm::chain_speculative_sampling``: it pads non-accepted + positions with ``-1``. Mask like modules/embedding.py does, so such a + position contributes no bias instead of tripping a device assert. + + Args: + prev_tokens: previous token ids (draft vocab), any shape. + markov_w1: [vocab, rank]. + Returns: + ``prev_tokens.shape + (rank,)`` in ``markov_w1``'s dtype. + """ + prev_tokens, invalid = get_masked_input_and_mask(prev_tokens.long(), 0, markov_w1.shape[0]) + return F.embedding(prev_tokens, markov_w1).masked_fill(invalid, 0) + + +def dspark_markov_step_bias( + prev_tokens: torch.Tensor, markov_w1: torch.Tensor, markov_w2: torch.Tensor +) -> torch.Tensor: + """Vanilla Markov head logit bias for one intra-block draft step. + + Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ + markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where + markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is + nn.Linear(rank, vocab, bias=False). With both weights stored + [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. + + Args: + prev_tokens: [B] long, previous token per request (draft vocab). + markov_w1: [vocab, rank]. + markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). + Returns: + [B, vocab_or_shard] bias in the markov weights' dtype. + """ + return F.linear(markov_prev_embeddings(prev_tokens, markov_w1), markov_w2) + + +def dspark_markov_chain( + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + step_bias_fn, + *, + hidden_states: Optional[torch.Tensor] = None, + next_token_fn=None, + cast_bias_to_logits: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """The intra-block Markov refinement loop, shared by every DSpark head. + + Reference: DeepSpec ``VanillaMarkov.sample_block_tokens``. For step i, + ``logits_i += bias(prev_i)`` with ``prev_0`` = the anchor token (the last + accepted token, block slot 0) and ``prev_{i>0}`` = the token drawn from + step i-1's *biased* logits. llama.cpp PR #25173 implements the same chain. + + The chain being sequential is load-bearing and is NOT derivable from a + checkpoint: published DSpark drafters ship only the training-time + ``apply_block_logits``, which biases the whole block in one teacher-forced + pass. SGLang ``srt/models/dspark.py:34-64 run_markov_block`` settles it -- + it steps the same way, feeding back the token drawn from the biased logits. + + Args: + base_logits: [B, K, vocab_or_shard] shared-lm_head logits. + first_prev_tokens: [B] long, anchor token ids (draft vocab). + step_bias_fn: ``(prev_tokens [B], step_hidden or None) -> bias``. A + closure, so a stateful head (RNN) can carry its recurrent state + across positions without a second loop. + hidden_states: [B, K, d] fed to ``step_bias_fn`` one position at a + time; None for the memoryless heads. + next_token_fn: ``([B, vocab_or_shard]) -> [B]`` token ids in the FULL + draft vocab; defaults to a plain argmax. Drafters whose draft + logits are TP vocab-sharded pass a shard-aware argmax here. + cast_bias_to_logits: cast the bias down to ``base_logits.dtype`` + before adding. The DFlash drafter does; the V4-Pro drafter does + not (its Markov weights and its logits already agree), and adding + the cast there would silently narrow its accumulation dtype. + Returns: + sampled_tokens [B, K], corrected_logits [B, K, vocab_or_shard]. + Greedy per-position argmax of the corrected logits reproduces the + reference sampled chain exactly. + """ + batch_size, block_size = base_logits.shape[:2] + if block_size == 0: + empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device) + return empty, base_logits + sampled, corrected = [], [] + prev = first_prev_tokens.long() + for k in range(block_size): + step_hidden = None if hidden_states is None else hidden_states[:, k] + bias = step_bias_fn(prev, step_hidden) + if cast_bias_to_logits: + bias = bias.to(base_logits.dtype) + step_logits = base_logits[:, k] + bias + corrected.append(step_logits.unsqueeze(1)) + if next_token_fn is None: + prev = torch.argmax(step_logits, dim=-1) + else: + prev = next_token_fn(step_logits).long() + sampled.append(prev) + return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) + + +def draft_helix_a2a_enabled() -> bool: + """Whether the drafter's helix mapping inherits the target's cp_config.""" + return _features().ull2_draft_a2a + + +def draft_helix_store_enabled() -> bool: + """Whether a helix drafter takes the fused block RoPE with an owner-masked pool store.""" + return _features().ull2_draft_store + + +def draft_helix_tail( + mapping: Mapping, + rows: torch.Tensor, + stats: torch.Tensor, + w_uv: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor], +) -> Optional[torch.Tensor]: + """One drafter layer's helix V up-projection, exchange and combine in one launch, or None. + + Args: + mapping: The drafter's helix mapping (tp 1, cp = the repurposed group). + rows: BF16 ``[T, H, kv_lora_rank]`` (or any view with T * H contiguous rows) latent + attention output of this rank's pages, every head. + stats: FP32 ``[T, H, 2]`` ``(max, sum)`` of ``rows``, as ``_helix_post_process`` takes them. + w_uv: BF16 ``[H, v_head_dim, kv_lora_rank]`` V up-projection. + zero_kv_mask: Optional bool ``[T]``, True where this rank holds no key of the token. + + Returns: + BF16 ``[T, H / cp * v_head_dim]``: this rank's heads after the merge, what the BMM + + ``_helix_post_process`` return; None when the switch is off or the call does not fit the + kernel (the caller keeps that path). The decision depends only on shapes and settings, so + every rank takes the same branch. + """ + if not draft_helix_a2a_enabled(): + return None + if not mapping.has_cp_helix() or mapping.cp_config.get("use_nccl_for_alltoall", True): + return None + if get_sm_version() not in (100, 103, 107): + return None + from .kernels import k3_helix_tail as kernel + from .mla import K3HelixTail + + if w_uv.dim() != 3 or rows.dtype != torch.bfloat16 or w_uv.dtype != torch.bfloat16: + return None + num_heads, v_dim, latent_dim = w_uv.shape + cp_size = mapping.cp_size + if not kernel.can_run(cp_size, num_heads, latent_dim, v_dim): + return None + if not rows.is_contiguous() or not stats.is_contiguous() or not w_uv.is_contiguous(): + return None + num_tokens = rows.numel() // (num_heads * latent_dim) + if num_tokens == 0 or stats.numel() != num_tokens * num_heads * 2: + return None + tail = K3HelixTail._cache.get(mapping) + if tail is None: + if torch.cuda.is_current_stream_capturing(): + # The receive buffer is allocated collectively, outside graph capture; an eager + # forward on every rank creates it first. + logger.warning_once( + "DSpark drafter helix tail: no receive buffer before graph capture, the " + "captured drafter keeps the separate exchange.", + key="dspark_ull2_tail_capture", + ) + return None + tail = K3HelixTail.get(mapping) + if num_tokens * cp_size * kernel.record_bytes(num_heads // cp_size, v_dim) > tail.half_bytes: + return None + return tail( + rows.reshape(num_tokens, num_heads, latent_dim), + stats.reshape(num_tokens, num_heads, 2), + w_uv, + zero_kv_mask, + ) + + +k3_dspark_mla_block_rope = None +k3_dspark_mla_ctx_norm_rope = None +if IS_CUTLASS_DSL_AVAILABLE: + from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( + is_fused_dspark_rmsnorm_rope_supported, + ) + + # Main's DSpark RMSNorm + RoPE kernel compiled for row-strided input, and + # TLLM_DSPARK_TAIL_V2's MLA RoPE kernels. + from .ops.spec.k3_dspark_mla_block_rope import k3_dspark_mla_block_rope + from .ops.spec.k3_dspark_mla_ctx_norm_rope import k3_dspark_mla_ctx_norm_rope + from .ops.spec.k3_dspark_rmsnorm_rope import k3_dspark_rmsnorm_rope + from .ops.spec.k3_dspark_rmsnorm_rope_into import ( + build_rmsnorm_rope_into, + k3_dspark_rmsnorm_rope_into, + ) + +# ---------------------------------------------------------------------------- +# Shared helpers of the block decode (the upper half of modeling_dspark.py). +# ---------------------------------------------------------------------------- + + +def apply_dspark_rotary_batched( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Per-row (batched) DeepSpec ``apply_rotary_emb``. + + ``freqs_cis`` carries a leading batch axis so each row of ``x`` is rotated + by its own per-request phases. ``x`` is the rope-dim slice only: + ``[G, s, rd]`` (3D) or ``[G, s, h, rd]`` (4D), with ``freqs_cis`` of shape + ``[G, s, rd // 2]``. + """ + orig_dtype = x.dtype + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + g, s, half = freqs_cis.shape + if xc.ndim == 3: + fc = freqs_cis.view(g, s, half) + else: + fc = freqs_cis.view(g, s, 1, half) + out = torch.view_as_real(xc * fc).flatten(-2) + return out.to(orig_dtype) + + +def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """RMSNorm matching the DeepSpec reference (fp32 reduce, then * weight).""" + dtype = x.dtype + xf = x.float() + xf = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) + return (weight.float() * xf).to(dtype) + + +def _rope_last_dims_batched( + t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply per-request RoPE to the last dimensions, preserving the non-RoPE prefix.""" + nope = t[..., :-rope_head_dim] + rope = apply_dspark_rotary_batched(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) + return torch.cat([nope, rope], dim=-1) + + +class _PendingDown(NamedTuple): + """A drafter MLP whose down projection runs at the all-reduce site that consumes it.""" + + h2: torch.Tensor + mlp: GatedMlp + + def partial(self) -> torch.Tensor: + """The un-reduced down projection: the pinned split-K GEMM where it applies, else + cuBLAS.""" + weight = self.mlp.down_proj.weight + out = draft_gemm(self.h2, weight, "down") + return out if out is not None else bf16_linear(self.h2, weight) + + +def _draft_mlp( + mlp: GatedMlp, + x: torch.Tensor, + down_all_reduce_params: Optional[AllReduceParams], + defer_down: bool = False, +): + """The drafter MLP on ``x``: ``mlp(x, final_all_reduce_params=down_all_reduce_params)``. + + Where the caller reduces the down projection itself (``enable_allreduce=False``) the gate_up / + down GEMMs are the pinned split-K ones of ``dspark_draft_gemm`` where they apply, around the + same ``silu_and_mul``, and with ``defer_down`` the down projection is returned unrun, as a + :class:`_PendingDown`, for the all-reduce site that consumes it. + """ + h1 = None + if down_all_reduce_params is not None and not down_all_reduce_params.enable_allreduce: + weight = mlp.gate_up_proj.weight + h1 = draft_gemm(x, weight, "gate_up") + if h1 is None and defer_down: + h1 = bf16_linear(x, weight) + if h1 is None: + return mlp(x, final_all_reduce_params=down_all_reduce_params) + pending = _PendingDown(gated_activation(h1, "silu"), mlp) + return pending if defer_down else pending.partial() + + +def _mla_rope_into( + q: torch.Tensor, + freqs_cis: torch.Tensor, + out: torch.Tensor, + num_heads: int, + rope_head_dim: int, + out_rope_offset: int, +) -> bool: + """Rotate q's rope tail into ``out[..., out_rope_offset:]``; False if unavailable. + + Returns a bool rather than raising so the caller keeps the eager path for + builds without the fused kernel, the same way _rmsnorm_rope_batched does. + """ + if not (IS_CUTLASS_DSL_AVAILABLE and is_sm_100f()): + return False + freqs_real = torch.view_as_real(freqs_cis).reshape(-1, freqs_cis.shape[-1], 2) + if not is_fused_dspark_rmsnorm_rope_supported( + q, None, freqs_real, num_heads, rope_head_dim, q.shape[-1] - rope_head_dim + ): + return False + torch.ops.trtllm.cute_dsl_dspark_rope_into( + q, freqs_real, out, num_heads, rope_head_dim, out_rope_offset + ) + return True + + +def _rmsnorm_rope_batched( + t: torch.Tensor, + weight: Optional[torch.Tensor], + eps: float, + rope_head_dim: int, + freqs_cis: torch.Tensor, + *, + num_heads: int = 1, + apply_weight: bool = True, + apply_rmsnorm: bool = True, + inverse_rope: bool = False, + norm_dim: Optional[int] = None, +) -> torch.Tensor: + """Fuse DSpark RMSNorm and last-dimension RoPE when supported. + + ``weight`` may be None when ``apply_weight`` is False -- no path reads it + then, fused or eager. + + ``norm_dim`` bounds the RMSNorm (and the weight) to a prefix of the row. + None means the whole row, which is DSpark's own convention. The other + supported value is ``t.shape[-1] - rope_head_dim``: DeepSeek-style MLA + normalizes the kv_lora_rank latent and leaves k_pe raw, and then ``weight`` + spans only that latent. + """ + if IS_CUTLASS_DSL_AVAILABLE and is_sm_100f(): + freqs_real = torch.view_as_real(freqs_cis).reshape(-1, freqs_cis.shape[-1], 2) + if is_fused_dspark_rmsnorm_rope_supported( + t, weight, freqs_real, num_heads, rope_head_dim, norm_dim + ): + # Reads row-strided views (the fused [q_a | kv_a] GEMM's latent + # columns) in place. + return k3_dspark_rmsnorm_rope( + t, + weight, + freqs_real, + num_heads, + rope_head_dim, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + norm_dim, + ) + + split_norm = norm_dim is not None and norm_dim != t.shape[-1] + head, tail = (t[..., :norm_dim], t[..., norm_dim:]) if split_norm else (t, None) + if apply_rmsnorm: + if apply_weight: + head = _rmsnorm(head, weight, eps) + else: + head = head * torch.rsqrt(head.square().mean(-1, keepdim=True) + eps) + elif apply_weight: + head = (head.float() * weight.float()).to(head.dtype) + t = head if tail is None else torch.cat([head, tail], dim=-1) + if rope_head_dim > 0: + t = _rope_last_dims_batched(t, rope_head_dim, freqs_cis, inverse=inverse_rope) + return t + + +def _runtime_position_cap(model_config, pretrained_config, slack: int = 0) -> int: + """Fallback RoPE table length: the served length, not the advertised one. + + The drafter sizes a position table from the checkpoint's + max_position_embeddings. K3 advertises 1,048,576: ~256 MiB per rank as + complex64, ~768 MiB transient via full-size fp32 cos/sin, while the context + cache is bounded by the runtime max_seq_len. Only the table length moves; + YaRN's correction range comes from original_max_position_embeddings. + + Shrinking the table this way is also what made its bound matter: sized from + max_position_embeddings it could not be indexed past, sized from + model_config.max_seq_len it can. Hence _runtime_position_ceiling, which + the drafter prefers over this value whenever a worker has published it. + """ + runtime = getattr(model_config, "max_seq_len", None) + return int(runtime or getattr(pretrained_config, "max_position_embeddings", 163840)) + slack + + +def _dspark_markov_chain_logits( + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + argmax_fn=None, +) -> torch.Tensor: + """Raw-tensor entry to ``dspark_markov_chain``, corrected logits only. + + The chain of ``dspark_markov`` driven by its vanilla step bias. ``markov_w2`` + may already be sliced down to this rank's TP vocab shard, in which case + ``argmax_fn`` must map a shard-local row back to a full-vocab token id. + """ + + def _step_bias(prev_tokens, step_hidden): + del step_hidden + return dspark_markov_step_bias(prev_tokens, markov_w1, markov_w2) + + _, corrected = dspark_markov_chain( + base_logits, + first_prev_tokens, + _step_bias, + next_token_fn=argmax_fn, + cast_bias_to_logits=True, + ) + return corrected + + +# ---------------------------------------------------------------------------- +# MLA-shaped standalone drafter backbone (Inferact/Kimi-K3-DSpark). +# +# Weights only (layers.py containers): nothing here registers with the +# attention backend or the KV cache manager, and the modules have no +# ``forward``. ``K3MLADSparkForCausalLM.dflash_forward`` runs the block decode +# over them. +# ---------------------------------------------------------------------------- + + +def _yarn_get_mscale(scale: float, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def build_dspark_mla_yarn_rope( + *, + dim: int, + base: float, + scaling_factor: float, + original_max_position_embeddings: int, + max_position_embeddings: int, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + device: str = "cuda", +) -> tuple[torch.Tensor, torch.Tensor]: + """YaRN cos/sin tables for an MLA drafter's ``qk_rope_head_dim`` slice. + + Repacks ``RopeEmbeddingUtils.create_sinusoidal_positions_yarn``, the same HF + DeepSeek-V2 transcription the drafters are distilled under. NOT reused: + ``RopeParams.create_rope_const_params()`` pins the fused MLA kernel's GPT-J + packing and hard-codes ``device='cuda'``. + + ``duplicate_data=True`` though only the first half is read as a complex + pair: the util's two modes disagree by 1 ulp on 6 of 8224 entries, and the + duplicated one is what every measured drafter run used. + + Returns ``(cos, sin)``, both ``[max_position_embeddings, dim]`` fp32, with + the frequency half duplicated so ``rotate_half`` applies. + """ + # The util returns a flat numpy [pos][slot](cos, sin); split the pair. + _, packed = RopeEmbeddingUtils.create_sinusoidal_positions_yarn( + num_pos=max_position_embeddings, + dim=dim, + base=base, + scaling_factor=scaling_factor, + original_max_position_embeddings=original_max_position_embeddings, + beta_fast=beta_fast, + beta_slow=beta_slow, + mscale=mscale, + mscale_all_dim=mscale_all_dim, + duplicate_data=True, + ) + table = torch.from_numpy(packed).reshape(max_position_embeddings, dim, 2) + return ( + table[..., 0].contiguous().to(device), + table[..., 1].contiguous().to(device), + ) + + +def build_dspark_mla_yarn_freqs_cis(*, device: str = "cuda", **kwargs) -> torch.Tensor: + """The same YaRN table as :func:`build_dspark_mla_yarn_rope`, adjacent-pair. + + Returns ``[max_position_embeddings, dim // 2]`` complex64: the convention the + fused ``cute_dsl_dspark_rmsnorm_rope`` kernel consumes, where the last dim of + the rotated slice is read as (re, im) pairs. + + HF DeepSeek de-interleaves and rotates halves instead, writing the same + values to lanes i and i + dim/2 rather than 2i and 2i+1. That permutation + cancels inside ``q_rope . k_rope``, so either convention works -- but every + producer of a rope slice must use the SAME one, and mixing them only shows + up as lower AL. ``mscale`` rides in the modulus, not the phase. + """ + cos, sin = build_dspark_mla_yarn_rope(device=device, **kwargs) + half = cos.shape[-1] // 2 + return torch.complex(cos[..., :half], sin[..., :half]).contiguous() + + +def apply_dspark_mla_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Rotate the trailing ``qk_rope_head_dim`` of ``x`` the HF DeepSeek way. + + ``x`` is ``[..., dim]`` and ``cos``/``sin`` are ``[..., dim]`` already + gathered at the query positions. HF interleaves the rope slice, so it first + reshapes ``(d//2, 2) -> transpose -> (d,)`` and only then applies the + half-split ``rotate_half``; the net rotation is GPT-J style. + + Equal to ``RotaryEmbedding.apply_rotary_pos_emb(..., is_neox=False)`` up to + an even/odd lane split of its output (checked exactly, max diff 0.0). Kept + separate because that helper unsqueezes cos/sin itself, so routing through + it would push a rank contract onto every caller for no numerical gain. + """ + d = x.shape[-1] + x = x.reshape(*x.shape[:-1], d // 2, 2).transpose(-1, -2).reshape(*x.shape[:-1], d) + x1, x2 = x[..., : d // 2], x[..., d // 2 :] + rotated = torch.cat((-x2, x1), dim=-1) + return x * cos + rotated * sin + + +def _draft_mappings(mapping: Mapping) -> tuple[Mapping, Optional[Mapping]]: + """``(dense, helix)`` of the drafter's mapping. + + ``helix`` is the helix CP mapping attention keeps (the latent cache split by pages over the CP + ranks), None off helix; ``dense`` is the mapping of everything else, the helix CP ranks + repurposed as TP ranks. + """ + if mapping.has_cp_helix(): + return mapping.repurpose_helix_cp_to_tp(), mapping + if mapping.cp_size > 1: + raise NotImplementedError( + f"The MLA DSpark drafter does not support context parallelism " + f"(cp_size={mapping.cp_size})." + ) + return mapping, None + + +class K3MLADSparkAttention(nn.Module): + """The MLA weights of one drafter layer. + + The low-rank ``*_a`` projections are replicated (every head reads the same latent); + ``q_b_proj`` / ``kv_b_proj`` are column-sharded by head over the attention TP group and + ``o_proj`` is row-sharded. Under attention DP the attention TP group is one rank, which is + also why the drafter's latent cache is unsharded. Under helix the latent cache is split by + pages over the CP ranks, so every rank holds every head (``q_b`` / ``kv_b`` replicated over + the CP group) and ``o_proj`` is row-parallel over the repurposed (CP folded into TP) group. + """ + + def __init__( + self, + config, + dense: Mapping, + helix: Optional[Mapping], + allreduce_strategy: AllReduceStrategy, + ): + super().__init__() + dtype = config.torch_dtype + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + # What the drafter's KV cache stores per token: one MLA latent, no V. + self.kv_cache_head_dim = self.kv_lora_rank + self.qk_rope_head_dim + + mapping = dense if helix is None else helix + tp_size, dp_size = mapping.tp_size, 1 + if mapping.enable_attention_dp: + tp_size, dp_size = 1, mapping.tp_size + elif helix is not None: + # CP ranks hold identical head projections: fold them like DP. + dp_size = mapping.cp_size + if self.num_heads % tp_size != 0: + raise ValueError( + f"DSpark MLA drafter has {self.num_heads} heads, " + f"not divisible by tp_size {tp_size}." + ) + self.num_heads_tp = self.num_heads // tp_size + self.num_heads_tp_cp = self.num_heads_tp // mapping.cp_size + self.mapping = Mapping( + world_size=mapping.pp_size * dp_size * tp_size, + tp_size=tp_size, + pp_size=mapping.pp_size * dp_size, + rank=mapping.rank, + gpus_per_node=mapping.gpus_per_node, + enable_attention_dp=mapping.enable_attention_dp, + ) + o_mapping = self.mapping if helix is None else dense + # Row-parallel o_proj reduces only when the heads are actually split. + reduce_output = not dense.enable_attention_dp and dense.tp_size > 1 + eps = config.rms_norm_eps + by_head = dict( + dtype=dtype, + tp_size=self.mapping.tp_size, + tp_rank=self.mapping.tp_rank, + tp_mode=TensorParallelMode.COLUMN, + ) + + self.q_a_proj = Bf16Linear(self.hidden_size, self.q_lora_rank, dtype=dtype) + self.q_a_layernorm = RMSNormWeight(self.q_lora_rank, eps, dtype) + self.q_b_proj = Bf16Linear(self.q_lora_rank, self.num_heads * self.qk_head_dim, **by_head) + self.kv_a_proj_with_mqa = Bf16Linear(self.hidden_size, self.kv_cache_head_dim, dtype=dtype) + self.kv_a_layernorm = RMSNormWeight(self.kv_lora_rank, eps, dtype) + self.kv_b_proj = Bf16Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + **by_head, + ) + self.o_proj = Bf16Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + dtype=dtype, + tp_size=o_mapping.tp_size, + tp_rank=o_mapping.tp_rank, + tp_mode=TensorParallelMode.ROW, + all_reduce=( + TpAllReduce(o_mapping, dtype, allreduce_strategy) if reduce_output else None + ), + ) + + +class K3MLADSparkDecoderLayer(nn.Module): + """One drafter layer: the MLA weights, the dense SwiGLU MLP and the two RMSNorms.""" + + def __init__( + self, + config, + layer_idx: int, + dense: Mapping, + helix: Optional[Mapping], + allreduce_strategy: AllReduceStrategy, + ): + super().__init__() + dtype, hidden, eps = config.torch_dtype, config.hidden_size, config.rms_norm_eps + self.layer_idx = layer_idx + self.self_attn = K3MLADSparkAttention(config, dense, helix, allreduce_strategy) + # Under attention DP every rank runs the whole MLP: TP groups of one rank. + self.mlp = GatedMlp( + hidden, + config.intermediate_size, + activation="silu", + dtype=dtype, + mapping=subgroup_mapping(dense, 1) if dense.enable_attention_dp else dense, + allreduce_strategy=allreduce_strategy, + ) + self.input_layernorm = RMSNormWeight(hidden, eps, dtype) + self.post_attention_layernorm = RMSNormWeight(hidden, eps, dtype) + + +class K3MLADSparkModel(nn.Module): + """The drafter's embedding, layers and final norm.""" + + def __init__(self, model_config): + super().__init__() + config = model_config.pretrained_config + dense, helix = _draft_mappings(model_config.mapping) + strategy = model_config.allreduce_strategy + # The drafter ships its own trained embedding rather than sharing the target's (verified + # different weights on Inferact/Kimi-K3-DSpark), so it is loaded from the drafter + # checkpoint. Vocab-parallel over the dense TP group, the lookup all-reduced at the + # default strategy. + self.embed_tokens = EmbeddingTable( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + mapping=dense, + vocab_parallel=True, + ) + self.layers = nn.ModuleList( + [ + K3MLADSparkDecoderLayer(config, layer_idx, dense, helix, strategy) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNormWeight(config.hidden_size, config.rms_norm_eps, config.torch_dtype) + + +class K3MLADSparkBackbone(nn.Module): + """The drafter backbone (``K3MLADSparkForCausalLM.draft_model_full``): weights only. + + ``K3MLADSparkForCausalLM.dflash_forward`` runs the block decode over ``model``. No ``lm_head``: + the drafter shares the target's (``K3MLADSparkForCausalLM.load_weights_from_target_model``). + """ + + def __init__(self, model_config): + super().__init__() + self.model_config = model_config + self.config = model_config.pretrained_config + self.model = K3MLADSparkModel(model_config) + self.lm_head = None + + +# ``DFlashForCausalLM.__init__`` builds its backbone as the model class registered for +# ``pretrained_config.architectures[0]``. For the base ``__init__`` of the MLA drafter that is +# this package's backbone, under a name no checkpoint carries. +_BACKBONE_ARCH = "K3RubinMLADSparkBackbone" +# What the drafter config names afterwards: the draft weight loader picks its checkpoint +# mapper by it (the checkpoint itself says "K3DSparkModel", registered nowhere). +_CHECKPOINT_ARCH = "K3DsparkForCausalLM" + + +@contextlib.contextmanager +def _builds_k3_backbone(pretrained_config) -> Iterator[None]: + """Within the block, the model registry resolves the drafter config to + :class:`K3MLADSparkBackbone`.""" + if _BACKBONE_ARCH in MODEL_CLASS_MAPPING: + raise RuntimeError(f"model architecture {_BACKBONE_ARCH!r} is already registered") + MODEL_CLASS_MAPPING[_BACKBONE_ARCH] = K3MLADSparkBackbone + pretrained_config.architectures = [_BACKBONE_ARCH] + try: + yield + finally: + MODEL_CLASS_MAPPING.pop(_BACKBONE_ARCH, None) + pretrained_config.architectures = [_CHECKPOINT_ARCH] + + +# ---------------------------------------------------------------------------- +# Standalone DSpark drafter. +# +# The drafter ships as its own checkpoint instead of living in the target's +# ``mtp.*`` namespace, so it shares nothing with the target but the vocabulary +# and the captured hidden states. It is DFlash plus a Markov logit bias, a +# confidence head and the shift_label slot convention, so it subclasses +# ``DFlashForCausalLM`` and adds those three, on an MLA backbone with its own +# block decode. +# ---------------------------------------------------------------------------- + +# Both published drafters (RadixArk/Kimi-K3-DSpark, Inferact/Kimi-K3-DSpark) +# name the head tensors after the submodules that own them: markov_head is a +# VanillaMarkov, confidence_head an AcceptRatePredictor whose linear is ``proj``. +# The bare spellings are kept for drafters exported without that nesting. +_DSPARK_HEAD_WEIGHT_ALIASES = { + "markov_w1.weight": ("markov_head.markov_w1.weight", "markov_w1.weight"), + "markov_w2.weight": ("markov_head.markov_w2.weight", "markov_w2.weight"), + "confidence_proj.weight": ("confidence_head.proj.weight", "confidence_proj.weight"), + "confidence_proj.bias": ("confidence_head.proj.bias", "confidence_proj.bias"), +} + + +# Lower bound of the cute-dsl MLA decode's workspace, grown to the runner's padded size. +_MLA_DECODE_WORKSPACE_BYTES = 256 * 1024 * 1024 +# Fixed cost of one split-KV piece, in KV tiles, in the per-request split plan +# of the cute-dsl block decode (_launch_mla_split_plan): 64 heads folded two +# tokens per M tile, fitted on R200 sweeps over 7.5K-200K tokens of context, +# batch 1-32, equal and mixed lengths. +_MLA_SPLIT_PLAN_PIECE_COST_TILES = 30 +# The planner kernel is one program over the batch, padded to a power of two of at least one +# warp: each of its candidate cuts is a reduction over the batch. +_MLA_SPLIT_PLAN_MIN_BLOCK = 32 +_MLA_SPLIT_PLAN_LANES_PER_WARP = 256 + + +class _MLABlockFixup(NamedTuple): + """Layer-invariant operands of the MLA block decode. + + The block's own latents are written into the pool at ``ctx_len + j``. Those + slots belong to this request -- the draft KV manager adds + ``max(draft_len, max_total_draft_tokens)`` tokens per generation request each + step (resource_manager.py:1060) -- and the accepted tokens overwrite them next + step. The decode then reads context and block in one pass, every block token + bounded at ``ctx_len + block_size`` (``kv_bounds_i32``). Every field depends on + ctx_len, the page table and the block geometry alone, so it is built once per + forward. + """ + + pages: torch.Tensor # [B, block] page holding each block position + offsets: torch.Tensor # [B, block] offset of that position inside its page + valid: torch.Tensor # [1, block, 1, block] strict upper triangle (unread by the decode) + page_tables_i32: torch.Tensor # kernel operands, cast once + seq_lens_i32: torch.Tensor # ctx_len + block_size + kv_bounds_i32: torch.Tensor # [B * block] seq_lens repeated per query token + # Helix: [B, block, 1] whether this rank's pages hold each block position; + # pages/offsets and the lengths above are then rank-local. + owned: Optional[torch.Tensor] = None + + +@lru_cache(maxsize=1) +def cute_dsl_mla_decode_unavailability_reason() -> Optional[str]: + """Return why the cute-dsl MLA decode cannot run here, or None. + + The op itself only reports by raising, and ``attention_backend="CUTEDSL"`` + has to fail at construction rather than five layers into the first draft + step. + + The ``kv_bounds`` probe is not paranoia. The op exists upstream without + per-token KV bounds, and this backend's whole shape -- one pass over + context and block, no fixup -- is expressed through them. Without the probe + the mismatch surfaces as a torch dispatch TypeError at the first draft + step, naming an argument count rather than a missing kernel feature. + """ + if (sm := get_sm_version()) not in (100, 103, 107): + return f"requires SM100, SM103 or SM107, got SM{sm}" + try: + import cutlass # noqa: F401 + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLNVMlaDecodeBlackwellRunner, # noqa: F401 + ) + except ImportError as error: + return f"the cute-dsl kernels are not importable: {error}" + try: + op = torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell + except AttributeError: + return "trtllm::cute_dsl_mla_decode_fp16_blackwell is not registered" + if not any( + arg.name == "kv_bounds" + for overload in op.overloads() + for arg in getattr(op, overload)._schema.arguments + ): + return ( + "the registered cute_dsl_mla_decode_fp16_blackwell takes no kv_bounds, " + "so per-token KV bounds are unavailable in this build" + ) + return None + + +def _all_reduce_add_norm( + all_reduce: TpAllReduce, partial: torch.Tensor, residual: torch.Tensor, norm: RMSNormWeight +) -> tuple[torch.Tensor, torch.Tensor]: + """``norm(all_reduce(partial) + residual)`` and the new residual, fused.""" + return all_reduce( + partial, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm.weight, + eps=norm.variance_epsilon, + ), + ) + + +def _block_upper_triangle(block_size: int, device) -> torch.Tensor: + """``[1, block, 1, block]``: query i sees block key j > i (the fixup's part).""" + idx = torch.arange(block_size, device=device) + return (idx.view(1, block_size) > idx.view(block_size, 1)).view(1, block_size, 1, block_size) + + +def _build_mla_block_fixup( + ctx_len, page_tables, block_size, page_size, shard: Optional[HelixDraftShard] = None +) -> _MLABlockFixup: + """Where the block's latents go, and the upper-triangle mask, for all layers. + + With a helix ``shard`` every position goes to its rank-local slot and the + lengths count only this rank's positions. + """ + device = ctx_len.device + pos = ctx_len.view(-1, 1) + torch.arange(block_size, device=device) + seq_lens = (ctx_len + block_size).to(torch.int32) + owned = None + if shard is not None: + pos, owned = shard.local_slots(pos) + owned = owned.unsqueeze(-1) + seq_lens = shard.local_lengths(seq_lens) + return _MLABlockFixup( + owned=owned, + pages=page_tables.gather(1, pos // page_size), + offsets=pos % page_size, + valid=_block_upper_triangle(block_size, device), + page_tables_i32=page_tables.to(torch.int32), + seq_lens_i32=seq_lens, + # repeat_interleave, not repeat: the kernel indexes kv_bounds at + # b * seq_len_q + q (mla_decode_fp16.py:400-407), so the batch must be + # the slow axis. B == 1 cannot tell the two apart. + kv_bounds_i32=seq_lens.repeat_interleave(block_size), + ) + + +def _resolve_dspark_mla_rope_params(pretrained_config, model_config=None, slack: int = 0) -> Dict: + """RoPE settings for an MLA drafter, from either HF spelling. + + Transformers v5 checkpoints (TorchSpec) carry ``rope_parameters``; older + ones carry ``rope_scaling``. Only YaRN is accepted: both published MLA + drafters use it, and silently treating an unknown scaling as plain RoPE + would cost acceptance without an error. + """ + scaling = ( + getattr(pretrained_config, "rope_parameters", None) + or getattr(pretrained_config, "rope_scaling", None) + or {} + ) + rope_type = scaling.get("rope_type") or scaling.get("type") or "default" + if rope_type != "yarn": + raise ValueError( + f"MLA DSpark drafter declares rope_type={rope_type!r}; only 'yarn' is " + "supported (the block decode builds its own YaRN table)." + ) + theta = scaling.get("rope_theta") or getattr(pretrained_config, "rope_theta", 10000.0) + return { + "theta": float(theta), + "scaling_factor": float(scaling["factor"]), + "original_max_positions": int(scaling["original_max_position_embeddings"]), + # Table length, not YaRN math: the correction range below uses + # original_max_positions, so capping this at the served length is free. + # K3 advertises 1,048,576, which would build a ~256 MiB complex64 table + # per rank -- through full-size fp32 cos/sin, so ~768 MiB transient -- + # for a drafter whose point is 5760 B/token instead of 20480. + "max_positions": _runtime_position_cap(model_config, pretrained_config, slack), + "beta_fast": float(scaling.get("beta_fast", 32.0)), + "beta_slow": float(scaling.get("beta_slow", 1.0)), + "mscale": float(scaling.get("mscale", 1.0)), + "mscale_all_dim": float(scaling.get("mscale_all_dim", 0.0)), + } + + +def _launch_mla_split_plan( + seq_lens: torch.Tensor, + plan: torch.Tensor, + *, + num_heads: int, + seq_len_q: int, + piece_cost: int, +) -> None: + """Write the per-request split counts of the cute-dsl MLA decode at ``piece_cost`` into + ``plan`` (one launch). + + The CuTe DSL MLA decode's planner kernel (``mla_split_plan._mla_split_kv_plan_kernel``) on + the decode runner's geometry (:func:`_mla_split_plan_geometry`), launched with a block of the + batch rounded up to a power of two, at this decode's piece cost. + """ + import triton + + from .kernels.spec import mla_split_plan as mla + + num_requests = seq_lens.shape[0] + if not 0 < num_requests <= mla._SPLIT_PLAN_MAX_REQUESTS: + raise ValueError(f"split plan supports 1..{mla._SPLIT_PLAN_MAX_REQUESTS} requests") + if piece_cost < 0: + raise ValueError(f"split plan piece cost must be >= 0, got {piece_cost}") + for name, tensor in (("seq_lens", seq_lens), ("plan", plan)): + if ( + tensor.dtype != torch.int32 + or tensor.ndim != 1 + or tensor.stride(0) != 1 + or tensor.numel() < num_requests + or not tensor.is_cuda + ): + raise ValueError(f"{name} must be a contiguous CUDA int32 vector of >= {num_requests}") + rows, wave_clusters = _mla_split_plan_geometry(num_heads, seq_len_q) + block = max(_MLA_SPLIT_PLAN_MIN_BLOCK, triton.next_power_of_2(num_requests)) + mla._mla_split_kv_plan_kernel[(1,)]( + seq_lens, + plan, + num_requests, + max(rows, 1), + max(wave_clusters, 1), + mla._SPLIT_PLAN_MAX_SPLIT, + KV_TILE=mla._SPLIT_PLAN_KV_TILE, + MIN_TILES=mla._SPLIT_PLAN_MIN_TILES, + PIECE_COST=int(piece_cost), + MAX_PIECES=mla._SPLIT_PLAN_MAX_SPLIT, + BLOCK=block, + num_warps=max(1, min(4, block // _MLA_SPLIT_PLAN_LANES_PER_WARP)), + ) + + +def _mla_split_plan_geometry(num_heads: int, seq_len_q: int) -> tuple[int, int]: + """``(rows_per_request, wave_clusters)`` of the CuTe DSL MLA decode, for the planner.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + runner = CuteDSLNVMlaDecodeBlackwellRunner + rows_per_request = runner.get_num_q_tiles(num_heads, seq_len_q) + wave_clusters = max(runner._get_max_active_blocks() // runner._CLUSTER_SHAPE_MNK[0], 1) + return rows_per_request, wave_clusters + + +@lru_cache(maxsize=1) +def _helix_post_process_fn() -> tuple[Callable, bool]: + """``_helix_post_process`` and whether it takes ``compile_nccl_reformat``.""" + from .drafter import _helix_post_process + + takes_flag = "compile_nccl_reformat" in inspect.signature(_helix_post_process).parameters + return _helix_post_process, takes_flag + + +def _draft_helix_post_process( + partial_o: torch.Tensor, + stats: torch.Tensor, + mapping: Mapping, + num_heads_tp_cp: int, + value_dim: int, + zero_kv_mask: Optional[torch.Tensor], +) -> torch.Tensor: + """``_helix_post_process`` for the drafter: NCCL reformats eager where supported. + + The drafter meets each batch bucket's shape first in a warmup forward, + after the target's layers: no compile between NCCL all-to-alls (it can + block every rank). ``drafter._helix_post_process`` has no + ``compile_nccl_reformat`` switch, so its NCCL path compiles. + """ + fn, takes_flag = _helix_post_process_fn() + kwargs = {"zero_kv_mask": zero_kv_mask} + if takes_flag: + kwargs["compile_nccl_reformat"] = False + return fn(partial_o, stats, mapping, num_heads_tp_cp, value_dim, **kwargs) + + +class _DSparkHeadMixin: + """The DSpark head set of a standalone drafter wrapper. + + DSpark is DFlash plus three things common to both backbone shapes: the + vanilla Markov intra-block logit bias, the ``shift_label`` slot convention + (block slot j predicts draft token j+1, so slot 0 holds the anchor) and the + confidence head. The block decode has nothing in common between the shapes, + which is why this is a mixin and not a base class. + + Confidence-scheduled verification is not implemented yet: ``confidence_proj`` + is loaded but unused, and drafting always proposes the full K tokens. + """ + + def _init_dspark_heads(self, cfg) -> None: + # Defaults on, unlike the DFlash base: the shift_label slot layout is + # part of what DSpark *is*, and both published drafters set + # block_size == max_draft_len, where the DFlash layout (slots 1..K) + # runs one slot past the block and reads the next request's anchor. + # An explicit false still selects the legacy layout. + shift_label = resolve_dspark_head_config(cfg, "shift_label") + self.shift_label = True if shift_label is None else bool(shift_label) + self._dspark_markov_rank = int(resolve_dspark_head_config(cfg, "markov_rank") or 0) + self._dspark_markov_head_type = str( + resolve_dspark_head_config(cfg, "markov_head_type") or "vanilla" + ).lower() + self._dspark_use_confidence_head = bool( + resolve_dspark_head_config(cfg, "use_confidence_head") or False + ) + # Plain None placeholders rather than nn.Parameter/buffer: the shapes + # ([vocab, rank]) are checkpoint-dependent, so nothing is pre-allocated + # and nothing is constructed in the module's default dtype. Using the + # checkpoint tensor as-is is also what keeps the head in the checkpoint's + # dtype instead of an nn.Module default. load_weights() fills them in; + # consumers treat None as "head absent". + self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) + self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) + self.confidence_proj_weight = None # loaded, unused (follow-up MR) + self.confidence_proj_bias = None + + if self._dspark_markov_rank > 0 and self._dspark_markov_head_type != "vanilla": + raise ValueError( + f"DSpark drafter declares markov_head_type=" + f"'{self._dspark_markov_head_type}'; only 'vanilla' is " + "supported (gated/rnn heads need per-step hidden features)." + ) + # The block decode only supports the non-causal DSpark convention. + # Legacy DFlash drafter configs (e.g. Laguna) also carry a causal field + # and handle it in the legacy decode path, which is why this check lives + # here rather than in the DFlash base. + if resolve_dspark_head_config(cfg, "causal"): + raise ValueError( + "DSpark drafter sets causal=true; the block decode only " + "supports the non-causal DSpark convention." + ) + if self._dspark_use_confidence_head: + logger.warning( + "DSpark drafter declares use_confidence_head; " + "confidence-scheduled verification is not implemented yet " + "(confidence_proj weights are loaded but unused, drafting " + "always proposes the full K tokens)." + ) + + @property + def has_markov_head(self) -> bool: + return self._dspark_markov_rank > 0 and self.markov_w1 is not None + + def apply_markov_chain_logits( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + tp_exchange=None, + argmax_fn=None, + vocab_slice: Optional[slice] = None, + ) -> torch.Tensor: + """Apply the vanilla-Markov intra-block bias to block logits. + + No-op (returns ``base_logits`` unchanged) when the checkpoint ships no + Markov head. See :func:`dspark_markov_chain` for the semantics. Two + calling conventions, one per worker: + + * main's ``DSparkWorker``: float32 ``base_logits`` of the full vocab, or + this rank's vocab shard with the TP group's ``tp_exchange`` -- + :func:`markov_chain_logits` runs the chain; + * the K3 ``K3DFlashWorker``: ``argmax_fn`` returning full-vocab token ids + and, for a TP vocab shard, this rank's ``vocab_slice`` (to shard the + markov_w2 rows identically) -- the eager chain runs. + """ + if not self.has_markov_head: + return base_logits + if argmax_fn is None and vocab_slice is None: + return markov_chain_logits( + base_logits, first_prev_tokens, self.markov_w1, self.markov_w2, tp_exchange + ) + if tp_exchange is not None: + raise ValueError( + "apply_markov_chain_logits takes either tp_exchange or argmax_fn / vocab_slice, " + "not both." + ) + markov_w2 = self.markov_w2 if vocab_slice is None else self.markov_w2[vocab_slice] + return _dspark_markov_chain_logits( + base_logits, first_prev_tokens, self.markov_w1, markov_w2, argmax_fn=argmax_fn + ) + + def _take_dspark_head_weights(self, weights: Dict) -> tuple[Dict, Dict]: + """Pull the head tensors out of ``weights`` and load them onto self. + + The head keys are pulled out before the backbone remap: left in, they + would pick up a ``model.`` prefix and be dropped by partial loading. + Returns the remaining weights and the extracted head weights. + """ + dspark_weights = {} + consumed = set() + for canonical, aliases in _DSPARK_HEAD_WEIGHT_ALIASES.items(): + for name in aliases: + if name in weights: + dspark_weights[canonical] = weights[name] + consumed.add(name) + break + if consumed: + weights = {k: v for k, v in weights.items() if k not in consumed} + # The inverse of the missing-weights check below. Without it, a config + # whose head switches this build cannot resolve loads the drafter with + # the heads silently dropped -- correct output, lower acceptance. + if self._dspark_markov_rank <= 0 and "markov_w1.weight" in dspark_weights: + raise ValueError( + "DSpark drafter ships markov_w1/markov_w2 but markov_rank resolved to 0. " + "The checkpoint's head switches were not found in dspark_config, " + "dflash_config, or at the top level; loading it would drop the Markov " + "head silently." + ) + if self._dspark_markov_rank > 0: + vocab = self.config.vocab_size + rank = self._dspark_markov_rank + for k in ("markov_w1.weight", "markov_w2.weight"): + if k not in dspark_weights: + raise ValueError( + f"DSpark drafter declares markov_rank=" + f"{self._dspark_markov_rank} but the checkpoint is " + f"missing {k}." + ) + if tuple(dspark_weights[k].shape) != (vocab, rank): + raise ValueError( + f"DSpark {k} has shape " + f"{tuple(dspark_weights[k].shape)}, expected " + f"[vocab, markov_rank] = ({vocab}, {rank})." + ) + self.markov_w1 = dspark_weights["markov_w1.weight"].to("cuda") + self.markov_w2 = dspark_weights["markov_w2.weight"].to("cuda") + if "confidence_proj.weight" in dspark_weights: + self.confidence_proj_weight = dspark_weights["confidence_proj.weight"].to("cuda") + if "confidence_proj.bias" in dspark_weights: + self.confidence_proj_bias = dspark_weights["confidence_proj.bias"].to("cuda") + return weights, dspark_weights + + +class _K3DFlashDrafterMixin: + """The capture projection of the K3 drafter: ``hidden_norm(fc(target hidden))``. + + ``fc`` and ``hidden_norm`` are built from the checkpoint by ``load_weights``, ``fc`` with its + whole weight. Under plain TP (no attention DP) whose group divides the fc input width, every + rank multiplies its slice of the captured features by its slice of the ``fc`` columns and one + all-reduce sums the partials, with ``hidden_norm`` fused in over MNNVL. + """ + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project captured target hidden states into the draft hidden space.""" + hidden_states = hidden_states.to(self.fc.weight.dtype) + if hidden_states.shape[0] > 0: + projected = self._tp_sharded_project_target_hidden(hidden_states) + if projected is not None: + return projected + norm = self.hidden_norm + return F.rms_norm( + bf16_linear(hidden_states, self.fc.weight), + (norm.weight.shape[0],), + norm.weight, + norm.eps, + ) + + def _resolve_fc_shard(self): + """``(weight slice, begin, end, all_reduce)`` of the TP-sharded fc, or None. + + Resolved once: the shard exists when the drafter runs plain TP (no attention DP) and the + fc input width splits evenly over it. The TP group is o_proj's all-reduce group: under + helix the drafter's ``model_config.mapping`` carries the CP split of its context KV + instead. + """ + if getattr(self, "_fc_shard", None) is None: + self._fc_shard = False + all_reduce = self.model.layers[0].self_attn.o_proj.all_reduce + mapping = None if all_reduce is None else all_reduce.mapping + in_features = self.fc.weight.shape[1] + if ( + mapping is not None + and mapping.tp_size > 1 + and not mapping.enable_attention_dp + and in_features % mapping.tp_size == 0 + ): + width = in_features // mapping.tp_size + begin = mapping.tp_rank * width + weight = self.fc.weight.data[:, begin : begin + width].contiguous() + self._fc_shard = (weight, begin, begin + width, all_reduce) + return self._fc_shard or None + + def _tp_sharded_project_target_hidden(self, hidden_states: torch.Tensor): + """The TP-sharded fc (``_resolve_fc_shard``) and its all-reduce with ``hidden_norm``, or + None when the fc is not sharded.""" + shard = self._resolve_fc_shard() + if shard is None: + return None + weight, begin, end, all_reduce = shard + partial = bf16_linear(hidden_states[:, begin:end], weight) + norm = self.hidden_norm + if all_reduce.mnnvl_allreduce is not None: + return all_reduce( + partial, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, norm_weight=norm.weight, eps=norm.eps + ), + ) + return F.rms_norm(all_reduce(partial), (norm.weight.shape[0],), norm.weight, norm.eps) + + +def _latent_cache_helix_mapping( + mapping: Mapping, helix_mapping: Optional[Mapping], inherit_cp_config: bool = False +) -> Optional[Mapping]: + """The mapping that splits a one-latent-per-token cache by pages over a helix group. + + ``mapping`` is the drafter's: the target's helix group with its CP ranks repurposed + as TP ranks. ``helix_mapping`` is the target's helix CP mapping, or None. The result + keeps every head on every rank (tp = 1) and shards the sequence over cp = + ``mapping.tp_size`` consecutive ranks, the target's CP groups. Returns None where + that does not apply: no helix, attention DP, pipeline parallelism, a single rank, or + a draft TP group that does not divide the target's CP (target tp > 1). + + ``inherit_cp_config=True`` gives the result ``helix_mapping``'s cp_config (its + all-to-all transport, e.g. ``use_nccl_for_alltoall: false`` for the MNNVL exchange) + with cp_type HELIX; otherwise its cp_config is ``{cp_type: HELIX}`` alone, the NCCL + all-to-all. + """ + if ( + helix_mapping is None + or not helix_mapping.has_cp_helix() + or mapping.enable_attention_dp + or mapping.pp_size != 1 + or mapping.tp_size <= 1 + or helix_mapping.cp_size % mapping.tp_size != 0 + ): + return None + cp_config = {"cp_type": CpType.HELIX} + if inherit_cp_config: + cp_config = {**helix_mapping.cp_config, **cp_config} + return Mapping( + world_size=mapping.world_size, + rank=mapping.rank, + gpus_per_node=mapping.gpus_per_node, + cp_size=mapping.tp_size, + cp_config=cp_config, + tp_size=1, + pp_size=1, + moe_tp_size=mapping.moe_tp_size, + moe_ep_size=mapping.moe_ep_size, + moe_cluster_size=mapping.moe_cluster_size, + enable_lm_head_tp_in_adp=mapping.enable_lm_head_tp_in_adp, + ) + + +def _shard_context_kv_by_helix(draft_config, helix_mapping: Optional[Mapping]) -> Optional[Mapping]: + """Split the latent cache by pages over the whole repurposed helix group. + + One latent per token cannot be split by head, so attention runs at tp = 1 with every head + on every rank, and the CuTe DSL decode takes the rank-local bounds this needs. + ``helix_mapping`` is the target's helix CP mapping (None outside Helix); the sharding is + :func:`_latent_cache_helix_mapping` of it, which replaces ``draft_config.mapping``. In a + unified draft cache the pages follow the target manager's helix ledger, which the DFlash + worker requires to be this mapping's. Returns the sharding, or None. + """ + # ull2_draft_a2a: exchange on the target's helix transport. + mapping = _latent_cache_helix_mapping( + draft_config.mapping, helix_mapping, inherit_cp_config=draft_helix_a2a_enabled() + ) + if mapping is None: + return None + frozen, draft_config._frozen = draft_config._frozen, False + draft_config.mapping = mapping + draft_config._frozen = frozen + transport = "NCCL" if mapping.cp_config.get("use_nccl_for_alltoall", True) else "MNNVL" + logger.info( + f"MLA DSpark: latent cache pages sharded over helix cp={mapping.cp_size} " + f"({transport} all-to-all)" + ) + return mapping + + +class K3MLADSparkForCausalLM(_DSparkHeadMixin, _K3DFlashDrafterMixin, DFlashForCausalLM): + """DSpark drafter on an MLA-shaped backbone, from a standalone checkpoint. + + A ``DFlashForCausalLM`` for the DFlash worker contract -- dual-source KV, non-causal block + attention, one precomputed context cache -- with its own weights (:class:`K3MLADSparkBackbone`, + built by the base ``__init__``) and its own block decode, since DFlash's splits a fused + ``qkv_proj`` by (num_heads, num_kv_heads, head_dim) and fuses per-head K/V across layers, + neither of which MLA has. + + The cache stores one ``kv_lora_rank + qk_rope_head_dim`` latent per token + per layer and no V, which is the point: 5 x 576 x 2B = 5760 B/token/rank + against the GQA drafter's 20480 under attention-DP, where KV heads are + unsharded. It lives in the draft KV manager's paged pool, BF16. + + The block decode is one cute-dsl pass over context and block with per-token + ``kv_bounds`` (``attention_backend="CUTEDSL"``, the only one). + + Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. + """ + + # The block decode is this class's own, so neither worker OP SET applies and + # neither backend's shape checks bind -- the absorbed shape (64 heads : 1 KV + # head, head_dim 576) would fail both. + _uses_worker_attention_backend = False + # One cute-dsl pass over context and block through the per-token kv_bounds of + # cute_dsl_mla_decode_fp16_blackwell, no fixup. No AUTO degrade: a build + # without the kernel raises in __init__ with the reason from + # cute_dsl_mla_decode_unavailability_reason. + _default_attention_backend = "CUTEDSL" + _supported_attention_backends = ("CUTEDSL",) + _auto_fallback_attention_backend = None + + # The context KV comes out of the manager's pool (5760 B/token/rank) instead + # of an arena dense in max_seq_len that free_gpu_memory_fraction never bounds. + _paged_ctx_cache = True + + def __init__( + self, + draft_config, + *, + dflash_attention_backend: str = "AUTO", + helix_mapping: Optional[Mapping] = None, + ): + cfg = draft_config.pretrained_config + if dflash_attention_backend == "AUTO": + dflash_attention_backend = type(self)._resolve_auto_attention_backend() + type(self).check_valid_attention_backend(dflash_attention_backend) + # Before the backbone is built: attention keeps the sharding's CP mapping. + helix = _shard_context_kv_by_helix(draft_config, helix_mapping) + with _builds_k3_backbone(cfg): + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + # The TP-sharded fc, resolved on first use. + self._fc_shard = None + # The helix page sharding of the latent cache, or None. + self._helix_mapping = helix + self._init_dspark_heads(cfg) + + attn0 = self.model.layers[0].self_attn + self.qk_nope_head_dim = attn0.qk_nope_head_dim + self.qk_rope_head_dim = attn0.qk_rope_head_dim + self.kv_lora_rank = attn0.kv_lora_rank + self.v_head_dim = attn0.v_head_dim + # The cache contract the worker sizes its arena from. Config-derived, so + # it is published here rather than waiting for the lazy buffer build: + # one MLA latent per token per layer, no V half. + self._num_attn_layers = len(self.model.layers) + self._num_heads = attn0.num_heads_tp + self._num_kv_heads = 1 + self._head_dim = attn0.kv_cache_head_dim + self._kv_factor = 1 + + # No slack: the lookahead lives in dflash_position_ceiling, which the + # worker publishes as _runtime_position_ceiling before any forward and + # _mla_freqs_cis builds the table from. This config-derived cap is only + # the fallback for direct construction, where no worker runs and the + # positions asked for are whatever the caller passes. + rope = _resolve_dspark_mla_rope_params(cfg, draft_config) + self._mla_rope_params = rope + self._mla_workspace = None + self._mla_freqs = None + self._mla_freqs_cap = None + self._cute_dsl_checked_shape = None + # (fixup, plan) of the last _mla_split_kv_plan. + self._mla_split_plan = None + # HF DeepSeek folds YaRN's attention scaling into the softmax scale: + # softmax_scale = mscale^2 / sqrt(qk_nope + qk_rope). NOT 1/sqrt(576) -- + # the absorbed product reproduces the 192-dim un-absorbed score. + mscale = _yarn_get_mscale(rope["scaling_factor"], rope["mscale_all_dim"]) + self.softmax_scale = (mscale * mscale) / math.sqrt( + self.qk_nope_head_dim + self.qk_rope_head_dim + ) + + # -- shape / buffers --------------------------------------------------- + + def _validate_gqa_shape(self): + """Replace the base's GQA precondition with the MLA one.""" + layers = getattr(self.model, "layers", None) + if not layers: + return + attn0 = layers[0].self_attn + keys = ( + "num_heads_tp", + "qk_nope_head_dim", + "qk_rope_head_dim", + "kv_lora_rank", + "v_head_dim", + ) + base = tuple(getattr(attn0, k) for k in keys) + for idx, layer in enumerate(layers[1:], start=1): + attn = getattr(layer, "self_attn", None) + if attn is None or not hasattr(attn, "kv_a_proj_with_mqa"): + raise ValueError( + f"MLA DSpark block decode needs self_attn.kv_a_proj_with_mqa on every " + f"draft layer, but layer {idx} of {type(self.config).__name__} has none." + ) + if tuple(getattr(attn, k) for k in keys) != base: + raise ValueError( + "MLA DSpark fuses the latent projection across layers and needs one " + f"uniform MLA shape, but layer {idx} differs from layer 0 " + f"({dict(zip(keys, base))})." + ) + + def _init_rope(self): + """No shared RoPE cache: the MLA path rotates only the rope slice.""" + self._rope_initialized = True + + @property + def reads_helix_sharded_ctx_pages(self) -> bool: + """The CuTe DSL decode reads rank-local pages and returns softmax stats; + ``_shard_context_kv_by_helix`` shards the cache for it.""" + return self._helix_mapping is not None + + def reads_managed_ctx_pages(self, tokens_per_block: int) -> bool: + """The context is paged (``_paged_ctx_cache``); the CuTe DSL decode takes a page size + that divides 128, above 1.""" + return tokens_per_block > 1 and 128 % tokens_per_block == 0 + + @classmethod + def _attention_backend_unavailability_reason(cls, backend: str) -> Optional[str]: + """Probe THIS family's kernel, not the base's GQA op set (never loaded here).""" + if backend == "CUTEDSL": + return cute_dsl_mla_decode_unavailability_reason() + return None + + def _mla_decode_workspace(self, device, min_bytes: int = 0): + want = max(_MLA_DECODE_WORKSPACE_BYTES, int(min_bytes)) + if self._mla_workspace is None or self._mla_workspace.numel() < want: + self._mla_workspace = torch.zeros(want, dtype=torch.uint8, device=device) + return self._mla_workspace + + def _assert_cute_dsl_can_implement(self, batch, block_size, page_size, dtype) -> None: + """Reject a shape the kernel cannot serve, here, naming the reason. + + Without this the rejection is silent until it is not: ``get_valid_tactics`` + returns [] (cute_dsl_custom_ops.py:8508), the AutoTuner falls back to its -1 + sentinel, and ``forward`` compiles an unvalidated config -- so a draft pool + with ``tokens_per_block=256`` (``128 % 256 != 0``, mla_decode_fp16.py:3798) + surfaces as a CUTLASS error out of the fifth drafter layer. + + Page size is the only operand here that is not fixed by the checkpoint, + and it comes from the draft KV pool, so this cannot move to __init__. + """ + key = (batch, block_size, page_size, dtype) + if self._cute_dsl_checked_shape == key: + return + import cutlass + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLNVMlaDecodeBlackwellRunner, + ) + + cutlass_dtype = cutlass.Float16 if dtype == torch.float16 else cutlass.BFloat16 + runner = CuteDSLNVMlaDecodeBlackwellRunner + # Mirror the runner's own probe (cute_dsl_custom_ops.py:8542) through its + # own constants, so this cannot drift from the tactic it will pick. Its + # candidate tiler list has one entry; split_kv=1 is always in range. + ok = runner._KERNEL_CLASS_BY_DTYPE[cutlass_dtype].can_implement( + batch, + block_size, + page_size, + self._num_heads, + self.kv_lora_rank, + self.qk_rope_head_dim, + cutlass_dtype, + cutlass_dtype, + cutlass.Float32, + cutlass.Float32, + (128, 128), + (128, 256), + 1, + False, + runner._IS_VAR_SEQ, + runner._IS_VAR_SPLIT_KV, + page_size, + ) + if not ok: + raise ValueError( + f"attention_backend='CUTEDSL' cannot serve this shape: batch={batch}, " + f"block_size={block_size}, page_size={page_size}, " + f"num_heads={self._num_heads}, dtype={dtype}." + ) + self._cute_dsl_checked_shape = key + + def _mla_split_kv_plan( + self, fixup, block_size, stream: Optional[torch.cuda.Stream] = None + ) -> Optional[torch.Tensor]: + """Per-request split-KV counts for the block decode of ``fixup``'s step. + + The decode kernel's grid holds the uniform split of its tactic, sized + from the batch alone; the plan lets each request use as many of those + slots as its KV length is worth (:func:`_launch_mla_split_plan`, the + decode's planner kernel at this decode's piece cost). Every + layer of a step decodes the same lengths, so the first layer launches + the plan and the others reuse it; a new fixup -- the next step, or a + CUDA graph capture -- plans again. None (the uniform split) for a + batch beyond the planner's request capacity. ``stream``: + launch the plan there (the output is allocated on the current stream; + the caller joins ``stream`` before a decode reads it). + """ + cached = self._mla_split_plan + if cached is not None and cached[0] is fixup: + return cached[1] + from .kernels.spec.mla_split_plan import _SPLIT_PLAN_MAX_REQUESTS + + seq_lens = fixup.seq_lens_i32 + batch = seq_lens.shape[0] + plan = None + if batch <= _SPLIT_PLAN_MAX_REQUESTS: + plan = torch.empty(batch, dtype=torch.int32, device=seq_lens.device) + with torch.cuda.stream(stream) if stream is not None else contextlib.nullcontext(): + _launch_mla_split_plan( + seq_lens, + plan, + num_heads=self._num_heads, + seq_len_q=block_size, + piece_cost=_MLA_SPLIT_PLAN_PIECE_COST_TILES, + ) + self._mla_split_plan = (fixup, plan) + return plan + + def _cute_dsl_block_decode(self, query, layer_cache, fixup, block_size): + """One cute-dsl MLA decode over context AND block. + + The kernel takes ``kv_bounds``, a per-query-token KV length REPLACING the + implicit causal bound (mla_decode_fp16.py:400-407), so one pass sees the + whole block from every block token. + + No fallback: what the kernel cannot serve raises. Build-level + availability is checked at construction through + ``cute_dsl_mla_decode_unavailability_reason``; the shape depends on the + draft pool's page size, which is not known until here, so it is checked + on the first call. + + Each request splits its KV as the step's plan says + (``_mla_split_kv_plan``); for a given split the result is the uniform + launch's bit for bit. Under helix also returns the softmax stats. + """ + import cutlass + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLNVMlaDecodeBlackwellRunner, + ) + + batch, _, num_heads, _ = query.shape + device = query.device + pool = layer_cache[:, 0, 0] # [pages, page_size, head_dim] + page_size = pool.shape[1] + d_latent = self.kv_lora_rank + self._assert_cute_dsl_can_implement(batch, block_size, page_size, query.dtype) + + # stride[1] == 1 is the kernel's only layout demand (mla_decode_fp16.py:465-470), + # kept by a 576-wide row sliced at 512. detach() because from_dlpack refuses + # grad-requiring tensors, which the pool write adds outside inference_mode. + kernel_query = query.detach() + q_latent = kernel_query[..., :d_latent].permute(2, 3, 1, 0) + q_rope = kernel_query[..., d_latent:].permute(2, 3, 1, 0) + c_latent = pool.detach()[..., :d_latent].permute(1, 2, 0) + c_rope = pool.detach()[..., d_latent:].permute(1, 2, 0) + + out = torch.empty(batch, block_size, num_heads, d_latent, dtype=query.dtype, device=device) + # Helix: (max, sum) per query token and head for the merge of the rank partials. + stats = ( + None + if fixup.owned is None + else torch.empty(batch * block_size, num_heads, 2, dtype=torch.float32, device=device) + ) + workspace_numel = CuteDSLNVMlaDecodeBlackwellRunner.get_max_padded_workspace_size( + num_heads, block_size, d_latent, batch, cutlass.Float32 + ) + workspace = self._mla_decode_workspace(device, workspace_numel) + torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell( + q_latent, + q_rope, + c_latent, + c_rope, + # .t() WITHOUT .contiguous(): the kernel wants (max_blocks, B) with dim 0 + # stride-1. contiguous() gives stride (B, 1) and it raises "Expected + # strides[leading_dim] == 1, but got B" for B > 1 (probe job 3005553). + fixup.page_tables_i32.t(), + fixup.seq_lens_i32, + out.permute(2, 3, 1, 0), + workspace, + num_heads, + block_size, + page_size, + self.softmax_scale, + 1.0, + batch, + stats, + # Every block token sees the whole block, not just its causal prefix. + fixup.kv_bounds_i32, + self._mla_split_kv_plan(fixup, block_size), + ) + return out if stats is None else (out, stats) + + def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size): + """Store the block's latents into the pool at ``ctx_len + j`` (see _MLABlockFixup), + then decode context and block in one pass (``_cute_dsl_block_decode``). + + The store precedes the decode, which reads the block back out of the pool. + Returns the latent-space output ``[B, block, heads, kv_lora_rank]`` (and the + softmax stats under helix). + """ + pool = layer_cache[:, 0, 0] # [pages, page_size, head_dim] + pages, offsets = fixup.pages.reshape(-1), fixup.offsets.reshape(-1) + blk_rows = blk_kv.reshape(-1, blk_kv.shape[-1]) + if fixup.owned is not None: + # Helix: positions on other ranks' pages rewrite their local slot. + pool[pages, offsets] = torch.where( + fixup.owned.reshape(-1, 1), blk_rows, pool[pages, offsets] + ) + elif spec_tail_fusion_enabled() and blk_kv.is_contiguous(): + k3_scatter_rows_to_pool(blk_kv.view(-1, blk_kv.shape[-1]), pool, pages, offsets) + else: + pool[pages, offsets] = blk_rows + return self._cute_dsl_block_decode(query, layer_cache, fixup, block_size) + + def _mla_freqs_cis(self, device): + """Adjacent-pair YaRN table, cached. See build_dspark_mla_yarn_freqs_cis. + + Sized from the ceiling the worker publishes once it knows the runtime one + (_runtime_position_ceiling, set in K3DFlashWorker._lazy_init_ctx_buffers before + any forward), because that is what ctx_len is clamped to and it exceeds + model_config.max_seq_len whenever spec decoding is on. The config cap is the + fallback for direct construction in tests. + """ + rope = dict(self._mla_rope_params) + runtime_cap = self._runtime_position_ceiling + if runtime_cap is not None: + # Outright, not max(): the published ceiling is already + # min(runtime, max_position_embeddings) + lookahead, so it is + # both sufficient and the only thing that keeps the table small + # when max_seq_len is unset and the config cap is the + # checkpoint's advertised 1,048,576. + rope["max_positions"] = int(runtime_cap) + # Keyed on the cap, not just built once: _lazy_init_ctx_buffers + # re-publishes the ceiling when the KV-estimation probe manager is + # swapped for the real one, and drafter forwards do run during + # estimation. A table built short for a later, larger cap would index + # out of range on freqs[positions]. + if self._mla_freqs is None or self._mla_freqs_cap != rope["max_positions"]: + self._mla_freqs = build_dspark_mla_yarn_freqs_cis( + dim=self.qk_rope_head_dim, + base=rope["theta"], + scaling_factor=rope["scaling_factor"], + original_max_position_embeddings=rope["original_max_positions"], + max_position_embeddings=rope["max_positions"], + beta_fast=rope["beta_fast"], + beta_slow=rope["beta_slow"], + mscale=rope["mscale"], + mscale_all_dim=rope["mscale_all_dim"], + device=device, + ) + self._mla_freqs_cap = rope["max_positions"] + return self._mla_freqs + + def _mla_freqs_real(self, device) -> torch.Tensor: + """``_mla_freqs_cis`` as a ``[positions, rope_dim // 2, 2]`` FP32 (cos, sin) view.""" + freqs = self._mla_freqs_cis(device) + return torch.view_as_real(freqs).reshape(freqs.shape[0], freqs.shape[-1], 2) + + def _mla_tail_v2_supported(self, device) -> bool: + """Whether the shapes fit the fused MLA RoPE kernels of ``TLLM_DSPARK_TAIL_V2``. + + False without the CuTe DSL (no ``k3_dspark_mla_block_rope`` / + ``k3_dspark_mla_ctx_norm_rope``). + """ + return ( + k3_dspark_mla_block_rope is not None + and k3_dspark_mla_ctx_norm_rope is not None + and IS_CUTLASS_DSL_AVAILABLE + and is_sm_100f() + and self.kv_lora_rank % 32 == 0 + and (self.qk_rope_head_dim // 2) % 32 == 0 + and getattr(self, "_kv_a_norm_stacked", None) is not None + and self._kv_a_norm_stacked.dtype == torch.bfloat16 + ) + + def _build_fused_kv_buffers(self) -> None: + """Stack the per-layer latent projection and its RMSNorm weight. + + The MLA analogue of the base's fused K/V GEMM: one ``[L*576, hidden]`` + weight, plus the ``[L, kv_lora_rank]`` layernorm scales applied after. + """ + if self._fused_kv_weight is not None: + return + layers_attn = [layer.self_attn for layer in self.model.layers] + attn0 = layers_attn[0] + self._fused_kv_weight = torch.cat( + [a.kv_a_proj_with_mqa.weight for a in layers_attn], dim=0 + ).contiguous() + self._fused_kv_bias = None + self._kv_a_norm_stacked = torch.stack([a.kv_a_layernorm.weight.data for a in layers_attn]) + self._kv_a_norm_eps = attn0.kv_a_layernorm.variance_epsilon + self._input_ln_eps = None + self._has_qk_norm = False + self._use_fused_qk_norm_rope = False + + # Split kv_b_proj into the absorption operands once. K absorbs into the + # query, V un-absorbs the attention output. Block decode only. + nh = attn0.num_heads_tp + for a in layers_attn: + w = a.kv_b_proj.weight.view(nh, self.qk_nope_head_dim + self.v_head_dim, -1) + a._k_b_proj = w[:, : self.qk_nope_head_dim].contiguous() + a._v_b_proj = w[:, self.qk_nope_head_dim :].contiguous() + + logger.debug( + f"MLA DSpark: fused latent projection built for {self._num_attn_layers} layers " + f"(weight={tuple(self._fused_kv_weight.shape)}, head_dim={self._head_dim})" + ) + + # -- context cache ----------------------------------------------------- + + def precompute_context_kv(self, projected_hidden, positions): + """Post-norm / post-RoPE MLA latent for ALL drafter layers in one GEMM. + + Returns ``(latent, None)``: ``[N, L, 1, kv_lora_rank + qk_rope_head_dim]`` + and no V half, which is what ``_kv_factor == 1`` means to the worker. + """ + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + n = projected_hidden.shape[0] + num_layers = self._num_attn_layers + weight_dtype = self._fused_kv_weight.dtype + if projected_hidden.dtype != weight_dtype: + projected_hidden = projected_hidden.to(weight_dtype) + + latent = F.linear(projected_hidden, self._fused_kv_weight) + latent = latent.view(n, num_layers, self._head_dim) + + # One fused RMSNorm+RoPE per layer, not one eager chain for all of them: + # the kernel takes a single norm weight per call and the layers do not + # share one. Same op the block path uses, so the two halves of an + # attention stay numerically consistent -- which is the point here, not + # speed (this runs per request, not per decode step). + freqs = self._mla_freqs_cis(latent.device)[positions.view(-1).long()] + out = torch.empty_like(latent) + for layer_idx in range(num_layers): + # 3D throughout: the eager fallback's rotary indexes [G, s, rd]. + out[:, layer_idx] = _rmsnorm_rope_batched( + latent[:, layer_idx].unsqueeze(0).contiguous(), + self._kv_a_norm_stacked[layer_idx], + self._kv_a_norm_eps, + self.qk_rope_head_dim, + freqs.unsqueeze(0), + norm_dim=self.kv_lora_rank, + ).squeeze(0) + return out.unsqueeze(2).contiguous(), None + + def precompute_context_kv_layer_major(self, projected_hidden, positions): + """``precompute_context_kv``'s latents laid out ``[L, N, head_dim]``. + + Same GEMM, same frequency gather and the same fused RMSNorm+RoPE kernel + per layer, but the kernel reads its layer's slice of the GEMM output in + place and writes straight into a layer-major buffer, so no per-layer + staging copies run. Returns None when the fused kernel is unavailable + (the caller then uses ``precompute_context_kv``). + """ + if not (IS_CUTLASS_DSL_AVAILABLE and is_sm_100f()): + return None + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + n = projected_hidden.shape[0] + num_layers = self._num_attn_layers + weight_dtype = self._fused_kv_weight.dtype + if projected_hidden.dtype != weight_dtype: + projected_hidden = projected_hidden.to(weight_dtype) + latent = draft_gemm(projected_hidden, self._fused_kv_weight, "ctx_latent") + if latent is None: + latent = F.linear(projected_hidden, self._fused_kv_weight) + latent = latent.view(n, num_layers, self._head_dim) + if tail_v2_enabled() and self._mla_tail_v2_supported(latent.device): + # Every layer in one launch, the RoPE table read by position in the + # kernel; bitwise the per-layer launches below. + return k3_dspark_mla_ctx_norm_rope( + latent, + self._kv_a_norm_stacked, + positions.view(-1).long(), + self._mla_freqs_real(latent.device), + self._kv_a_norm_eps, + ) + freqs = self._mla_freqs_cis(latent.device)[positions.view(-1).long()] + freqs_real = torch.view_as_real(freqs).reshape(-1, freqs.shape[-1], 2) + if not is_fused_dspark_rmsnorm_rope_supported( + latent[:, 0], + self._kv_a_norm_stacked[0], + freqs_real, + 1, + self.qk_rope_head_dim, + self.kv_lora_rank, + ): + return None + compiled = build_rmsnorm_rope_into( + self._head_dim, + self.qk_rope_head_dim, + 1, + self._kv_a_norm_eps, + True, + True, + False, + self.kv_lora_rank, + ) + out = torch.empty((num_layers, n, self._head_dim), dtype=latent.dtype, device=latent.device) + for layer_idx in range(num_layers): + # Each layer's slice of the GEMM output is read in place. + k3_dspark_rmsnorm_rope_into( + latent[:, layer_idx], + self._kv_a_norm_stacked[layer_idx], + freqs_real, + out[layer_idx], + compiled, + ) + return out + + # -- block decode ------------------------------------------------------ + + def dflash_forward( + self, + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_k_cache, + ctx_v_cache, + ctx_cache_batch_idx, + ctx_kv_cache=None, + ctx_page_table=None, + block_placement: Optional[DraftBlockPlacement] = None, + block_fixup: Optional[_MLABlockFixup] = None, + ): + """Absorbed-MLA block decode of every drafter layer over the paged latent cache. + + Args mirror the base: ``ctx_kv_cache`` holds each layer's ``[pages, 1, 1, page, 576]`` + latent pool and ``ctx_page_table`` the batch's pages (``ctx_k_cache`` / ``ctx_v_cache`` + are unused). Either worker may hand over the block's placement: ``block_placement`` + (main's worker) is where the worker already placed the block in the pages; + ``block_fixup`` (the K3 worker's fused prepare) is the decode operands outright, and wins + when both are given. + + Per layer: input RMSNorm (fused with the previous layer's MLP all-reduce under TP), + ``[q_a | kv_a]`` GEMM (or the two), q_a RMSNorm, q_b GEMM, query RoPE and the + ``W_UK`` absorption, block latent RMSNorm + RoPE and its store into the pool, the cute-dsl + decode, the helix merge, the ``W_UV`` up-projection, o_proj with its all-reduce + residual + + RMSNorm, and the MLP. Then the final RMSNorm. + """ + del ctx_k_cache, ctx_v_cache + if ctx_page_table is None: + raise NotImplementedError( + "The MLA DSpark block decode reads the draft KV manager's paged latent cache " + "(_paged_ctx_cache); this context has no page table." + ) + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + + batch, block_size = noise_embedding.shape[0], noise_embedding.shape[1] + device = noise_embedding.device + num_heads = self._num_heads + latent_dim = self.kv_lora_rank + # draft3_glue: per-step glue off the critical path (same values). + glue = draft3_glue_enabled() + # A caller-built fixup carries what the decode reads; the slots are not read then. + slots = None if glue and block_fixup is not None else ctx_cache_batch_idx.long() + helix = self._helix_mapping + # dspark_tail_v2: per layer, one kernel rotates the query's rope half, + # normalizes and rotates the block latents and stores them in the pool + # (reading the RoPE table by position), and the cute-dsl decode follows + # directly. Same values as the separate launches. Its store writes every + # block position, so a helix-sharded cache, where other ranks' positions + # must keep their slot's value, takes the separate launches, unless + # ull2_draft_store stores the rows through scratch with the rank's + # ownership mask. + fused_block_rope = ( + tail_v2_enabled() + and (helix is None or draft_helix_store_enabled()) + and self._mla_tail_v2_supported(device) + ) + if fused_block_rope: + blk_freqs = None + freqs_real = self._mla_freqs_real(device) + flat_positions = query_positions.reshape(-1).long() + else: + # [B, blk, rope/2] complex: one phase per draft position, shared by + # the query and the block latent of every layer. + blk_freqs = self._mla_freqs_cis(device)[query_positions.long()] + # Context pages for this batch, resolved once: rows are batch positions + # of the manager's per-request block table, which is exactly what + # ctx_cache_batch_idx is. + page_tables = None if block_fixup is not None else ctx_page_table.index_select(0, slots) + page_size = ctx_kv_cache[0].shape[-2] + if helix is not None and block_size > page_size: + raise NotImplementedError( + f"Helix-sharded MLA DSpark context needs block_size <= page size; " + f"got {block_size} > {page_size}." + ) + if ctx_kv_cache[0].dtype != noise_embedding.dtype: + raise ValueError( + f"The MLA DSpark block decode reads a {noise_embedding.dtype} latent pool, " + f"got {ctx_kv_cache[0].dtype}." + ) + # Under helix the fused RoPE kernel writes the block latents to scratch + # rows, and one scatter stores the positions this rank's pages hold. + block_scratch = fused_block_rope and helix is not None + if block_scratch: + num_block_tokens = batch * block_size + blk_scratch = torch.empty( + (num_block_tokens, 1, self._head_dim), dtype=torch.bfloat16, device=device + ) + scratch_rows, scratch_offsets = self._block_scratch_index( + num_block_tokens, device, cached=glue + ) + # The decode operands are layer-invariant: built once here. + if block_fixup is not None: + if helix is not None and block_fixup.owned is None: + raise ValueError("Helix-sharded MLA DSpark needs a rank-local block fixup.") + fixup = block_fixup + elif block_placement is not None and helix is None: + # Main's worker placed the block already; under helix its global + # placement is not this rank's, so the fixup is rebuilt below. + fixup = _MLABlockFixup( + pages=block_placement.pages, + offsets=block_placement.offsets, + valid=_block_upper_triangle(block_size, device), + page_tables_i32=page_tables, + seq_lens_i32=block_placement.seq_lens, + kv_bounds_i32=block_placement.kv_bounds, + ) + else: + fixup = _build_mla_block_fixup( + num_ctx_per_req[:batch], + page_tables.long(), + block_size, + page_size, + None if helix is None else HelixDraftShard(helix.cp_size, helix.cp_rank, page_size), + ) + logger.info_once( + "MLA DSpark block decode: one cute-dsl pass over context and block, no fixup", + key="mla_dspark_block_decode_variant", + ) + # Helix: the tokens this rank holds no key of, layer-invariant. Built + # once here when an ull2_draft_* feature is on, else per layer. + helix_zero_kv = ( + fixup.kv_bounds_i32 == 0 + if helix is not None and (draft_helix_a2a_enabled() or draft_helix_store_enabled()) + else None + ) + # glue: the first decode's split-KV plan on the drafter's side stream, beside + # layer 0's projections, behind the context store the worker may have put + # there (draft3_side_launch); joined before layer 0's first pool write. + plan_stream = self._draft3_side_stream() if self._draft3_side_keep else None + if glue: + plan_stream = self._draft3_side_stream() + plan_stream.wait_stream(torch.cuda.current_stream()) + self._mla_split_kv_plan(fixup, block_size, stream=plan_stream) + hidden_states = noise_embedding + residual = None + # A TP-sharded drafter over MNNVL: o_proj and down_proj return their + # un-reduced partials and each reduction is fused with the residual add + # and the following RMSNorm (post_attention_layernorm, the next layer's + # input_layernorm, or the final norm). + fused_ar = self._draft_fused_all_reduce() + # o_proj and down_proj push their partials into that all-reduce and one + # site kernel finishes it with the residual and norm, where a site applies. + sites = self._draft_ar_sites(fused_ar) + # The previous layer's un-reduced MLP output. + partial = None + for layer_idx, layer in enumerate(self.model.layers): + attn = layer.self_attn + if partial is None: + hs_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + if residual is None: + # glue: the noise embedding is this forward's own buffer, read only + # here; the pushed sites write new residual tensors. + residual = hidden_states if glue and sites is not None else hidden_states.clone() + hs_normed = layer.input_layernorm(hs_flat) + elif partial is not None: + res_flat = residual.reshape(-1, residual.shape[-1]) + hs_normed, res_flat = self._reduce_partial( + fused_ar, sites, partial, res_flat, layer.input_layernorm + ) + residual = res_flat.reshape(batch, block_size, -1) + else: + res_flat = residual.reshape(-1, residual.shape[-1]) + hs_normed, res_flat = layer.input_layernorm(hs_flat, res_flat) + residual = res_flat.reshape(batch, block_size, -1) + + fused_a = self._draft_fused_a_proj(layer_idx) if fused_ar is not None else None + if fused_a is not None: + # q_a and kv_a read the same normed input: one GEMM, split after. + # Both halves are consumed as row-strided views (q_a_layernorm + # and the fused RMSNorm+RoPE read strided rows), so no copies. + q_kv_a = draft_gemm(hs_normed, fused_a, "q_kv_a") + if q_kv_a is None: + q_kv_a = bf16_linear(hs_normed, fused_a) + q_a = q_kv_a[:, : attn.q_lora_rank] + blk_latent_rows = q_kv_a[:, attn.q_lora_rank :] + else: + q_a = attn.q_a_proj(hs_normed) + blk_latent_rows = None + q_a_normed = attn.q_a_layernorm(q_a) + q = draft_gemm(q_a_normed, attn.q_b_proj.weight, "q_b") + if q is None: + q = attn.q_b_proj(q_a_normed) + q = q.view(batch, block_size, num_heads, self.qk_nope_head_dim + self.qk_rope_head_dim) + if fused_block_rope: + query = torch.empty( + (batch, block_size, num_heads, self.kv_lora_rank + self.qk_rope_head_dim), + dtype=q.dtype, + device=q.device, + ) + if blk_latent_rows is None: + blk_latent_rows = attn.kv_a_proj_with_mqa(hs_normed) + num_tokens = batch * block_size + kv_pool = ctx_kv_cache[layer_idx][:, 0, 0] + if plan_stream is not None and not block_scratch: + plan_stream = self._draft3_join_side() + k3_dspark_mla_block_rope( + q.view(num_tokens, num_heads, q.shape[-1]), + query.view(num_tokens, num_heads, query.shape[-1]), + blk_latent_rows, + attn.kv_a_layernorm.weight, + flat_positions, + freqs_real, + blk_scratch if block_scratch else kv_pool, + scratch_rows if block_scratch else fixup.pages.reshape(-1), + scratch_offsets if block_scratch else fixup.offsets.reshape(-1), + attn.kv_a_layernorm.variance_epsilon, + ) + if plan_stream is not None: + plan_stream = self._draft3_join_side() + if block_scratch: + k3_scatter_rows_to_pool( + blk_scratch.view(num_tokens, -1), + kv_pool, + fixup.pages.reshape(-1), + fixup.offsets.reshape(-1), + rows_mask=fixup.owned, + ) + bmm_out( + q[..., : self.qk_nope_head_dim] + .view(num_tokens, num_heads, self.qk_nope_head_dim) + .transpose(0, 1), + attn._k_b_proj.to(q.dtype), + query[..., : self.kv_lora_rank] + .view(num_tokens, num_heads, self.kv_lora_rank) + .transpose(0, 1), + ) + ctx_out = self._cute_dsl_block_decode( + query, ctx_kv_cache[layer_idx], fixup, block_size + ) + else: + # One buffer rather than cat([absorbed, rope]): q's rope half is a 64-slice + # of a 192-wide row, so cat fell off CatArrayBatchedCopy onto the generic + # kernel -- 17.3 us vs 1.6 us, five launches/iter (measured job 2986420). + query = torch.empty( + (batch, block_size, num_heads, self.kv_lora_rank + self.qk_rope_head_dim), + dtype=q.dtype, + device=q.device, + ) + # Rotate the trailing rope slice of every head straight into the + # destination's rope columns. No norm and no weight here -- + # q_a_layernorm already ran. + q_rope_into = _mla_rope_into( + q, blk_freqs, query, num_heads, self.qk_rope_head_dim, self.kv_lora_rank + ) + if not q_rope_into: + q = _rmsnorm_rope_batched( + q, + None, + 0.0, + self.qk_rope_head_dim, + blk_freqs, + num_heads=num_heads, + apply_weight=False, + apply_rmsnorm=False, + ) + query[..., self.kv_lora_rank :] = q[..., self.qk_nope_head_dim :] + q_nope = q[..., : self.qk_nope_head_dim] + # Absorb the nope half into latent space so the block attends the stored + # latent directly (MQA) instead of expanding per head. view(), never + # reshape(): both are strided views and reshape would stage a silent copy. + bmm_out( + q_nope.view(batch * block_size, num_heads, self.qk_nope_head_dim).transpose( + 0, 1 + ), + attn._k_b_proj.to(q.dtype), + query[..., : self.kv_lora_rank] + .view(batch * block_size, num_heads, self.kv_lora_rank) + .transpose(0, 1), + ) + + # The block's own latent comes from the draft hidden; the context's + # from the projected target hidden (precompute_context_kv). That + # dual-source split is the DFlash contract, not an MLA detail. + if blk_latent_rows is None: + blk_latent_rows = attn.kv_a_proj_with_mqa(hs_normed) + blk_latent = blk_latent_rows.view(batch, block_size, -1) + # RMSNorm over the latent, RoPE over k_pe, one kernel. norm_dim keeps + # k_pe out of the reduction, which is what DeepSeek-style MLA wants + # and what DSpark's own convention (whole row) does not do. + blk_kv = _rmsnorm_rope_batched( + blk_latent, + attn.kv_a_layernorm.weight, + attn.kv_a_layernorm.variance_epsilon, + self.qk_rope_head_dim, + blk_freqs, + norm_dim=latent_dim, + ) + if plan_stream is not None: + plan_stream = self._draft3_join_side() + ctx_out = self._mla_paged_attention( + query, blk_kv, ctx_kv_cache[layer_idx], fixup, block_size + ) + stats = None + merged = None + if helix is not None: + ctx_out, stats = ctx_out + zero_kv = helix_zero_kv if helix_zero_kv is not None else fixup.kv_bounds_i32 == 0 + # ull2_draft_a2a: V up-projection, exchange and combine in one launch. + merged = draft_helix_tail(helix, ctx_out, stats, attn._v_b_proj, zero_kv) + if merged is not None: + attn_out = merged + elif spec_tail_fusion_enabled(): + # Same batched GEMM as the einsum, written straight into the + # [tokens, heads, v] rows o_proj reads, so no reshape copy. + attn_out = torch.empty( + (batch * block_size, num_heads, self.v_head_dim), + dtype=ctx_out.dtype, + device=ctx_out.device, + ) + bmm_out( + ctx_out.reshape(batch * block_size, num_heads, -1).transpose(0, 1), + attn._v_b_proj.to(ctx_out.dtype).transpose(1, 2), + attn_out.transpose(0, 1), + ) + else: + attn_out = torch.einsum("bshc,hvc->bshv", ctx_out, attn._v_b_proj.to(ctx_out.dtype)) + if helix is not None and merged is None: + # V is linear, so the per-rank partials merge after it: the + # all-to-all moves v_head_dim (128), not kv_lora_rank (512). + attn_out = _draft_helix_post_process( + attn_out.reshape(batch * block_size, -1), + stats, + helix, + attn.num_heads_tp_cp, + self.v_head_dim, + zero_kv, + ) + + res_flat = residual.reshape(-1, residual.shape[-1]) + if fused_ar is not None: + no_reduce = AllReduceParams(enable_allreduce=False) + site = None + if sites is not None: + site = sites.linear_ar_norm( + attn_out.reshape(batch * block_size, -1), + attn.o_proj.weight, + "o_proj", + res_flat, + layer.post_attention_layernorm, + ) + if site is not None: + hidden_out, res_flat = site + else: + hidden_out = attn.o_proj( + attn_out.reshape(batch * block_size, -1), all_reduce_params=no_reduce + ) + hidden_out, res_flat = self._fused_ar_residual_norm( + fused_ar, hidden_out, res_flat, layer.post_attention_layernorm + ) + partial = _draft_mlp(layer.mlp, hidden_out, no_reduce, defer_down=sites is not None) + hidden_states = partial + else: + hidden_out = attn.o_proj(attn_out.reshape(batch * block_size, -1)) + hidden_out, res_flat = layer.post_attention_layernorm(hidden_out, res_flat) + hidden_out = layer.mlp(hidden_out) + hidden_states = hidden_out.reshape(batch, block_size, -1) + residual = res_flat.reshape(batch, block_size, -1) + + res_flat = residual.reshape(-1, residual.shape[-1]) + if partial is not None: + hidden_states_out, _ = self._reduce_partial( + fused_ar, sites, partial, res_flat, self.model.norm + ) + else: + hidden_states_out, _ = self.model.norm( + hidden_states.reshape(-1, hidden_states.shape[-1]), res_flat + ) + if plan_stream is not None: + self._draft3_join_side() + return hidden_states_out + + def _block_scratch_index(self, num_tokens: int, device, cached: bool): + """Pages ``arange(num_tokens)`` and offsets ``0`` of the block scratch rows (a + ``[num_tokens, 1, D]`` pool view). ``cached``: built once per shape outside graph + capture and reused (they depend on the shape only).""" + cache = self.__dict__.setdefault("_draft3_scratch_index", {}) + key = (num_tokens, device) + if cached and key in cache: + return cache[key] + rows = torch.arange(num_tokens, device=device) + index = (rows, torch.zeros_like(rows)) + if cached and not torch.cuda.is_current_stream_capturing(): + cache[key] = index + return index + + def _draft3_side_stream(self) -> torch.cuda.Stream: + stream = self.__dict__.get("_draft3_stream") + if stream is None: + stream = self.__dict__["_draft3_stream"] = torch.cuda.Stream() + return stream + + @property + def _draft3_side_keep(self) -> list: + return self.__dict__.setdefault("_draft3_keep", []) + + def draft3_side_launch(self, launch: Callable[[], None], keep=()) -> None: + """Run ``launch()`` on the drafter's side stream after the current stream's work + (``draft3_glue``). The next ``dflash_forward`` joins the side stream before its first + draft-pool write; ``keep`` (the tensors the launch reads) stays referenced until then, so + the caching allocator does not hand their memory to the current stream while the side + stream may still read it.""" + stream = self._draft3_side_stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + launch() + self._draft3_side_keep.append(tuple(keep)) + + def _draft3_join_side(self) -> None: + """The current stream waits for the side stream; the kept tensors are released.""" + torch.cuda.current_stream().wait_stream(self._draft3_side_stream()) + self._draft3_side_keep.clear() + return None + + def _draft_fused_all_reduce(self) -> Optional[TpAllReduce]: + """The all-reduce the block decode fuses with residual + RMSNorm, or None. + + For a TP-sharded drafter whose o_proj reduces over MNNVL (the fusion the + MNNVL all-reduce implements). + """ + all_reduce = self.model.layers[0].self_attn.o_proj.all_reduce + if all_reduce is None or all_reduce.mnnvl_allreduce is None: + return None + return all_reduce + + def _draft_ar_sites(self, fused_ar) -> Optional[DraftArSites]: + """The pushed all-reduce sites on the fused all-reduce, or None. + + The sites push into the Lamport buffers of the all-reduce's MNNVL + workspace (``layers.lamport_push_target``); a call a site does not + take runs the fused all-reduce. + """ + if fused_ar is None: + return None + sites = self.__dict__.get("_ar_sites") + if sites is None or sites.all_reduce is not fused_ar: + sites = self.__dict__["_ar_sites"] = DraftArSites(fused_ar) + return sites + + def _reduce_partial(self, fused_ar, sites, partial, residual, norm): + """``norm(all_reduce(partial) + residual)`` and the new residual, where + ``partial`` is a down projection's output or a ``_PendingDown`` that a + pushed site projects and reduces itself.""" + if isinstance(partial, _PendingDown): + site = None + if sites is not None: + site = sites.linear_ar_norm( + partial.h2, partial.mlp.down_proj.weight, "down", residual, norm + ) + if site is not None: + return site + partial = partial.partial() + return self._fused_ar_residual_norm(fused_ar, partial, residual, norm) + + def _draft_fused_a_proj(self, layer_idx: int) -> Optional[torch.Tensor]: + """The ``[q_a_proj; kv_a_proj_with_mqa]`` weight of a layer, for one GEMM. + + Both replicated projections read the same normed input, so one GEMM + replaces two. The separate projections' weights become row views of the + fused weight, so no memory is added. Built once, outside graph capture + (None if the first call is under capture). + """ + cache = self.__dict__.setdefault("_fused_a_weights", {}) + if layer_idx in cache: + return cache[layer_idx] + attn = self.model.layers[layer_idx].self_attn + q_a, kv_a = attn.q_a_proj, attn.kv_a_proj_with_mqa + fused = None + if q_a.weight.dtype == kv_a.weight.dtype and not torch.cuda.is_current_stream_capturing(): + fused = torch.cat([q_a.weight.data, kv_a.weight.data], dim=0).contiguous() + q_a.weight.data = fused[: q_a.out_features] + kv_a.weight.data = fused[q_a.out_features :] + cache[layer_idx] = fused + return fused + + @staticmethod + def _fused_ar_residual_norm(all_reduce, partial, residual, norm): + """``norm(all_reduce(partial) + residual)`` and the new residual, in one kernel.""" + return _all_reduce_add_norm(all_reduce, partial, residual, norm) + + # -- weights ----------------------------------------------------------- + + #: Only the head is shared. This drafter ships its own embed_tokens whose + #: values differ from the target's -- see load_weights_from_target_model -- + #: so a checkpoint missing it must fail rather than run on random weights. + WEIGHTS_SHARED_WITH_TARGET = ("lm_head",) + + #: The checkpoint's TorchSpec spellings of the capture projection and the final norm + #: (the SpecForge drafters name them ``fc`` / ``hidden_norm`` / ``norm``). + _TORCHSPEC_RENAMES = { + "context_proj.weight": "fc.weight", + "context_norm.weight": "hidden_norm.weight", + "final_norm.weight": "norm.weight", + } + + def _backbone_tensors(self) -> Iterator[tuple[str, Callable[..., None], tuple[str, ...]]]: + """``(target, load, checkpoint names)`` of every backbone tensor, ``load`` taking the + named full checkpoint tensors in order.""" + + def copy_into(param): + return lambda tensor: param.data.copy_(tensor) + + backbone = self.model + yield "embed_tokens", backbone.embed_tokens.load, ("embed_tokens.weight",) + yield "norm", copy_into(backbone.norm.weight), ("norm.weight",) + for idx, layer in enumerate(backbone.layers): + prefix = f"layers.{idx}." + attn = layer.self_attn + for name in ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "o_proj"): + yield ( + f"{prefix}self_attn.{name}", + getattr(attn, name).load, + (f"{prefix}self_attn.{name}.weight",), + ) + for name in ("q_a_layernorm", "kv_a_layernorm"): + yield ( + f"{prefix}self_attn.{name}", + copy_into(getattr(attn, name).weight), + (f"{prefix}self_attn.{name}.weight",), + ) + yield ( + f"{prefix}mlp", + layer.mlp.load, + tuple(f"{prefix}mlp.{p}.weight" for p in ("gate_proj", "up_proj", "down_proj")), + ) + for name in ("input_layernorm", "post_attention_layernorm"): + yield ( + f"{prefix}{name}", + copy_into(getattr(layer, name).weight), + (f"{prefix}{name}.weight",), + ) + + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): + """Load the drafter checkpoint into the K3-owned weights, each tensor by its name. + + The DSpark heads go to :meth:`_take_dspark_head_weights`; ``fc`` / ``hidden_norm`` (or + their TorchSpec names ``context_proj`` / ``context_norm``) build the capture projection; + every backbone tensor (HF MLA names, with or without the ``model.`` prefix; the final + norm also as ``final_norm``) is copied into this rank's shard of its weight. A checkpoint + missing any of them raises. ``weight_mapper`` is not used. + """ + del weight_mapper, kwargs + weights, _ = self._take_dspark_head_weights(weights) + weights = {self._TORCHSPEC_RENAMES.get(k, k): v for k, v in weights.items()} + weights = { + k[len("model.") :] if k.startswith("model.") else k: v for k, v in weights.items() + } + + # Without fc the drafter has no capture projection and drafts from an empty + # context forever. + wrapper_missing = [k for k in self.WRAPPER_OWNED_WEIGHTS if k not in weights] + if wrapper_missing: + raise ValueError( + f"{type(self).__name__}: checkpoint is missing {wrapper_missing}, " + "which this wrapper owns and builds from the checkpoint." + ) + tensors = list(self._backbone_tensors()) + missing = sorted( + target for target, _, names in tensors if any(n not in weights for n in names) + ) + if missing: + raise ValueError( + f"{type(self).__name__}: checkpoint provides no weights for " + f"{missing[:8]}{' ...' if len(missing) > 8 else ''}. These are " + f"not in WEIGHTS_SHARED_WITH_TARGET " + f"({', '.join(self.WEIGHTS_SHARED_WITH_TARGET)}), so loading " + "would leave them uninitialized." + ) + + fc = weights["fc.weight"] + self.fc = Bf16Linear(fc.shape[1], fc.shape[0], dtype=fc.dtype, device="cuda") + self.fc.load(fc) + norm = weights["hidden_norm.weight"] + self.hidden_norm = RMSNormWeight( + norm.shape[0], getattr(self.config, "rms_norm_eps", 1e-6), norm.dtype, device="cuda" + ) + self.hidden_norm.weight.data.copy_(norm) + used = set(self.WRAPPER_OWNED_WEIGHTS) + for _, load, names in tensors: + load(*(weights[n] for n in names)) + used.update(names) + unused = sorted(set(weights) - used) + if unused: + logger.warning( + f"{type(self).__name__}: checkpoint tensors not loaded: " + f"{unused[:8]}{' ...' if len(unused) > 8 else ''}" + ) + + def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: + """Share only ``lm_head``; the drafter has its own trained embedding. + + Inferact/Kimi-K3-DSpark ships ``embed_tokens`` whose values differ from + the target's, so adopting the target's would feed the block decode + embeddings it was not distilled against -- silently, at some acceptance + cost. + """ + self.draft_model_full.lm_head = target_model.lm_head + self.lm_head = target_model.lm_head + + +def build_k3_dspark_draft(model_config, draft_config, lm_head, model) -> nn.Module: + """Build the Kimi-K3 standalone MLA DSpark drafter. + + The K3 shell routes only the standalone MLA drafter here; every other + DSpark drafter (embedded DeepSeek-V4, standalone GQA) stays with main's + registry builder. + + Args: + model_config: the target engine's ``ModelConfig``. + draft_config: the drafter's own ``ModelConfig``. + lm_head: unused; the drafter shares the target's head at weight-load time. + model: unused; kept for the draft-builder signature. + + Returns: + The :class:`K3MLADSparkForCausalLM` drafter. + + Raises: + ValueError: the draft is embedded in the target checkpoint, or its + backbone is not MLA-shaped. + """ + del lm_head, model + spec_config = model_config.spec_config + if getattr(spec_config, "draft_is_embedded_in_target", False): + raise ValueError( + "build_k3_dspark_draft builds the standalone MLA DSpark drafter only; this draft " + "is embedded in the target checkpoint (main's DSpark builder serves it)." + ) + if not is_mla(draft_config.pretrained_config): + raise ValueError( + "build_k3_dspark_draft builds the standalone MLA DSpark drafter only; this " + "drafter's backbone is not MLA-shaped (main's DSpark builder serves it)." + ) + return K3MLADSparkForCausalLM( + draft_config, + dflash_attention_backend=spec_config.attention_backend, + # The target's helix CP original, set while its modules are built on the repurposed + # mapping (KimiK3DefaultForCausalLM._setup_helix_mappings). + helix_mapping=getattr(model_config, "_helix_mapping_with_cp", None), + ) + + +__all__ = [ + "K3MLADSparkForCausalLM", + "build_k3_dspark_draft", + "K3MLADSparkBackbone", + "K3MLADSparkModel", + "build_dspark_mla_yarn_rope", + "build_dspark_mla_yarn_freqs_cis", + "apply_dspark_mla_rope", + "install_k3_dspark", + "uses_k3_dspark", +] + + +class K3DSparkWorker(K3DFlashWorker): + """Worker for a *standalone* DSpark drafter (DFlash lineage). + + DSpark is DFlash plus two extra heads, so the drafting plumbing is + inherited wholesale from :class:`K3DFlashWorker` -- paged context K/V, + slot management, the mask-token block forward -- and only the two + head-driven policies are overridden here: the block-output slot + convention (``shift_label``) and the Markov intra-block logit bias. + + Mirrors the model side, where ``GQADSparkForCausalLM`` extends + ``DFlashForCausalLM`` with the same two heads. + + Naming: this is the unqualified DSpark worker because a separately + shipped drafter is the ordinary case; :class:`DSv4DSparkWorker` carries + the qualifier because a draft embedded in the target checkpoint is the + special one. Workers are classified by *deployment form*, never by draft + backbone -- so there is no ``Qwen3DSparkWorker``. Note the name meant the + embedded worker before this split; both the rebind and the rename to + ``DSv4DSparkWorker`` land in one commit so the swap reads as a unit. + + A worker is agnostic to the draft backbone: everything backbone-shaped is + supplied by the draft model, which reports its own shapes + (``_num_attn_layers``, ``_num_heads``, ``_num_kv_heads``, ``_head_dim``) + and owns the operators (``_build_fused_kv_buffers``, + ``precompute_context_kv``, ``dflash_forward``, + ``apply_markov_chain_logits``, ``project_target_hidden``). The worker only + allocates against the reported shapes and sequences the calls. An MLA + drafter therefore reuses this class unchanged; its differences (fused-QKV + assumptions, a 576-latent K/V layout) land in its own draft-model + subclass. Naming workers by backbone would produce N classes with + identical bodies. + + Deployment form is the axis the runtime state actually splits on: paged + draft K/V here, a worker-owned rolling window in + :class:`DSv4DSparkWorker`. + """ + + def set_draft_model(self, draft_model) -> None: + """Reject an unsupported vocab mapping here rather than mid-decode. + + ``d2t`` is model-static, so a config mistake should surface at load and + not as a ``NotImplementedError`` raised per decode step, possibly during + CUDA-graph capture. + """ + super().set_draft_model(draft_model) + if self._d2t is not None and getattr(draft_model, "has_markov_head", False): + raise NotImplementedError( + "DSpark Markov head requires a shared draft/target vocab " + "(d2t vocab mapping is not supported); drafter " + f"{type(draft_model).__name__} declares one." + ) + + def _draft_block_width(self, draft_model) -> int: + """Block width under the dspark ``shift_label`` convention. + + shift_label reads slots 0..K-1, so K draft tokens fit in K slots and + the base class' K+1 over-demands by one -- enough to reject a block-7 + checkpoint at max_draft_len=7, which is how both published DSpark + drafters are meant to run. + """ + if draft_model.shift_label: + return self.max_draft_len + return super()._draft_block_width(draft_model) + + def _draft_first_slot(self, draft_model) -> int: + return 0 if draft_model.shift_label else 1 + + def _draft_slot_ids( + self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int + ) -> torch.Tensor: + """Block-output slots under the dspark ``shift_label`` convention. + + The drafter checkpoint declares the convention, so it is read off the + draft model rather than assumed: a DSpark drafter trained with the + legacy DFlash slot layout keeps the base class' slots 1..K. + """ + return dflash_draft_slot_ids( + num_gens, block_size, num_draft_tokens, draft_model.shift_label, device="cuda" + ) + + def _draft_vocab_parallel_head(self, draft_model, spec_metadata): + """Also require Markov weights the vocab-parallel chain can index: + contiguous ``[vocab, rank]`` tensors of one dtype.""" + head = super()._draft_vocab_parallel_head(draft_model, spec_metadata) + if head is None or not getattr(draft_model, "has_markov_head", False): + return head + w1, w2 = draft_model.markov_w1, draft_model.markov_w2 + if ( + w1.is_contiguous() + and w2.is_contiguous() + and w1.shape == w2.shape + and w1.dtype == w2.dtype + and w2.shape[0] == head.vocab_size + ): + return head + return None + + def _vocab_parallel_markov_weights(self, draft_model): + if not getattr(draft_model, "has_markov_head", False): + return None, None + return draft_model.markov_w1, draft_model.markov_w2 + + def _refine_block_logits( + self, + draft_model, + gen_logits: torch.Tensor, + inputs: dict, + spec_metadata, + ) -> torch.Tensor: + """Add the greedy-chained Markov intra-block bias to the block logits. + + A DSpark drafter checkpoint may omit the Markov head (``markov_rank`` + 0), which loads as a drafter without one; that case falls through to + the unmodified backbone logits. + """ + if not getattr(draft_model, "has_markov_head", False): + return gen_logits + return self._apply_dspark_markov_bias( + draft_model, gen_logits, inputs["first_prev_tokens"], spec_metadata + ) + + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata=None, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + The logits are full-vocab: the drafter shares the target's gathered + LM head (vocab-sharded selection runs the chain through + :meth:`_draft_vocab_parallel_head` instead). Full-vocab logits go through Jonas's fused chain + (``dspark_markov_fused``) where it supports them, instead of the + drafter's ``apply_markov_chain_logits``; the same chain bit for bit. + """ + # The d2t guard lives in set_draft_model: it is model-static, so raising + # it here would surface a load-time config error per decode step. + # This one keys on the runtime logits width. + full_vocab = draft_model.markov_w2.shape[0] + if gen_logits.shape[-1] != full_vocab: + raise NotImplementedError( + f"DSpark Markov head: draft logits width {gen_logits.shape[-1]} does not " + f"match the drafter vocab {full_vocab} (the shared LM head gathers its output)." + ) + if ( + spec_metadata is None or not self._draft_logits_are_sharded(gen_logits, spec_metadata) + ) and fused_markov_supported(gen_logits, draft_model.markov_w1, draft_model.markov_w2): + return k3_fused_markov_chain_logits( + gen_logits, first_prev_tokens, draft_model.markov_w1, draft_model.markov_w2 + ) + return draft_model.apply_markov_chain_logits(gen_logits, first_prev_tokens) + + +def _helix_sanitize_empty_kv( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Force zero-local-KV rows to a no-op contribution for the Helix combine. + + A CP rank that owns no KV blocks for a token attends to zero keys, so the + attention kernel normalizes by a zero softmax sum and yields NaN. The combine + weights each rank by sum * exp(max - global_max), so such rows must contribute + softmax_stats of (-inf, 0) and a zeroed partial_o to act as a no-op. The rows + are selected by zero_kv_mask, which is robust regardless of what the kernel + wrote. Passing None disables sanitization. + + Args: + partial_o: Partial attention output, shape [num_tokens, ...]. + softmax_stats: Per (token, head) (max, sum), shape [num_tokens, num_heads, 2]. + zero_kv_mask: Bool tensor of shape [num_tokens], True where this rank has + zero local KV. + """ + if zero_kv_mask is None: + return partial_o, softmax_stats + num_tokens = partial_o.shape[0] + mask = zero_kv_mask[:num_tokens].view(-1, 1) + # masked_fill overwrites masked rows regardless of their current (possibly NaN) + # value and is CUDA-graph safe due to static shapes. + partial_o = partial_o.masked_fill(mask, 0.0) + sm_max = softmax_stats[..., 0].masked_fill(mask, float("-inf")) + sm_sum = softmax_stats[..., 1].masked_fill(mask, 0.0) + softmax_stats = torch.stack([sm_max, sm_sum], dim=-1) + return partial_o, softmax_stats + + +# Distinct input specializations the NCCL reformat helpers below have been +# asked to compile. ``dynamic=False`` specializes on the input shapes and +# ``cp_size``, so this set tracks the same thing dynamo's per-frame cache does. +_HELIX_NCCL_SHAPES: set = set() +_HELIX_NCCL_LIMIT_REPORTED = False + + +def _helix_note_nccl_specialization( + partial_o: torch.Tensor, softmax_stats: torch.Tensor, cp_size: int +) -> None: + """Warn once if the Helix NCCL reformat is past dynamo's recompile budget. + + Past ``cache_size_limit`` dynamo stops compiling and runs the frame eagerly + without raising, which silently gives back the fusion these helpers exist + for. Report it once per process so the regression is visible in the log + instead of only in a kernel trace. + """ + global _HELIX_NCCL_LIMIT_REPORTED + if _HELIX_NCCL_LIMIT_REPORTED: + return + key = (tuple(partial_o.shape), tuple(softmax_stats.shape), cp_size) + if key in _HELIX_NCCL_SHAPES: + return + _HELIX_NCCL_SHAPES.add(key) + try: + import torch._dynamo as _dynamo + + limit = getattr(_dynamo.config, "cache_size_limit", None) + except ImportError: + limit = None + if limit is None or len(_HELIX_NCCL_SHAPES) <= limit: + return + _HELIX_NCCL_LIMIT_REPORTED = True + logger.warning( + "Helix NCCL all-to-all reformat has seen %d distinct input " + "specializations, above torch._dynamo.config.cache_size_limit=%d. " + "Dynamo stops recompiling past that limit and runs these helpers " + "eagerly without raising, losing the sanitize/transpose fusion. " + "Reduce the number of CUDA-graph batch buckets or raise " + "cache_size_limit.", + len(_HELIX_NCCL_SHAPES), + limit, + ) + + +# NCCL needs each all-to-all chunk to be contiguous. Inductor can pad the +# transposed buffers that the chunks are split from, even with +# ``pad_outputs=False``, so turn padding off for this helper. This keeps the +# fused kernels and avoids a copy of each chunk after the compiled call. +@torch.compile(dynamic=False, options={"comprehensive_padding": False}) +def _helix_nccl_pre_alltoall( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor], + cp_size: int, +) -> List[torch.Tensor]: + """Sanitize zero-local-KV rows and reformat into the alltoall send layout. + + ``_helix_sanitize_empty_kv`` and the transpose and split into the send + layout live in one compiled region so inductor folds the fill into the + transposed store instead of writing ``partial_o`` and reading it straight + back. Dynamo inlines the call, so the sanitize has exactly one definition + and the fusion is unaffected. + + ``dynamic=False`` is deliberate. With dynamic shapes inductor emits a much + slower transposed store, and picks different grids on different ranks for + the same shape, which cancels the win at the collective. + + The cost is one specialization per CUDA-graph batch bucket. Exceeding + dynamo's ``cache_size_limit`` falls back to eager SILENTLY -- the symptom is + ``triton_poi_fused_*`` disappearing from the trace, not an error, so + ``_helix_note_nccl_specialization`` counts the distinct shapes at the call + site and says so once. + """ + partial_o, softmax_stats = _helix_sanitize_empty_kv(partial_o, softmax_stats, zero_kv_mask) + chunks = [] + for t in (partial_o, softmax_stats): + t = t.transpose(1, 0).contiguous() + chunks.extend(torch.split(t, t.shape[0] // cp_size)) + return chunks + + +@torch.compile(dynamic=False) +def _helix_nccl_post_alltoall(gathered: List[torch.Tensor]) -> List[torch.Tensor]: + """Reformat the gathered partials into the helix_post_process layout.""" + return [t.transpose(1, 2).contiguous() for t in gathered] + + +def _helix_post_process( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + mapping: Mapping, + num_heads_tp_cp: int, + value_dim: int, + aux_stream: Optional[torch.cuda.Stream] = None, + ln_events: Optional[list] = None, + zero_kv_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Helix CP post-processing: all-to-all exchange and combine partial + attention outputs across CP ranks. + + *value_dim* is the per-head width of *partial_o*: ``head_dim`` for MHA; + for MLA ``kv_lora_rank``, or ``v_head_dim`` when the V up-projection runs + before the exchange. + + zero_kv_mask marks tokens this CP rank owns no KV for; those rows are forced + to a no-op contribution before the exchange, exactly once per backend: + NCCL in _helix_nccl_pre_alltoall, fused with the reformat + fifo v2 in the all-to-all sender, while the entry is in shared memory + fifo v1 here, via _helix_sanitize_empty_kv + + When *aux_stream* and *ln_events* are provided the two + ``.contiguous()`` calls in the FIFO-v1 path are overlapped on + separate CUDA streams for better performance. + """ + if mapping.cp_config.get("use_nccl_for_alltoall", True): + # NCCL path. Sanitize is folded into _helix_nccl_pre_alltoall so + # inductor can fuse it into the reformat. + _helix_note_nccl_specialization(partial_o, softmax_stats, mapping.cp_size) + chunks = _helix_nccl_pre_alltoall(partial_o, softmax_stats, zero_kv_mask, mapping.cp_size) + gathered = alltoall_helix(chunks, mapping.cp_group) + gathered = _helix_nccl_post_alltoall(gathered) + return torch.ops.trtllm.helix_post_process(gathered[0], gathered[1], 1.0) + else: + # FIFO-based implementation using MNNVL workspace. + helix = HelixAllToAllNative.get(mapping) + num_tokens = partial_o.shape[0] + cp_size = mapping.cp_size + fifo_version = mapping.cp_config.get("fifo_version", 2) + + if fifo_version == 1: + partial_o, softmax_stats = _helix_sanitize_empty_kv( + partial_o, softmax_stats, zero_kv_mask + ) + + def reshape_o(): + return ( + partial_o.view(num_tokens, cp_size, num_heads_tp_cp, value_dim) + .transpose(1, 2) + .contiguous() + ) + + def reshape_s(): + return ( + softmax_stats.view(num_tokens, cp_size, num_heads_tp_cp, 2) + .transpose(1, 2) + .contiguous() + ) + + if aux_stream is not None and ln_events is not None: + partial_o, softmax_stats = maybe_execute_in_parallel( + reshape_o, + reshape_s, + ln_events[0], + ln_events[1], + aux_stream, + ) + else: + partial_o = reshape_o() + softmax_stats = reshape_s() + + partial_o_out, softmax_stats_out = helix.alltoall_native(partial_o, softmax_stats) + return torch.ops.trtllm.helix_post_process_native( + partial_o_out, softmax_stats_out, 1.0, 2 + ) + else: + # fifo_v2: one entry is one token, and the sender already streams + # every byte through shared memory, so the sanitize rides along + # there instead of 6 separate elementwise kernels (~40% of the block). + partial_o = partial_o.view(num_tokens, cp_size, num_heads_tp_cp * value_dim) + softmax_stats = softmax_stats.view(num_tokens, cp_size, num_heads_tp_cp * 2) + partial_o_out, softmax_stats_out = helix.alltoall_native( + partial_o, + softmax_stats, + None if zero_kv_mask is None else zero_kv_mask[:num_tokens], + ) + gathered_o = partial_o_out.view(num_tokens, cp_size, num_heads_tp_cp, value_dim) + gathered_stats = softmax_stats_out.view(num_tokens, cp_size, num_heads_tp_cp, 2) + return torch.ops.trtllm.helix_post_process_native(gathered_o, gathered_stats, 1.0, 1) + + +def uses_k3_dspark(spec_config) -> bool: + """Whether ``spec_config`` is one-engine DSpark with a separately shipped drafter. + + Whether that drafter is MLA-shaped, and so served by the private classes, is only + known from its loaded config (:func:`install_k3_dspark`). + """ + return ( + spec_config is not None + and spec_config.spec_dec_mode.use_one_engine() + and spec_config.spec_dec_mode.is_dspark() + and not bool(getattr(spec_config, "draft_is_embedded_in_target", False)) + ) + + +def _replace_in_epilogue(target: nn.Module, old: Optional[nn.Module], new: nn.Module) -> None: + """Put ``new`` where main's shell appended ``old`` (by identity), else append it.""" + for i, module in enumerate(target.epilogue): + if module is old: + target.epilogue[i] = new + return + target.epilogue.append(new) + + +def install_k3_dspark(target: nn.Module, model_config: ModelConfig) -> bool: + """Replace main's DSpark drafter and worker on ``target`` with the private ones. + + Args: + target: a ``SpecDecOneEngineForCausalLM`` whose ``__init__`` just ran (it reads + ``draft_config``, ``draft_model``, ``spec_worker``, ``use_separate_draft_kv_cache``, + ``epilogue``, ``lm_head``, ``logits_processor`` and ``model``). + model_config: the target's ``ModelConfig``, with the mapping main's drafter was + built under still in place. + + Returns: + True when the private drafter and worker were installed; False when the + configuration keeps main's (no DSpark, embedded draft, GQA drafter). + """ + spec_config = getattr(model_config, "spec_config", None) + if not uses_k3_dspark(spec_config): + return False + # Main's shell built the drafter's own ModelConfig for the external drafter + # (external_drafter_config_kwargs: mapping_for_full_rows(), the draft KV quant + # algo, the target's extra_attrs); the private drafter is built on the same one, + # which the KV cache creator also reads. + draft_config = target.draft_config + if draft_config is None or not is_mla(draft_config.pretrained_config): + logger.info("K3 DSpark: not an MLA drafter; main's DSpark drafter and worker serve it") + return False + + # Imported here: only an MLA DSpark deployment pays for the private drafter and + # worker modules (and the kernel modules they import) at import time. + from .drafter import K3DSparkWorker, build_k3_dspark_draft + + draft_model = build_k3_dspark_draft(model_config, draft_config, target.lm_head, target.model) + if spec_config.spec_dec_mode.is_parallel_draft(): + draft_model.logits_processor = target.logits_processor + # The drafter runs on every row: the mapping main hands get_spec_worker. + spec_worker = K3DSparkWorker( + spec_config, + model_config.mapping.mapping_for_full_rows(), + target.use_separate_draft_kv_cache, + ) + + _replace_in_epilogue(target, target.draft_model, draft_model) + _replace_in_epilogue(target, target.spec_worker, spec_worker) + target.draft_model = draft_model + target.spec_worker = spec_worker + # As SpecDecOneEngineForCausalLM.__init__: cache the drafter's static vocab map + # and bind the worker to both models (vocab-parallel selection binds the heads + # in setup_aliases / post_load_weights, after loading). + spec_worker.set_draft_model(draft_model) + spec_worker.set_target_model(target) + logger.info( + f"K3 DSpark: private drafter {type(draft_model).__module__}.{type(draft_model).__name__}" + f" and worker {type(spec_worker).__module__}.{type(spec_worker).__name__}" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/features.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/features.py new file mode 100644 index 000000000000..03449b8a5e81 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/features.py @@ -0,0 +1,371 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Which optimizations the Kimi K3 SM107 target runs, per deployment. Nothing here reads the environment. + +``routing.py`` maps each deployment shape to one profile (its ``parallel`` value: ``tp1``, ``helix8``, +``helix16_moetp4ep4``, ``dep16``), and each profile turns on a fixed set of optimizations. Only optimizations that +some deployment able to reach them leaves off are listed here; the others run wherever their code is reached. +:data:`RETIRED_ENV` lists the environment variables the target does not read; :func:`select` warns about any that is +set. + +``route()`` calls :func:`select` with the deployment's profile before it imports the target's modules, so every +module of the target reads :func:`active` after the selection, at import (module-level constants) or when called. +Before any selection :func:`active` is ``K3Features()``: every optimization off, the plain code path. Unit tests that +import the modules directly get that; they change a field with ``monkeypatch.setattr(features, "_ACTIVE", +dataclasses.replace(features.active(), =))`` (fields read when called) or by patching the constant in +the module that reads it (fields read at import). +""" + +from __future__ import annotations + +import dataclasses +import os +import sys +from typing import Dict, Optional + + +@dataclasses.dataclass(frozen=True) +class K3Features: + """One deployment's settings. Each default is the plain code path.""" + + # Token embedding and the speculative decode step's head (modeling.py) + embedding_gather: bool = False # TLLM_SPEC_HEAD_FUSION + # Decode step inputs: one upload and one scatter at the graph replay (modeling.py) + pack_step_inputs: bool = False # TRTLLM_K3_ULL_STEP_PREP and TRTLLM_BATCHED_INPUT_COMMIT + # MoE (moe_runtime.py, fp8_weight_read.py, modeling.py) + fuse_moe_residual: bool = False # TRTLLM_K3_FUSE_MOE_RESIDUAL + fp8w_lowm_gemm: bool = False # TRTLLM_K3_FP8W_LOWM_GEMM + ld_fused: bool = False # TRTLLM_K3_LD_FUSED=1 + shared_sm: int = 0 # TRTLLM_K3_SHARED_SM + fc12_die_major: bool = False # TRTLLM_K3_FC12_DIE_LAYOUT and TRTLLM_K3_ULL_EXPERTS_DIE_MAJOR + router_exact_kda: bool = False # TRTLLM_K3_ROUTER_EXACT=kda + latent_ar_own_ws: bool = False # TRTLLM_K3_LATENT_AR_OWN_WS + # Attention residual and the decoder layer's all-reduces (attn_res.py, modeling.py, moe_runtime.py, kernels/) + attn_res_pdl: bool = False # TRTLLM_K3_ATTN_RES_PDL=2 + attn_res_pdl_max_m: int = 64 # TRTLLM_K3_ATTN_RES_PDL_MAX_M + attn_res_mxfp8: bool = False # TRTLLM_K3_ATTN_RES_MXFP8 + attn_res_early_trigger: bool = False # TRTLLM_K3_ATTN_RES_EARLY_TRIGGER + cute_carveout_fix: bool = False # TRTLLM_K3_CUTE_CARVEOUT_FIX + ull_oproj_sandwich: bool = False # TRTLLM_K3_ULL_OPROJ_SANDWICH + ull_tail_sandwich: bool = False # TRTLLM_K3_ULL_TAIL_SANDWICH + ull2_coll_cp_tail: bool = False # TRTLLM_K3_ULL2_COLL_CP_TAIL and TRTLLM_K3_ULL_TAIL1 + ull_hygiene: bool = False # TRTLLM_K3_ULL_HYGIENE + ull_front: bool = False # TRTLLM_K3_ULL_FRONT and TRTLLM_K3_ULL_FRONT_EARLY + ull4_mla_plan_early: bool = False # TRTLLM_K3_ULL4_MLA_PLAN_EARLY + # KDA (kda.py) + kda_prepad_prefill: bool = False # TLLM_KDA_PREPAD_PREFILL + kda_fused_conv_prefill: bool = False # TRTLLM_KDA_FUSED_CONV_PREFILL + kda_l2_prefetch: bool = False # TRTLLM_K3_KDA_L2_PREFETCH + # MLA (mla.py, mla.py, mla.py, mla.py) + ull2_mla_tail: bool = False # TRTLLM_K3_ULL2_MLA_TAIL + mla_q_norm_pdl: bool = False # TRTLLM_K3_MLA_Q_NORM_PDL + ull3_mla_gate_at: str = "decode" # TRTLLM_K3_ULL3_MLA_GATE_AT + ull2_mla_decode_pdl: bool = False # TRTLLM_K3_ULL2_MLA_DEC_PDL + # FP4 MLA generation decode (ops/attention/_fp4_mla_decode.py, kernels/fp4_mla_packed_reduce.py) + fp4_mla_page_plan: bool = True # TRTLLM_FP4_MLA_PAGE_PLAN (=0: False) + fp4_mla_reduce_pdl: bool = False # TRTLLM_FP4_MLA_REDUCE_PDL + fp4_mla_f16_partials: bool = False # TRTLLM_FP4_MLA_F16_PARTIALS + # DSpark (drafter.py, dflash.py) + draft3_glue: bool = False # TRTLLM_K3_DRAFT3_GLUE + dspark_draft_gemm: bool = False # TLLM_DSPARK_DRAFT_GEMM + ull2_draft_a2a: bool = False # TRTLLM_K3_ULL2_DRAFT_A2A and TRTLLM_K3_ULL2_DRAFT_TAIL + ull2_draft_store: bool = False # TRTLLM_K3_ULL2_DRAFT_STORE + spec_tail_fusion: bool = False # TLLM_SPEC_TAIL_FUSION + dspark_tail_v2: bool = False # TLLM_DSPARK_TAIL_V2 + + def changed(self) -> Dict[str, object]: + """The fields that differ from the plain path, by name.""" + base = K3Features() + return { + f.name: getattr(self, f.name) + for f in dataclasses.fields(self) + if getattr(self, f.name) != getattr(base, f.name) + } + + +# One profile per deployment (the names are routing.py's ``parallel`` values). +PROFILES: Dict[str, K3Features] = { + # Context servers (DWDP, one rank per process). + "tp1": K3Features( + kda_fused_conv_prefill=True, + kda_prepad_prefill=True, + ), + # Generation over Helix CP8. + "helix8": K3Features( + attn_res_early_trigger=True, + attn_res_mxfp8=True, + attn_res_pdl=True, + attn_res_pdl_max_m=128, + dspark_tail_v2=True, + embedding_gather=True, + fc12_die_major=True, + fp4_mla_reduce_pdl=True, + fp8w_lowm_gemm=True, + fuse_moe_residual=True, + kda_l2_prefetch=True, + latent_ar_own_ws=True, + ld_fused=True, + mla_q_norm_pdl=True, + router_exact_kda=True, + shared_sm=16, + spec_tail_fusion=True, + ), + # Generation over Helix CP16, routed experts MoE TP4 x EP4 (lowest latency). + "helix16_moetp4ep4": K3Features( + attn_res_early_trigger=True, + attn_res_mxfp8=True, + attn_res_pdl=True, + cute_carveout_fix=True, + draft3_glue=True, + dspark_draft_gemm=True, + dspark_tail_v2=True, + embedding_gather=True, + fc12_die_major=True, + fp4_mla_f16_partials=True, + fp4_mla_reduce_pdl=True, + fp8w_lowm_gemm=True, + fuse_moe_residual=True, + ld_fused=True, + mla_q_norm_pdl=True, + pack_step_inputs=True, + router_exact_kda=True, + spec_tail_fusion=True, + ull2_coll_cp_tail=True, + ull2_draft_a2a=True, + ull2_draft_store=True, + ull2_mla_decode_pdl=True, + ull2_mla_tail=True, + ull3_mla_gate_at="exchange:16,start", + ull4_mla_plan_early=True, + ull_front=True, + ull_hygiene=True, + ull_oproj_sandwich=True, + ull_tail_sandwich=True, + ), + # Generation under attention DP over 16 ranks. + "dep16": K3Features( + attn_res_pdl=True, + fp4_mla_page_plan=False, + fp8w_lowm_gemm=True, + fuse_moe_residual=True, + spec_tail_fusion=True, + ), +} + +_FIELD_ENV = { + "embedding_gather": "TLLM_SPEC_HEAD_FUSION", + "pack_step_inputs": "TRTLLM_K3_ULL_STEP_PREP", + "fuse_moe_residual": "TRTLLM_K3_FUSE_MOE_RESIDUAL", + "fp8w_lowm_gemm": "TRTLLM_K3_FP8W_LOWM_GEMM", + "ld_fused": "TRTLLM_K3_LD_FUSED", + "shared_sm": "TRTLLM_K3_SHARED_SM", + "fc12_die_major": "TRTLLM_K3_FC12_DIE_LAYOUT", + "router_exact_kda": "TRTLLM_K3_ROUTER_EXACT", + "latent_ar_own_ws": "TRTLLM_K3_LATENT_AR_OWN_WS", + "attn_res_pdl": "TRTLLM_K3_ATTN_RES_PDL", + "attn_res_pdl_max_m": "TRTLLM_K3_ATTN_RES_PDL_MAX_M", + "attn_res_mxfp8": "TRTLLM_K3_ATTN_RES_MXFP8", + "attn_res_early_trigger": "TRTLLM_K3_ATTN_RES_EARLY_TRIGGER", + "cute_carveout_fix": "TRTLLM_K3_CUTE_CARVEOUT_FIX", + "ull_oproj_sandwich": "TRTLLM_K3_ULL_OPROJ_SANDWICH", + "ull_tail_sandwich": "TRTLLM_K3_ULL_TAIL_SANDWICH", + "ull2_coll_cp_tail": "TRTLLM_K3_ULL2_COLL_CP_TAIL", + "ull_hygiene": "TRTLLM_K3_ULL_HYGIENE", + "ull_front": "TRTLLM_K3_ULL_FRONT", + "ull4_mla_plan_early": "TRTLLM_K3_ULL4_MLA_PLAN_EARLY", + "kda_prepad_prefill": "TLLM_KDA_PREPAD_PREFILL", + "kda_fused_conv_prefill": "TRTLLM_KDA_FUSED_CONV_PREFILL", + "kda_l2_prefetch": "TRTLLM_K3_KDA_L2_PREFETCH", + "ull2_mla_tail": "TRTLLM_K3_ULL2_MLA_TAIL", + "mla_q_norm_pdl": "TRTLLM_K3_MLA_Q_NORM_PDL", + "ull3_mla_gate_at": "TRTLLM_K3_ULL3_MLA_GATE_AT", + "ull2_mla_decode_pdl": "TRTLLM_K3_ULL2_MLA_DEC_PDL", + "fp4_mla_page_plan": "TRTLLM_FP4_MLA_PAGE_PLAN", + "fp4_mla_reduce_pdl": "TRTLLM_FP4_MLA_REDUCE_PDL", + "fp4_mla_f16_partials": "TRTLLM_FP4_MLA_F16_PARTIALS", + "draft3_glue": "TRTLLM_K3_DRAFT3_GLUE", + "dspark_draft_gemm": "TLLM_DSPARK_DRAFT_GEMM", + "ull2_draft_a2a": "TRTLLM_K3_ULL2_DRAFT_A2A", + "ull2_draft_store": "TRTLLM_K3_ULL2_DRAFT_STORE", + "spec_tail_fusion": "TLLM_SPEC_TAIL_FUSION", + "dspark_tail_v2": "TLLM_DSPARK_TAIL_V2", +} + +# Variables nothing reads: the code they switched runs wherever it is reached, or no deployment runs it. +_FIXED_ENV = ( + "KIMI_K3_ATTN_RES_TOPOLOGY", + "KIMI_K3_FUSED_ATTN_RES", + "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS", + "TLLM_DSPARK_AR_PUSH", + "TLLM_DSPARK_CTX_SIDE_STREAM", + "TLLM_DSPARK_DRAFT_FUSION", + "TLLM_DSPARK_FUSED_MARKOV", + "TLLM_DSPARK_MARKOV_V2", + "TLLM_DSPARK_MARKOV_V3", + "TLLM_K3_MLA_GEN_BACKEND", + "TLLM_KDA_BUF_CACHE_ENTRIES", + "TLLM_SPEC_VOCAB_PARALLEL", + "TLLM_SPEC_VOCAB_PARALLEL_COMM", + "TLLM_SPEC_VOCAB_PARALLEL_SAMPLING", + "TRTLLM_BATCHED_INPUT_COMMIT", + "TRTLLM_BATCHED_INPUT_COMMIT_VERIFY", + "TRTLLM_FP4_MLA_LPT_PLAN", + "TRTLLM_FP4_MLA_PACKED_WAVE_PLAN", + "TRTLLM_FP4_MLA_REDUCE_SPREAD_STATS", + "TRTLLM_FP4_MLA_SPLIT_KV", + "TRTLLM_HELIX_V_UP_BEFORE_A2A", + "TRTLLM_K3_AR_FINISH_FUSED", + "TRTLLM_K3_AR_PUSH", + "TRTLLM_K3_AR_PUSH_AMID", + "TRTLLM_K3_AR_PUSH_ONE_SHOT_MAX_BYTES", + "TRTLLM_K3_ATTN_RES_SM_RESERVE", + "TRTLLM_K3_BF16_OVER_PREQUANTIZED_MAX_M", + "TRTLLM_K3_DRAFT3_QHEAD", + "TRTLLM_K3_DRAFT3_REDUCE_UNROLL", + "TRTLLM_K3_FC12_DIE_SCHED", + "TRTLLM_K3_FC12_DIE_SCHED_MIN_TOKENS", + "TRTLLM_K3_FC12_PREZERO", + "TRTLLM_K3_FC12_SCHED_LOOKAHEAD", + "TRTLLM_K3_FC12_SM_RESERVE", + "TRTLLM_K3_FP4_MLA_SHARED_DECODE", + "TRTLLM_K3_FP8W_LOWM_FUSED_A", + "TRTLLM_K3_FUSED_FINISH_ONE_SHOT_MAX_BYTES", + "TRTLLM_K3_HELIX_TAIL_CTAS", + "TRTLLM_K3_HELIX_TAIL_FUSED", + "TRTLLM_K3_HELIX_TAIL_GATE", + "TRTLLM_K3_HELIX_TAIL_MB", + "TRTLLM_K3_KDA_DECODE_EARLY_TRIGGER", + "TRTLLM_K3_KDA_DECODE_FUSED_FAB", + "TRTLLM_K3_KDA_L2_PREFETCH_MAX_ROWS", + "TRTLLM_K3_KDA_MTP_V2", + "TRTLLM_K3_KDA_MX_PDL", + "TRTLLM_K3_KDA_QKVG_STAGE", + "TRTLLM_K3_KDA_QKVG_STAGE_RING", + "TRTLLM_K3_LATENT_AR_CUTE", + "TRTLLM_K3_LATENT_AR_ONESHOT_MAX_BYTES", + "TRTLLM_K3_LATENT_AR_OWN_WS_MAX_M", + "TRTLLM_K3_LATENT_AR_ROW", + "TRTLLM_K3_LATENT_AR_ROW_MAX_M", + "TRTLLM_K3_LATENT_STREAM_PRIORITY", + "TRTLLM_K3_LD_FUSED_MAX_M", + "TRTLLM_K3_LD_FUSED_MIN_M", + "TRTLLM_K3_LOWM_SF_SLOTS", + "TRTLLM_K3_LU_MID_MAX_M", + "TRTLLM_K3_LU_MID_MXFP8", + "TRTLLM_K3_MID_AR_PUSH", + "TRTLLM_K3_MID_AR_PUSH_MAX_M", + "TRTLLM_K3_MID_MXFP8_IN", + "TRTLLM_K3_MID_OPROJ_PC", + "TRTLLM_K3_MLA_CTX_FP8_KV_FUSED", + "TRTLLM_K3_MLA_CTX_TRITON_QUANT", + "TRTLLM_K3_MLA_GATE_SIDE", + "TRTLLM_K3_MLA_GATE_SIDE_MAX_M", + "TRTLLM_K3_MLA_Q_FUSED", + "TRTLLM_K3_MLA_Q_FUSED_KV_PRIO", + "TRTLLM_K3_MLA_Q_FUSED_MAX_TOKENS", + "TRTLLM_K3_MLA_Q_FUSED_NORM", + "TRTLLM_K3_OPROJ_PC", + "TRTLLM_K3_ROUTER_FIRST", + "TRTLLM_K3_ROUTER_FIRST_MAX_M", + "TRTLLM_K3_ROUTE_APPROX", + "TRTLLM_K3_ROUTE_GEMV_MAX_M", + "TRTLLM_K3_ROUTE_TAIL", + "TRTLLM_K3_ROUTE_TAIL_MAX_M", + "TRTLLM_K3_ROUTE_TAIL_ORDER", + "TRTLLM_K3_ROUTE_TAIL_ORDER_MIN_ROWS", + "TRTLLM_K3_ROUTING_STATS_CHUNK", + "TRTLLM_K3_ROUTING_STATS_DIR", + "TRTLLM_K3_ROUTING_STATS_EVERY", + "TRTLLM_K3_ROUTING_STATS_STRIDE", + "TRTLLM_K3_RP_TAIL", + "TRTLLM_K3_RP_TAIL_MAX_M", + "TRTLLM_K3_SHARED_STREAM_PRIORITY", + "TRTLLM_K3_SHARED_TAIL_MAX_M", + "TRTLLM_K3_SHARED_TAIL_MIN_M", + "TRTLLM_K3_SHARED_TAIL_PRIO", + "TRTLLM_K3_ULL2_COLL_AR_EARLY_STATE", + "TRTLLM_K3_ULL2_DRAFT_TAIL", + "TRTLLM_K3_ULL3_MLA_LATE_TAIL_MAX_TOKENS", + "TRTLLM_K3_ULL3_MLA_TAIL_TRIGGER", + "TRTLLM_K3_ULL_EXPERTS", + "TRTLLM_K3_ULL_EXPERTS_COMBINE", + "TRTLLM_K3_ULL_EXPERTS_CTAS", + "TRTLLM_K3_ULL_EXPERTS_DIE_MAJOR", + "TRTLLM_K3_ULL_EXPERTS_TRIGGER", + "TRTLLM_K3_ULL_FRONT_EARLY", + "TRTLLM_K3_ULL_FRONT_MAX_M", + "TRTLLM_K3_ULL_MOEBLK", + "TRTLLM_K3_ULL_MOETP", + "TRTLLM_K3_ULL_SHARED_PERSIST", + "TRTLLM_K3_ULL_TAIL1", +) + +# Variables read under another name -> that name. +_RENAMED_ENV = { + "TRTLLM_MNNVL_AR_ONESHOT_MAX_BYTES": "TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES", +} + +#: Environment variables the target does not read -> the field that replaces them (None: no field; +#: ``_RENAMED_ENV`` names the variable read instead). +RETIRED_ENV: Dict[str, Optional[str]] = { + **{env: field for field, env in _FIELD_ENV.items()}, + **{env: None for env in _FIXED_ENV}, + **{env: None for env in _RENAMED_ENV}, +} + +_ACTIVE = K3Features() +_SELECTED: Optional[str] = None +_PACKAGE = __name__.rpartition(".")[0] + + +def active() -> K3Features: + """The selected target's settings (``K3Features()`` before :func:`select`).""" + return _ACTIVE + + +def selected() -> Optional[str]: + """The selected target, or None.""" + return _SELECTED + + +def select(target: str) -> K3Features: + """Make ``target``'s profile the active one; once per process, before any other module of the target. + + Raises if another profile was selected, or if a module of the target was imported first (its + import-time settings would be the plain path's, not ``target``'s). + """ + global _ACTIVE, _SELECTED + if target not in PROFILES: + raise KeyError(f"no Kimi K3 feature profile {target!r}; have {sorted(PROFILES)}") + if _SELECTED is not None: + if _SELECTED != target: + raise RuntimeError( + f"Kimi K3 features: {_SELECTED!r} is selected in this process, cannot select {target!r}" + ) + return _ACTIVE + early = sorted(m for m in sys.modules if m.startswith(_PACKAGE + ".") and m != __name__) + if early: + raise RuntimeError( + f"Kimi K3 features: modules of the target were imported before the {target!r} profile was " + f"selected, with the plain path's settings: {', '.join(early)}" + ) + _ACTIVE, _SELECTED = PROFILES[target], target + + from tensorrt_llm.logger import logger + + changed = _ACTIVE.changed() + logger.info( + f"Kimi K3 features: {target} ({len(changed)} set): " + + ", ".join(f"{name}={value}" for name, value in sorted(changed.items())) + ) + stale = sorted(name for name in RETIRED_ENV if name in os.environ) + if stale: + names = ", ".join( + f"{name} (set {_RENAMED_ENV[name]} instead)" if name in _RENAMED_ENV else name + for name in stale + ) + logger.warning( + "Kimi K3 features: the target does not read these environment variables (its " + f"profile in features.py decides): {names}" + ) + return _ACTIVE diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/fp8_weight_read.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/fp8_weight_read.py new file mode 100644 index 000000000000..5c033e7777f9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/fp8_weight_read.py @@ -0,0 +1,929 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 on Rubin (SM107): projections that read FP8 block-scaled weights. + +``_Fp8BlockScaleWeightReadLinear`` keeps a bias-free projection as E4M3 codes with 128x128 +block scales: the checkpoint FP8 attention projections, and the MoE layers' BF16 MLPs when +they are quantized after loading. Low-M rows run the low-M GEMMs; the rest runs DeepGEMM on +SM100/SM103 and the native MXFP8 GEMM on SM107. ``_swap_linear_to_fp8_weight_read`` and +``_convert_moe_mlps_to_fp8_weight_read`` put it in place of a loaded linear. +""" + +from __future__ import annotations + +import gc +from typing import Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch.distributed import AllReduceParams +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from .features import active as _features +from .layers import ( + Bf16Linear, + Fp8BlockScaleLinear, + fp8_block_scale_gemm, + low_m_gemm_rows, + low_m_gemm_scale, + prepare_fp8_block_weight, +) +from .ops.gemm.k3_fp8w_lowm_gemm import k3_fp8w_lowm_gemm as k3_fp8w_lowm_gemm_op +from .ops.gemm.k3_fp8w_lowm_gemm_push import k3_fp8w_lowm_gemm_push as k3_fp8w_lowm_gemm_push_op +from .ops.gemm.k3_fp8w_lowm_mxfp8_gemm import k3_fp8w_lowm_mxfp8_gemm as k3_fp8w_lowm_mxfp8_gemm_op +from .ops.gemm.k3_fp8w_lowm_mxfp8_push import k3_fp8w_lowm_mxfp8_push as k3_fp8w_lowm_mxfp8_push_op +from .ops.gemm.k3_fp8w_mid_gemm import k3_fp8w_mid_gemm as k3_fp8w_mid_gemm_op +from .ops.gemm.k3_fp8w_mid_gemm_push import k3_fp8w_mid_gemm_push as k3_fp8w_mid_gemm_push_op +from .ops.gemm.k3_mxfp8_gemm import k3_mxfp8_gemm as k3_mxfp8_gemm_op + +_F = _features() + +# FP8 weight-read projections with few rows (per-shape, per-architecture limit, see +# k3_fp8w_lowm_gemm.max_m_for) run the CuTe low-M GEMM: on SM100/SM103 instead of DeepGEMM +# swap-AB, on SM107 instead of the CuTe MXFP8 GEMM (with a DeepGEMM-packed copy of the weight +# scale, built at post-load, see _Fp8BlockScaleWeightReadLinear.prepare_lowm_gemm). +_K3_FP8W_LOWM_GEMM = _F.fp8w_lowm_gemm + +# Fused latent projection (cute_dsl_kernels/blackwell/k3_latent_fp8_gemm.py), for token counts +# in [_K3_LD_FUSED_MIN_M, _K3_LD_FUSED_MAX_M]: +# TRTLLM_K3_LD_FUSED: latent_down + the routed experts' NVFP4 input quantization in one chain +# (the k3_latent_down_fp4 entry), bit-exact with fp8_swap_ab_gemm followed by fp4_quantize. +# It fuses only where the unfused projection would run fp8_swap_ab_gemm, not where it runs the +# low-M kernel (TRTLLM_K3_FP8W_LOWM_GEMM). On SM107 the op runs the low-M kernel's split-K GEMM +# with the fused epilogue up to k3_latent_fp8_gemm.lowm_route_max_m tokens (bit-exact with that +# GEMM at the same tactic followed by fp4_quantize) and the persistent kernel above; both read a +# DeepGEMM-packed copy of the weight scale built at post-load +# (_Fp8BlockScaleWeightReadLinear.prepare_lowm_gemm). +_K3_LD_FUSED = _F.ld_fused + + +# --------------------------------------------------------------------------- +# FP8 block-scale weight read for attention and replicated MoE-layer MLP projections. +# --------------------------------------------------------------------------- + + +def _rubin_mxfp8_is_available() -> bool: + """Whether this rank can run the SM107 native block-scaled MXFP8 GEMM. + + Recomputed rather than cached: the SM version cannot change within a run, + so weight preparation and forward always agree. + """ + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + return get_sm_version() == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE + + +def _is_fp8_block_projection(module: nn.Module) -> bool: + """A checkpoint FP8 128 x 128 block-scale projection: the loader reads its codes and scales + together and post-load swaps it for ``_Fp8BlockScaleWeightReadLinear``.""" + return isinstance(module, Fp8BlockScaleLinear) + + +def _k3_lowm_gemm_can_run( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + allow_partial_n_tile: bool = False, +) -> bool: + """Whether the K3 low-M GEMM (``ops.gemm.k3_fp8w_lowm_gemm``) serves this projection.""" + from .kernels import k3_fp8w_lowm_gemm + + return k3_fp8w_lowm_gemm.can_run( + x, weight, weight_scale, allow_partial_n_tile=allow_partial_n_tile + ) + + +class _Fp8BlockScaleWeightReadLinear(nn.Module): + """Bias-free FP8 block-scale GEMM for checkpoint attention and converted MLPs. + + Attention loads checkpoint E4M3 codes and 128x128 scales directly. The optional + MLP path quantizes BF16 weights once at load. On SM100/SM103 the UE8M0 + scales are packed for DeepGEMM, which quantizes BF16 activations inside the + GEMM. On Rubin (SM107) the scales are packed into the CuTe DSL K32 R128c4 + layout and the activation is quantized by the specialized CuTe quantizer + before the native MXFP8 GEMM, which runs with K3's fine-grained low-M tuning + buckets. On both, the low-M GEMM (``trtllm::fp8_block_scale_low_m_gemm``) serves + the rows ``layers.low_m_gemm_rows`` gives for the weight shape; it reads a + DeepGEMM-packed scale, which is ``weight_scale`` on SM100/SM103 and a copy + repacked from the MXFP8 scale on SM107. That dispatch is + ``layers.fp8_block_scale_gemm``. + + The Rubin agentic kernels sit in front of that dispatch: the K3 low-M GEMM + (``TRTLLM_K3_FP8W_LOWM_GEMM``, ``kernels/k3_fp8w_lowm_gemm.py``) with its gated, + sigmoid, MXFP8-input and all-reduce-push variants. With ``TRTLLM_K3_FP8W_LOWM_GEMM`` + set the K3 low-M GEMM owns the low-M rows exactly as on that branch and the generic low-M GEMM is not + prepared; with it unset the generic one serves them as on main. + """ + + def __init__( + self, weight_fp8: torch.Tensor, weight_scale: torch.Tensor, out_features: int + ) -> None: + super().__init__() + self.in_features = weight_fp8.shape[1] + self.out_features = out_features + self._weights_transformed = True + self._fused_projection = None + # Buffers (not parameters): these are the module's weights post-load; + # there is nothing further to load into them and they must not be + # touched by any later autocast/dtype move. + self.register_buffer("weight", weight_fp8, persistent=False) + self.register_buffer("weight_scale", weight_scale, persistent=False) + # Scales of the output-row slices :meth:`forward_slice` serves (an + # ``Fp8BlockScaleWeightSliceCache``), prepared by :meth:`prepare_output_slice`. + self._output_slice_cache: Optional[nn.Module] = None + # Rows the low-M GEMM serves (0: none) and the DeepGEMM-packed scale it reads; set at + # post-load. + self._low_m_max = 0 + self._low_m_weight_scale: Optional[torch.Tensor] = None + # Set by a caller for the next forward only: a contiguous BF16 ``[M, out_features]`` + # tensor that the low-M GEMM writes its result into. Other paths return their own + # tensor, so the caller compares the result with the buffer it handed in. + self.next_output: Optional[torch.Tensor] = None + # Set by a caller for the next forward only (every forward consumes and clears it): the + # ``LamportPushTarget`` (``layers.lamport_push_target``) of the all-reduce that + # follows this projection. Where the low-M GEMM without split-K (or the mid-M GEMM with + # its push epilogue) serves ``forward`` / ``forward_gated``, it stores its output into + # that all-reduce's Lamport buffers instead, marks the target pushed and returns the + # all-reduce's output tensor, which ``layers.finish_lamport_push`` fills. + self.next_ar_push = None + # Copied from the source ``Linear`` (the MLA fused q_a/kv_a projection sets it): the + # low-M GEMM may serve this weight although N is not a multiple of 128. + self.fp8w_lowm_partial_n_tile = False + # Whether the fused latent projection (TRTLLM_K3_LD_FUSED) may replace this one (set when + # it is swapped in). + self.k3_latent_fused = False + # SM107 with TRTLLM_K3_FP8W_LOWM_GEMM or ``k3_latent_fused``: the weight scale in the + # DeepGEMM-packed UE8M0 layout of the low-M GEMM and the fused latent projections + # (``weight_scale`` holds the CuTe MXFP8 layout there). + self.lowm_weight_scale: Optional[torch.Tensor] = None + + @property + def has_fp8_qdq(self) -> bool: + """Match the ``Linear`` interface consumed by ``GatedMLP``.""" + return False + + @property + def has_w4a8_nvfp4_fp8(self) -> bool: + """Match the ``Linear`` interface consumed by ``GatedMLP``.""" + return False + + @property + def supports_prequantized_input(self) -> bool: + """Whether a producer may hand this GEMM an FP8 activation in CuTe layout.""" + return _rubin_mxfp8_is_available() + + @property + def packed_weight_scale(self) -> Optional[torch.Tensor]: + """The weight scale in DeepGEMM's packed UE8M0 layout, as the low-M GEMM and the fused + latent projections read it: ``weight_scale`` on SM100/SM103, ``lowm_weight_scale`` on + SM107 (``None`` unless :meth:`prepare_lowm_gemm` built it).""" + if _rubin_mxfp8_is_available(): + return self.lowm_weight_scale + return self.weight_scale + + def takes_prequantized_input(self, num_rows: int) -> bool: + """Whether a producer that can emit this GEMM's scale layout should hand ``num_rows`` + rows to :meth:`forward_prequantized` rather than BF16 rows to :meth:`forward`. + + False where :meth:`forward` serves BF16 rows with the low-M GEMM (activation quantized + inside the kernel) faster than the CuTe MXFP8 GEMM runs on the prequantized operand. + """ + if not self.supports_prequantized_input: + return False + if self.lowm_weight_scale is None: + return True + from .kernels import k3_fp8w_lowm_gemm + + return num_rows > k3_fp8w_lowm_gemm.bf16_over_prequantized_max_m( + self.out_features, self.in_features + ) + + @staticmethod + def quantize_weight(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """BF16 ``[out, in]`` weight -> (FP8 weight, backend-ready scale). + + Both dims must be multiples of 128. Because the 128x128 block scale is + computed per block, concatenating several such weights along ``out`` + and quantizing the result is per-block identical to quantizing each + separately (no block crosses a 128-aligned boundary), so a fused + weight's row slices equal the individually quantized weights. + """ + # Lazy imports: only pulled in on the FP8 path. + from tensorrt_llm.deep_gemm.utils.math import per_block_cast_to_fp8 + + weight_fp8, weight_scale = per_block_cast_to_fp8(weight, use_ue8m0=False) + # Both GEMMs consume UE8M0 scales rather than the FP32 grid. + return prepare_fp8_block_weight(weight_fp8, weight_scale) + + @staticmethod + def _prepare_scale_layout(weight_scale: torch.Tensor, n: int, k: int) -> torch.Tensor: + from tensorrt_llm.quantization.utils.fp8_utils import ( + transform_k128_scales_to_cutedsl_mxfp8_layout, + transform_sf_into_required_layout, + ) + + if _rubin_mxfp8_is_available(): + # cute_dsl_mxfp8_gemm_rubin reads K32 R128c4 UE8M0 scales. + return transform_k128_scales_to_cutedsl_mxfp8_layout(weight_scale, mn=n, k=k) + # fp8_swap_ab_gemm with disable_ue8m0_cast=True consumes packed, + # TMA-aligned UE8M0 scales. + return transform_sf_into_required_layout( + weight_scale, + mn=n, + k=k, + recipe=(1, 128, 128), + is_sfa=False, + ) + + @classmethod + def from_linear(cls, linear: nn.Module) -> "_Fp8BlockScaleWeightReadLinear": + """The FP8 weight read of a loaded projection, sharing its codes and scales: + + * ``core.Fp8BlockScaleLinear`` (checkpoint E4M3 codes and FP32 128 x 128 scales, or their + ``prepare``d form); + * a bias-free BF16 projection (``nn.Linear`` or ``core.Bf16Linear``): quantized here. + """ + assert linear.bias is None, "FP8 weight read expects a bias-free Linear" + if isinstance(linear, Fp8BlockScaleLinear): + result = cls(linear.weight.detach(), linear.weight_scale.detach(), linear.out_features) + result.quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) + result.tp_size, result.tp_rank = linear.tp_size, linear.tp_rank + result._weights_transformed = getattr(linear, "prepared", False) + result.fp8w_lowm_partial_n_tile = getattr(linear, "fp8w_lowm_partial_n_tile", False) + return result + weight_fp8, weight_scale = cls.quantize_weight(linear.weight.data) + return cls(weight_fp8, weight_scale, linear.out_features) + + def transform_weights(self) -> None: + if not self._weights_transformed: + self._output_slice_cache = None + if self._fused_projection is not None: + owner_ref, start = self._fused_projection + owner = owner_ref() + if owner is None: + raise RuntimeError("KDA fused FP8 projection was released before its views") + if self.weight_scale.dtype != torch.float32 or self.weight_scale.ndim != 2: + raise ValueError("KDA fused FP8 views require a raw FP32 scale grid") + # Codes and raw scale grids alias the fused projection. Only its + # owner may resmooth them in place; each view packs the updated + # scales without modifying the shared codes again. + owner.transform_weights() + self.weight = owner.weight[start : start + self.out_features] + self.weight_scale = self._prepare_scale_layout( + self.weight_scale, self.out_features, self.in_features + ) + else: + self.weight, self.weight_scale = prepare_fp8_block_weight( + self.weight, self.weight_scale + ) + self._weights_transformed = True + + def post_load_weights(self) -> None: + self.transform_weights() + self.prepare_lowm_gemm() + self._prepare_low_m_gemm() + + def _prepare_low_m_gemm(self) -> None: + """Select the low-M GEMM for this weight where it has a DeepGEMM-packed scale. + + On SM107 that scale is a copy (N * K / 128 bytes) derived from the MXFP8 + ``weight_scale``; row views of a fused KDA projection have none (the fused module + serves decode). Not prepared under ``TRTLLM_K3_FP8W_LOWM_GEMM``: the K3 low-M GEMM + owns those rows then (see the class docstring). + """ + self._low_m_max, self._low_m_weight_scale = 0, None + if _K3_FP8W_LOWM_GEMM: + return + n, k = self.out_features, self.in_features + low_m_max = low_m_gemm_rows(n, k) + if not _rubin_mxfp8_is_available(): + self._low_m_max, self._low_m_weight_scale = low_m_max, self.weight_scale + elif low_m_max and self._fused_projection is None: + self._low_m_max = low_m_max + self._low_m_weight_scale = low_m_gemm_scale(self.weight_scale, n, k) + + def prepare_lowm_gemm(self) -> bool: + """Build ``lowm_weight_scale`` on SM107 if TRTLLM_K3_FP8W_LOWM_GEMM is set or a fused + latent projection may replace this one (``k3_latent_fused``). + + Derived from the transformed MXFP8 scale, so it follows weight reloads; rebuilt on + every call (post-load), before CUDA graph capture. Costs N * K / 128 bytes (0.8% of + the FP8 weight). SM100/SM103 run these GEMMs on ``weight_scale`` itself. + """ + self.lowm_weight_scale = None + if not ((_K3_FP8W_LOWM_GEMM or self.k3_latent_fused) and _rubin_mxfp8_is_available()): + return False + if self._fused_projection is not None: + # A row view of the KDA fused q/k/v/g weight; the fused module serves forward. + return False + self.transform_weights() + from .kernels import k3_fp8w_lowm_gemm + + n, k = self.weight.shape + if n % k3_fp8w_lowm_gemm.N_ALIGN or k % 128: + return False + self.lowm_weight_scale = k3_fp8w_lowm_gemm.pack_block_scale_from_mxfp8_layout( + self.weight_scale, n, k + ) + return True + + def _lowm_gemm_applies(self, x_flat: torch.Tensor, weight_scale: torch.Tensor) -> bool: + """Whether the ``k3_fp8w_lowm_gemm`` entry serves ``x_flat`` with ``weight_scale``.""" + return x_flat.dtype == torch.bfloat16 and _k3_lowm_gemm_can_run( + x_flat, + self.weight, + weight_scale, + allow_partial_n_tile=self.fp8w_lowm_partial_n_tile, + ) + + def _run_lowm_gemm( + self, + x_flat: torch.Tensor, + weight_scale: torch.Tensor, + out_buffer: Optional[torch.Tensor], + ) -> torch.Tensor: + """The low-M GEMM, into ``out_buffer`` when the caller handed in a usable one.""" + if ( + out_buffer is not None + and out_buffer.dtype == torch.bfloat16 + and out_buffer.shape == (x_flat.shape[0], self.out_features) + and out_buffer.is_contiguous() + ): + return k3_fp8w_lowm_gemm_op(x_flat, self.weight, weight_scale, out_buffer) + return k3_fp8w_lowm_gemm_op(x_flat, self.weight, weight_scale) + + def _take_next_output(self, num_rows: int) -> Optional[torch.Tensor]: + """Clear :attr:`next_output`; return it if it is a contiguous, 16-byte aligned BF16 + ``[num_rows, out_features]`` tensor.""" + out, self.next_output = self.next_output, None + if ( + out is None + or out.dtype != torch.bfloat16 + or tuple(out.shape) != (num_rows, self.out_features) + or not out.is_contiguous() + or out.data_ptr() % 16 + ): + return None + return out + + def _take_next_ar_push(self, x_flat: torch.Tensor, gated: bool = False): + """Clear :attr:`next_ar_push`; return it if the low-M GEMM can push this forward's + ``[M, out_features]`` output into it (low-M kernel without split-K; above the low-M + range, on ungated rows, the mid-M GEMM's push epilogue).""" + push, self.next_ar_push = self.next_ar_push, None + if ( + push is None + or x_flat.dtype != torch.bfloat16 + or tuple(push.output.shape) != (x_flat.shape[0], self.out_features) + ): + return None + from .kernels import k3_fp8w_lowm_gemm + + tactic = k3_fp8w_lowm_gemm.ar_push_tactic( + x_flat.shape[0], self.out_features, self.in_features, gated=gated + ) + if tactic is None and not gated: + tactic = k3_fp8w_lowm_gemm.mid_ar_push_tactic( + x_flat.shape[0], self.out_features, self.in_features + ) + return push if tactic is not None else None + + def mid_push_pdl_ready(self, num_rows: int) -> bool: + """``TRTLLM_K3_KDA_MX_PDL``: whether :meth:`forward_prequantized` on the producer's + ``[num_rows, in_features]`` R128c4 operand will push with the mid-M GEMM (the conditions + of :meth:`_mid_prequantized_push_tactic` that do not depend on the tensors), which reads + nothing the producer writes before its ``griddepcontrol.wait``: the producer may then + release it early.""" + push = self.next_ar_push + if push is None or self.lowm_weight_scale is None: + return False + if num_rows > 128: + return False + if tuple(push.output.shape) != (num_rows, self.out_features): + return False + from .kernels import k3_fp8w_lowm_gemm + + tactic = k3_fp8w_lowm_gemm.mid_ar_push_tactic(num_rows, self.out_features, self.in_features) + return tactic is not None and tactic.mma_n <= 128 + + def _mid_prequantized_push_tactic(self, activation: torch.Tensor, scale: torch.Tensor, push): + """The mid-M tactic :meth:`forward_prequantized` pushes the KDA verify's MXFP8 operand with, + or None. Up to 128 rows (one scale-factor block), where the R128c4 scales of the operand are + the mid-M GEMM's ``[K / 128][1][512 B]`` layout; the mid-M GEMM instead of the CuTe MXFP8 + GEMM, so the projection's FP32 accumulation order can differ there.""" + if push is None or self.lowm_weight_scale is None: + return None + if activation.dim() != 2 or not activation.is_contiguous(): + return None + m, k = activation.shape + if k != self.in_features or tuple(push.output.shape) != (m, self.out_features): + return None + from .kernels import k3_fp8w_lowm_gemm + + tactic = k3_fp8w_lowm_gemm.mid_ar_push_tactic(m, self.out_features, k) + if tactic is None or tactic.mma_n > 128: + return None + if scale.numel() != (k // 128) * 512 or not scale.is_contiguous(): + return None + return tactic + + def _run_ar_push(self, x_flat: torch.Tensor, push, gate=None) -> torch.Tensor: + """The low-M GEMM with the all-reduce push epilogue; returns the all-reduce output.""" + k3_fp8w_lowm_gemm_push_op( + x_flat, self.weight, self.packed_weight_scale, push.workspace, push.one_shot, gate=gate + ) + push.pushed = True + return push.output + + def forward( + self, + x: torch.Tensor, + *, + all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, + layer_idx: Optional[int] = None, + ) -> torch.Tensor: + out_shape = (*x.shape[:-1], self.out_features) + x_flat = x.reshape(-1, x.shape[-1]) + out_buffer, self.next_output = self.next_output, None + ar_push = self._take_next_ar_push(x_flat) + + # The K3 low-M GEMM (``TRTLLM_K3_FP8W_LOWM_GEMM``) up to its per-shape M gate, on SM107 + # with ``lowm_weight_scale``, on SM100 / SM103 with ``weight_scale``; else the low-M / + # MXFP8 (SM107) / swap-AB (SM100 / SM103) GEMM of ``layers.fp8_block_scale_gemm``. + lowm_scale = self.lowm_weight_scale if _rubin_mxfp8_is_available() else self.weight_scale + if ( + _K3_FP8W_LOWM_GEMM + and lowm_scale is not None + and self._lowm_gemm_applies(x_flat, lowm_scale) + ): + if ar_push is not None: + out = self._run_ar_push(x_flat, ar_push) + else: + out = self._run_lowm_gemm(x_flat, lowm_scale, out_buffer) + else: + out = fp8_block_scale_gemm( + x_flat, + self.weight, + self.weight_scale, + low_m_max=self._low_m_max, + low_m_scale=self._low_m_weight_scale, + ) + return out.reshape(out_shape) + + def prepare_output_slice(self, start: int, end: int) -> None: + """Prepare the scales of output features ``[start, end)`` for :meth:`forward_slice`. + + The shared MLA calls this at post-load for the context heads of a Helix rank. The slice + keeps the 128-row scale grid: it covers the whole scale tiles around ``[start, end)``. + """ + from tensorrt_llm.quantization.utils.fp8_utils import Fp8BlockScaleWeightSliceCache + + if not 0 <= start < end <= self.out_features: + raise ValueError(f"Invalid output feature range [{start}, {end})") + self.transform_weights() + if self._output_slice_cache is None: + self._output_slice_cache = Fp8BlockScaleWeightSliceCache() + self._output_slice_cache.prepare(self.weight, self.weight_scale, start, end) + + def forward_slice(self, x: torch.Tensor, start: int, end: int) -> torch.Tensor: + """Output features ``[start, end)`` of :meth:`forward`, contiguous. + + One GEMM over the weight rows of the slice's scale tiles (the general path of + :meth:`forward`: the CuTe MXFP8 GEMM on SM107, ``fp8_swap_ab_gemm`` elsewhere), cropped + to the slice. + """ + self.prepare_output_slice(start, end) + weight, scale, crop = self._output_slice_cache.get(self.weight, start, end) + x_flat = x.reshape(-1, x.shape[-1]) + if _rubin_mxfp8_is_available(): + act_fp8, act_scale = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(x_flat) + out = k3_mxfp8_gemm_op(act_fp8, weight, act_scale, scale) + else: + out = torch.ops.trtllm.fp8_swap_ab_gemm( + x_flat, weight, scale, output_dtype=x.dtype, disable_ue8m0_cast=True + ) + out = out.reshape(*x.shape[:-1], weight.shape[0]) + return out[..., crop : crop + end - start].contiguous() + + def _lowm_forward_scale(self, x_flat: torch.Tensor) -> Optional[torch.Tensor]: + """The weight scale :meth:`forward` runs the K3 low-M GEMM with on the BF16 rows + ``x_flat`` (``lowm_weight_scale`` on SM107, ``weight_scale`` on SM100/SM103), or None + where :meth:`forward` runs another GEMM.""" + if not _K3_FP8W_LOWM_GEMM or x_flat.dtype != torch.bfloat16: + return None + if self.lowm_weight_scale is not None: + scale = self.lowm_weight_scale + elif _rubin_mxfp8_is_available(): + return None + else: + scale = self.weight_scale + return scale if self._lowm_gemm_applies(x_flat, scale) else None + + def _runs_low_m_kernel(self, num_rows: int) -> bool: + """Whether :meth:`forward` on ``num_rows`` BF16 rows would run the low-M kernel (the + input-independent part of its dispatch: the low-M GEMM on, not the Rubin CuTe MXFP8 + path, ``num_rows`` within the low-M kernel's range). On SM107 forward runs it on + ``lowm_weight_scale`` up to the per-shape gate (``max_m_for``).""" + if not _K3_FP8W_LOWM_GEMM: + return False + if _rubin_mxfp8_is_available() and self.lowm_weight_scale is None: + return False + if self.weight.shape[0] % 128 and not self.fp8w_lowm_partial_n_tile: + return False + from .kernels import k3_fp8w_lowm_gemm + + limit = k3_fp8w_lowm_gemm.low_m_max_for(*self.weight.shape) + if self.lowm_weight_scale is not None: + limit = min(limit, k3_fp8w_lowm_gemm.max_m_for(*self.weight.shape)) + return 1 <= num_rows <= limit + + def takes_mxfp8_input(self, num_rows: int) -> bool: + """Whether a producer should hand :meth:`forward_mxfp8` the MXFP8 form of ``num_rows`` + rows: where :meth:`forward` would run the low-M kernel, whose in-kernel quantization the + MXFP8 GEMM replaces (the attention's first projection also asks + :meth:`takes_mid_mxfp8_input`).""" + return self._runs_low_m_kernel(num_rows) + + def takes_mid_mxfp8_input(self, num_rows: int) -> bool: + """Whether :meth:`forward` on ``num_rows`` BF16 rows runs the mid-M path of + ``k3_fp8w_lowm_gemm`` (above the low-M cut-over, up to the per-shape gate) with one + 128-row scale block, so that a producer may hand :meth:`forward_mxfp8` the rows' MXFP8 form + with the scales in that GEMM's layout (``k3_attn_res`` ``mx_atom``): the GEMM then runs + without its quantization kernel, bitwise the BF16 path.""" + return self.mid_operand_tactic(num_rows) is not None + + def mid_operand_tactic(self, num_rows: int): + """c240: the mid-M tactic :meth:`forward` runs on ``num_rows`` BF16 rows (SM107 low-M GEMM + switch on, above the low-M cut-over, up to the per-shape gate) when it has one 128-row + scale block, else None. A producer may then store the rows' MXFP8 form in that GEMM's + layout (``mid_prequant_plan``) for :meth:`forward_mxfp8` / :meth:`forward_mid_mxfp8`.""" + if not _K3_FP8W_LOWM_GEMM: + return None + if self.lowm_weight_scale is None: + return None + n, k = self.weight.shape + if n % 128 and not self.fp8w_lowm_partial_n_tile: + return None + from .kernels import k3_fp8w_lowm_gemm + + if ( + not k3_fp8w_lowm_gemm.low_m_max_for(n, k) + < num_rows + <= min(128, k3_fp8w_lowm_gemm.max_m_for(n, k)) + ): + return None + plan = k3_fp8w_lowm_gemm.mid_prequant_plan(num_rows, n, k) + return plan[0] if plan is not None and plan[1] == 1 else None + + def forward_mid_mxfp8( + self, x: torch.Tensor, xq: torch.Tensor, xs: torch.Tensor + ) -> torch.Tensor: + """``TRTLLM_K3_MID_OPROJ_PC`` (c240): :meth:`forward` on rows whose MXFP8 form a producer + stored with the scales in the mid-M GEMM's layout (``xs`` 1-D, see + :meth:`mid_operand_tactic`; ``x`` only gives the output shape): the mid-M GEMM without + its quantization kernel, with the all-reduce push epilogue where :attr:`next_ar_push` + takes it. Bitwise :meth:`forward` on the rows the producer quantized.""" + + m = xq.shape[0] + tactic = self.mid_operand_tactic(m) + if tactic is None or xs.dim() != 1: + raise RuntimeError("forward_mid_mxfp8 needs the mid-M operand layout") + x_flat = x.reshape(-1, x.shape[-1]) + push = self._take_next_ar_push(x_flat) + self.next_output = None + if push is not None: + k3_fp8w_mid_gemm_push_op( + xq.reshape(m, -1), + xs, + self.weight, + self.lowm_weight_scale, + push.workspace, + push.one_shot, + tactic, + ) + push.pushed = True + return push.output + out = torch.empty((m, self.out_features), dtype=torch.bfloat16, device=xq.device) + k3_fp8w_mid_gemm_op(xq.reshape(m, -1), xs, self.weight, self.lowm_weight_scale, out, tactic) + return out.reshape(x.shape[:-1] + (self.out_features,)) + + def _staged_above_cutover(self, num_rows: int) -> bool: + """Whether ``forward_mxfp8(..., stage_weights=True)`` serves ``num_rows`` rows above the + low-M kernel's cut-over with the staged MXFP8 low-M GEMM: :meth:`forward` runs the mid-M + path of ``k3_fp8w_lowm_gemm`` there, whose bits the staged kernel reproduces + (``k3_fp8w_lowm_mxfp8.staged_matches_mid_m``).""" + if not _K3_FP8W_LOWM_GEMM or self.lowm_weight_scale is None: + return False + from .kernels import k3_fp8w_lowm_gemm, k3_fp8w_lowm_mxfp8 + + n, k = self.weight.shape + if num_rows > k3_fp8w_lowm_gemm.max_m_for(n, k): + return False + return k3_fp8w_lowm_mxfp8.staged_matches_mid_m(num_rows, n, k) + + def takes_staged_mxfp8_input(self, num_rows: int) -> bool: + """Whether ``forward_mxfp8(..., stage_weights=True)`` runs the staged MXFP8 low-M GEMM on + the MXFP8 form of ``num_rows`` rows (bitwise :meth:`forward`): where :meth:`forward` runs + the low-M kernel, and above its cut-over where it runs the mid-M path with the same + split-K (up to 64 rows).""" + from .kernels import k3_fp8w_lowm_mxfp8 + + n, k = self.weight.shape + if self._runs_low_m_kernel(num_rows): + # Also where the shared-memory ring alone holds the CTA's whole weight slice. + return k3_fp8w_lowm_mxfp8.weights_on_chip(num_rows, n, k) + return self._staged_above_cutover(num_rows) + + def forward_mxfp8( + self, x: torch.Tensor, xq: torch.Tensor, xs: torch.Tensor, stage_weights: bool = False + ) -> torch.Tensor: + """:meth:`forward` on the BF16 rows ``x``, bit for bit, given their MXFP8 form ``(xq, + xs)`` (E4M3 ``[M, K]`` and one UE8M0 exponent per 1x32 group, e.g. from + the ``k3_attn_res_norm`` MXFP8 variants): the low-M GEMM on the prequantized operand where + :meth:`forward` would run the low-M kernel, else :meth:`forward` on ``x``. + ``stage_weights``: the low-M GEMM keeps its whole weight slice on chip before its + ``griddepcontrol.wait`` (``k3_fp8w_lowm_mxfp8.staged_mx_tactic``, the same bits), also + above the low-M cut-over where :meth:`takes_staged_mxfp8_input` says so.""" + from .kernels import k3_fp8w_lowm_mxfp8 + + x_flat = x.reshape(-1, x.shape[-1]) + if xs.dim() == 1 and self.takes_mid_mxfp8_input(x_flat.shape[0]): + # The producer stored the scales in the mid-M GEMM's layout. + from .kernels import k3_fp8w_lowm_gemm + + m = x_flat.shape[0] + tactic, _ = k3_fp8w_lowm_gemm.mid_prequant_plan(m, self.out_features, self.in_features) + self.next_ar_push = None + out = self._take_next_output(m) + if out is None: + out = torch.empty((m, self.out_features), dtype=torch.bfloat16, device=x.device) + k3_fp8w_mid_gemm_op( + xq.reshape(m, -1), xs, self.weight, self.lowm_weight_scale, out, tactic + ) + return out.reshape(x.shape[:-1] + (self.out_features,)) + allow_partial = self.fp8w_lowm_partial_n_tile + # The scale forward's low-M branch uses (lowm_weight_scale on SM107). + w_scale = self.packed_weight_scale + above_cutover = ( + stage_weights + and not self._runs_low_m_kernel(x_flat.shape[0]) + and self._staged_above_cutover(x_flat.shape[0]) + ) + if not ( + x.dtype == torch.bfloat16 + and (above_cutover or self._runs_low_m_kernel(x_flat.shape[0])) + and xq.shape[0] == x_flat.shape[0] + and _k3_lowm_gemm_can_run( + x_flat, self.weight, w_scale, allow_partial_n_tile=allow_partial + ) + and k3_fp8w_lowm_mxfp8.can_run_mxfp8( + xq, + xs, + self.weight, + w_scale, + allow_partial_n_tile=allow_partial, + above_cutover=above_cutover, + ) + ): + return self.forward(x) + out = self._take_next_output(x_flat.shape[0]) + # The MXFP8 GEMM has no all-reduce push epilogue: the caller runs the all-reduce. + self.next_ar_push = None + if out is None: + out = torch.empty( + (x_flat.shape[0], self.out_features), dtype=torch.bfloat16, device=x.device + ) + stage_weights = stage_weights and k3_fp8w_lowm_mxfp8.weights_on_chip( + x_flat.shape[0], self.out_features, self.in_features + ) + k3_fp8w_lowm_mxfp8_gemm_op(xq, xs, self.weight, w_scale, out, stage_weights=stage_weights) + return out.reshape(x.shape[:-1] + (self.out_features,)) + + def gated_applies(self, x: torch.Tensor, gate: torch.Tensor) -> bool: + """Whether :meth:`forward_gated` serves ``x`` and ``gate`` with one gated low-M GEMM. + + That is where :meth:`forward` would run the low-M GEMM on ``x * gate`` (same kernel + family and split), so the gated launch is bitwise equal to the multiply plus + :meth:`forward`. + """ + if x.dim() != 2 or not x.is_contiguous(): + return False + scale = self._lowm_forward_scale(x) + if scale is None: + return False + from .kernels import k3_fp8w_lowm_gemm + + return k3_fp8w_lowm_gemm.can_run_gated( + x, + gate, + self.weight, + scale, + allow_partial_n_tile=self.fp8w_lowm_partial_n_tile, + ) + + def forward_sigmoid(self, x: torch.Tensor) -> torch.Tensor: + """``forward(x).sigmoid()``, bit for bit; one low-M GEMM launch with a sigmoid + epilogue where :meth:`forward` would run the low-M GEMM.""" + x_flat = x.reshape(-1, x.shape[-1]) + scale = self._lowm_forward_scale(x_flat) + if scale is not None: + from .kernels import k3_fp8w_lowm_gemm + + if x_flat.shape[0] <= k3_fp8w_lowm_gemm.low_m_max_for( + self.out_features, self.in_features + ): + self.next_output = None + self.next_ar_push = None + out = k3_fp8w_lowm_gemm_op(x_flat, self.weight, scale, sigmoid=True) + return out.reshape(x.shape[:-1] + (self.out_features,)) + return self.forward(x).sigmoid() + + def forward_gated(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """``forward(x * gate)`` in one gated low-M GEMM launch (where :meth:`gated_applies`). + + ``x`` and ``gate`` are BF16 ``[M, in_features]``; the gate may be a row-strided view. + """ + ar_push = self._take_next_ar_push(x, gated=True) + if ar_push is not None: + self.next_output = None + return self._run_ar_push(x, ar_push, gate=gate) + return k3_fp8w_lowm_gemm_op( + x, + self.weight, + self.packed_weight_scale, + out=self._take_next_output(x.shape[0]), + gate=gate, + ) + + def gated_mxfp8_applies(self, x: torch.Tensor, gate: torch.Tensor) -> bool: + """Whether :meth:`forward_gated_mxfp8` runs the MXFP8 low-M GEMM on the producer's MXFP8 + form of ``bf16(x * gate)``: where :meth:`forward_gated` runs + the gated low-M GEMM, whose tiling the MXFP8 kernel shares (same bits).""" + if not self.gated_applies(x, gate): + return False + from .kernels import k3_fp8w_lowm_mxfp8 + + m, k = x.shape + try: + k3_fp8w_lowm_mxfp8.default_mx_tactic(m, self.out_features, k) + except ValueError: + return False + return k % 32 == 0 + + def forward_gated_mxfp8( + self, x: torch.Tensor, gate: torch.Tensor, xq: torch.Tensor, xs: torch.Tensor + ) -> torch.Tensor: + """:meth:`forward_gated` ``(x, gate)``, bit for bit, given the MXFP8 form ``(xq, xs)`` of + ``bf16(x * gate)`` (E4M3 ``[M, K]`` and one UE8M0 exponent per 1x32 group, the gated + low-M GEMM's own quantization): the MXFP8 low-M GEMM, with the all-reduce push epilogue + where :attr:`next_ar_push` is set, so the activation path after its + ``griddepcontrol.wait`` is two TMA loads instead of the BF16 loads, the gate multiply and + the quantization.""" + from .kernels import k3_fp8w_lowm_mxfp8 + + w_scale = self.packed_weight_scale + if w_scale is None or not k3_fp8w_lowm_mxfp8.can_run_mxfp8(xq, xs, self.weight, w_scale): + return self.forward_gated(x, gate) + ar_push = self._take_next_ar_push(x, gated=True) + if ar_push is not None: + self.next_output = None + k3_fp8w_lowm_mxfp8_push_op( + xq, xs, self.weight, w_scale, ar_push.workspace, ar_push.one_shot + ) + ar_push.pushed = True + return ar_push.output + out = self._take_next_output(x.shape[0]) + if out is None: + out = torch.empty( + (x.shape[0], self.out_features), dtype=torch.bfloat16, device=x.device + ) + k3_fp8w_lowm_mxfp8_gemm_op(xq, xs, self.weight, w_scale, out) + return out + + def forward_prequantized( + self, + activation: torch.Tensor, + activation_scale: torch.Tensor, + ) -> torch.Tensor: + """Project an FP8 activation a fused producer already quantized. + + Args: + activation: ``[..., in_features]`` ``float8_e4m3fn`` activation. + activation_scale: ``uint8`` UE8M0 scales in the CuTe K32 R128c4 + layout that ``fp8_quantize_1x128_cutedsl_ue8m0`` produces. + """ + push, self.next_ar_push = self.next_ar_push, None + if activation.dtype is not torch.float8_e4m3fn: + raise TypeError("prequantized Kimi K3 activation must be float8_e4m3fn") + if not self.supports_prequantized_input: + raise RuntimeError("prequantized Kimi K3 activation requires the Rubin CuTe GEMM") + if activation_scale.dtype is not torch.uint8: + raise TypeError("prequantized Kimi K3 activation scale must be uint8 R128c4") + rows = activation.reshape(-1, activation.shape[-1]) + tactic = self._mid_prequantized_push_tactic(rows, activation_scale, push) + if tactic is not None: + # The mid-M GEMM pushes into the o_proj all-reduce. + k3_fp8w_mid_gemm_push_op( + rows, + activation_scale.reshape(-1), + self.weight, + self.lowm_weight_scale, + push.workspace, + push.one_shot, + tactic, + ) + push.pushed = True + return push.output + out = k3_mxfp8_gemm_op( + activation.reshape(-1, activation.shape[-1]), + self.weight, + activation_scale, + self.weight_scale, + ) + return out.reshape(activation.shape[:-1] + (self.out_features,)) + + +def _swap_linear_to_fp8_weight_read( + parent: nn.Module, + attr: str, + linear_types: Tuple[type, ...] = (nn.Linear,), +) -> int: + """Replace ``parent.`` with an FP8 weight-read module if it is a + plain linear of one of ``linear_types``; return the number of modules + converted (0 or 1), so callers can accumulate a conversion count. + + Transfer checkpoint FP8 storage or release the original BF16 storage + after conversion. The loader's transient parameter map must not retain + an unused BF16 copy until loading completes. + """ + child = getattr(parent, attr, None) + if not isinstance(child, linear_types): + return 0 + converted = _Fp8BlockScaleWeightReadLinear.from_linear(child) + converted.k3_latent_fused = attr == "routed_expert_down_proj" and _K3_LD_FUSED + setattr(parent, attr, converted) + child.weight.data = child.weight.data.new_empty(0) + return 1 + + +def _has_weights(module: nn.Module) -> bool: + """False once ``modeling_utils.remove_weights()`` has stripped a module. + + Post-load finalization walks every decoder layer, so it must skip layers + whose parameters were dropped — the layer-wise benchmarks keep only the + profiled slice resident. + """ + return not getattr(module, "_weights_removed", False) + + +def _convert_moe_mlps_to_fp8_weight_read( + model: nn.Module, include_fused_gate_up: bool = True +) -> int: + """Swap the replicated MoE-layer MLP projections to an FP8 weight read. + + Targets the shared-expert MLP (gate/up/down) and the latent up/down + projection on every MoE layer — the bias-free BF16 projections that + attention data-parallelism re-reads in full each decode step. Attention + follows checkpoint quantization separately. Routed experts and the dense + layer-0 MLP are not converted. Returns the number of projections converted. + """ + count = 0 + + for layer in model.layers: + if not _has_weights(layer): + continue + moe = getattr(layer, "block_sparse_moe", None) + if moe is None: + continue + shared = getattr(moe, "shared_experts", None) + if shared is not None: + # GatedMlp fuses gate and up into gate_up_proj; keep the split + # names too so either MLP layout converts. The fused gate_up read + # only pays off when attention DP re-reads it per rank per step; + # under TP the bf16 GEMM overlaps on the aux stream and the FP8 + # quantize+GEMM would serialize onto the critical path. + shared_attrs = ( + ("gate_proj", "up_proj", "gate_up_proj", "down_proj") + if include_fused_gate_up + else ("gate_proj", "up_proj", "down_proj") + ) + for attr in shared_attrs: + child = getattr(shared, attr, None) + if isinstance(child, Bf16Linear) and child.tp_size != 1: + continue + count += _swap_linear_to_fp8_weight_read( + shared, attr, linear_types=(nn.Linear, Bf16Linear) + ) + for attr in ("routed_expert_down_proj", "routed_expert_up_proj"): + count += _swap_linear_to_fp8_weight_read(moe, attr) + + # Return the freed BF16 blocks to the driver so the raw (non-caching- + # allocator) allocations made during executor creation succeed on the + # tight DEP16 memory headroom. + if count: + gc.collect() + torch.cuda.empty_cache() + return count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda.py new file mode 100644 index 000000000000..2e1dfc33d450 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda.py @@ -0,0 +1,1858 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi Delta Attention (KDA) for the Kimi K3 (k3_rubin) text model. + +``K3KDALinearAttention`` holds the layer's weights in K3's own containers (``layers.py``) and +runs one explicit kernel chain per phase (prefill, decode, multi-token verify) on the package's +KDA kernels (``kernels/kda``, ``ops/ssm/_kda_prefill.py``). The +engine-owned cache manager stays main's ``modules/kimi_kda/cache_manager.py``. + +KDA is the linear-attention block used at ``linear_attn_config.kda_layers`` +positions in the Kimi K3 text-core. It carries a short-convolution state and a +delta-rule recurrent state per layer, so it follows the hybrid-cache / +mamba ownership pattern rather than the paged-KV FMHA attention-backend +interface. + +Kimi K3 KDA layer: K3-owned weights and one explicit kernel chain per phase. + +Weights (checkpoint names under ``model.layers..self_attn.``, this rank's head shard): + +* ``q_proj`` / ``k_proj`` / ``v_proj`` / ``g_proj`` / ``o_proj``: FP8 128 x 128 block-scale + projections (``layers.Fp8BlockScaleLinear``). :meth:`finalize_decode_weights` fuses + ``[q | k | v | g]`` into ``qkvg_proj`` and repoints the four to row views of it; the model loader + then swaps every FP8 projection to the K3 FP8 weight-read GEMM (``text_model``). +* ``f_a_proj`` / ``b_proj`` / ``f_b_proj``: BF16 (``layers.Bf16Linear``); ``[f_a | b]`` is + fused into ``_bfa_proj_weight`` (rows padded to 8). +* ``q_conv1d`` / ``k_conv1d`` / ``v_conv1d``: depthwise short-convolution weights + (``layers.ShortConvWeight``); ``o_norm``: the gated output RMSNorm weight + (``layers.RMSNormWeight``); ``A_log`` / ``dt_bias``. + +Kernel chains (``H`` heads of this rank, ``D = 128``; ``x`` the BF16 rows): + +* **Prefill** (:meth:`forward_prefill`): QKVG GEMM -> ``F.linear`` ``[f_a | b]`` and ``F.linear`` + ``f_b`` -> short convolution with its conv-pool write-back, SiLU and Q/K L2 norm + (``fused_kda_conv_prefill``; or ``extract_transpose_prefill_slice`` + ``causal_conv1d_fwd`` + + ``fused_kda_post_conv``) -> ``kda_prefill`` (fused K1-K3, A_kk inverse, persistent + K4) on fp32 state rows (a BF16 pool is staged through an fp32 copy and written back) -> gated + RMSNorm (``rms_norm_gated_token_major``) -> o_proj (and the TP all-reduce). +* **Decode** (:meth:`forward_decode`, one token per request): QKVG GEMM, next to ``[f_a | b]`` / + ``f_b`` on the projection aux stream -> the fused CUDA decode kernel (convolution, gate, + recurrence, gated RMSNorm, conv / state pools updated by slot) -> o_proj. +* **Verify** (:meth:`forward_verify_fused`, ``1 + num_spec`` tokens per request): QKVG GEMM (or its + MXFP8 hand-off form) next to the gate kernel on the aux stream (``k3_kda_fab_fa`` -> the forget + gate computed in the verify kernel; ``k3_kda_fab_proj``; or ``F.linear`` ``[f_a | b]`` and + ``f_b``), the L2 prefetch of the next KDA layer's weights -> ``kda_mtp_decode`` (forget gate + from ``f_a`` in the kernel, or BF16 core, or MXFP8 core for o_proj) on the replay caches -> + o_proj, with the all-reduce pushed into its Lamport buffers where the GEMM can. + +The kernels are called through the target's ``ops/ssm`` entries (and the family's +``causal_conv1d_fwd`` / ``rms_norm_gated_token_major``). + +Cache ownership: three short-convolution states in the combined ``[slots, 3D, W - 1]`` BF16 pool +and one delta-rule recurrent state in the ``[slots, H, V, K]`` pool (fp32 or BF16), plus the replay +caches of the fused verify kernel, all allocated by the engine's ``KimiK3HybridCacheManager``. + +Host-prepared scheduling metadata for packed KDA prefill. + +Kimi K3 Rubin merge: main's FlashInfer KDA prefill metadata (``kda_max_seqlen``, +``kda_seq_order``, ``kda_cu_chunks``, the 16-byte aligned ``generation_state_indices``) +plus the Rubin agentic branch's metadata for its CuTe DSL prefill (the package's default +prefill, the ``kda_prefill`` entry): the FLA-style ``kda_chunk_indices`` (filled on the +device), ``kda_varlen_is_aligned`` / ``kda_single_sequence_length``, chunk alignment 64 for +the warmup, the int64 ``state_indices_long`` mirror (packed into the step's input upload while +the engine packs them, ``pyexecutor.step_inputs``) and the mid-chunk snapshot plan +(``kda_snapshot``: the in-forward snapshots ``MambaHybridCacheManagerV2`` plans with block reuse, +see :meth:`K3MambaMetadata._prepare_inforward_split`; otherwise ``None``). +""" + +from __future__ import annotations + +import functools +import math +import weakref +from dataclasses import dataclass, replace +from itertools import accumulate +from typing import Dict, List, NamedTuple, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + eager_on_graph, + is_in_breakable_cuda_graph, +) +from tensorrt_llm._torch.utils import get_model_extra_attrs +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + +from ..ops.norm.rms_norm_gated_token_major import rms_norm_gated_token_major +from ..ops.ssm.causal_conv1d_fwd import causal_conv1d_fwd +from .features import active as _features +from .kernels.kda._kda_kernels import ( + PAD_SLOT_ID, + K3KDAKernelDispatch, + extract_transpose_prefill_slice, + kda_mtp_verify_ctas, + reset_recurrent_state_rows, +) +from .layers import ( + Bf16Linear, + Fp8BlockScaleLinear, + LamportPending, + RMSNormWeight, + ShortConvWeight, + TpAllReduce, + finish_lamport_push, + fork_join, + lamport_push_target, + lamport_push_tensor, +) + +_F = _features() + +# Prefill chunk of the kda_prefill entry. +_PREFILL_CHUNK = 64 +# features.kda_prepad_prefill: the post-convolution kernel writes a single unaligned sequence's +# Q/K/V/beta zero-padded to the prefill op's chunk instead of the op padding copies of them. +_KDA_PREPAD_PREFILL = _F.kda_prepad_prefill +# features.kda_fused_conv_prefill: the short convolution (with its conv-pool read and +# write-back), SiLU, Q/K L2 norm and the Q/K/V layout change run as one Triton pass over the QKVG +# projection output (fused_kda_conv_prefill) instead of extract_transpose_prefill_slice + +# causal_conv1d_fwd + fused_kda_post_conv. +_KDA_FUSED_CONV_PREFILL = _F.kda_fused_conv_prefill +# Single-token decode: [f_a | b] -> f_b runs on the projection aux stream next to the QKVG GEMM up +# to this many rows. +_KDA_BFA_MULTISTREAM_MAX_ROWS = 128 +# Decode verify without the forget-gate round trip. Up to _KDA_FAB_FUSED_MAX_ROWS rows, +# k3_kda_fab runs next to the QKVG GEMM on the aux stream: in f_a mode (BF16 [f_a | b] only, +# the verify kernel computes the forget gate from f_a itself) when that kernel runs in one wave +# or the QKVG GEMM has left the low-M kernel (more rows than _KDA_FUSED_FB_LOWM_ROWS), else as +# the full fab kernel (beta and the forget gate from one fused [f_a | b] -> f_b kernel). The f_b +# MMA costs every wave of the verify kernel, so multi-wave low-M batches keep the full fab kernel. +# Larger row counts run the [f_a | b] -> f_b chain on the aux stream as well. The fab kernels take +# at most 16 heads (the Helix head shards); other shapes keep the [f_a | b] -> f_b chain. +_KDA_FAB_FUSED_MAX_ROWS = 256 +_KDA_FUSED_FB_LOWM_ROWS = 224 +# Decode verify: once this layer's QKVG projection and [f_a | b] kernel are done, the projection +# aux stream prefetches the next KDA layer's QKVG, [f_a | b], f_b and o_proj weights into L2 +# (k3_l2_prefetch) while this layer's recurrence and all-reduce leave HBM idle. Hints only: no +# result depends on them. Up to _KDA_L2_PREFETCH_MAX_ROWS verify rows: the prefetched lines stay +# in L2 through the MoE block, whose weight streams then run about 1 % slower. The prefetch kernel +# stays on the aux stream, which the MoE block of the same layer reuses for its shared experts and +# joins back; no main-stream kernel waits for it. +_KDA_L2_PREFETCH = _F.kda_l2_prefetch +_KDA_L2_PREFETCH_MAX_ROWS = 16 +# KDA mixers of one model by layer index (keyed by the model's text config), for the prefetch of +# the next KDA layer's weights. +_KDA_LAYERS: Dict[int, "weakref.WeakValueDictionary[int, K3KDALinearAttention]"] = {} + +_FP8_PROJECTIONS = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") +_QKVG_PROJECTIONS = ("q_proj", "k_proj", "v_proj", "g_proj") + + +def _use_fused_fab( + x: torch.Tensor, bfa_weight: torch.Tensor, fb_weight: torch.Tensor, num_heads: int +) -> bool: + """Whether the ``k3_kda_fab_proj`` entry accepts these operands.""" + if not x.is_contiguous(): + return False + from .kernels import k3_kda_fab + + return k3_kda_fab.can_use_k3_kda_fab( + x.reshape(-1, x.shape[-1]), bfa_weight, fb_weight, num_heads + ) + + +@functools.cache +def _num_sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def _verify_kernel_single_wave(num_heads: int, num_requests: int, device: torch.device) -> bool: + """Whether the v2 verify kernel's grid (``kda_mtp_verify_ctas``) fits in one wave of one CTA + per SM.""" + ctas = kda_mtp_verify_ctas(num_heads, num_requests, device) + index = device.index if device.index is not None else torch.cuda.current_device() + return ctas <= _num_sms(index) + + +def _use_fab_fa( + x: torch.Tensor, bfa_weight: torch.Tensor, fb_weight: torch.Tensor, num_heads: int +) -> bool: + """Whether the verify gate comes from the ``k3_kda_fab_fa`` entry + the in-kernel f_b. + + ``x`` is the ``[requests, steps, hidden]`` verify input.""" + if not x.is_contiguous(): + return False + num_rows = x.shape[0] * x.shape[1] + if num_rows <= _KDA_FUSED_FB_LOWM_ROWS and not _verify_kernel_single_wave( + num_heads, x.shape[0], x.device + ): + return False + if ( + fb_weight.dtype != torch.bfloat16 + or not fb_weight.is_contiguous() + or tuple(fb_weight.shape) != (num_heads * 128, 128) + or fb_weight.data_ptr() % 16 + ): + return False + from .kernels import k3_kda_fab + + return k3_kda_fab.can_use_k3_kda_fab_fa( + x.reshape(-1, x.shape[-1]), bfa_weight, num_heads, fb_weight + ) + + +def _o_proj_takes_mxfp8(o_proj: nn.Module, num_rows: int) -> bool: + """Whether ``o_proj`` takes ``num_rows`` rows as the Rubin MXFP8 operand of the verify + kernel's epilogue (``forward_prequantized``) rather than BF16 (the K3 FP8 weight-read + projection decides per row count, see ``takes_prequantized_input``).""" + takes = getattr(o_proj, "takes_prequantized_input", None) + if takes is not None: + return bool(takes(num_rows)) + return bool(getattr(o_proj, "supports_prequantized_input", False)) + + +class _KdaPrequantizedCore(NamedTuple): + """MXFP8 KDA core emitted by the fused verify epilogue. + + ``activation`` is the ``[tokens, H * V]`` E4M3 post-o_norm core and + ``scale`` its UE8M0 1x128 scales in the flat uint8 R128c4 layout that + ``o_proj.forward_prequantized`` consumes. + """ + + activation: torch.Tensor + scale: torch.Tensor + + +@dataclass +class _KDAPrefillState: + """The fp32 state rows the ``kda_prefill`` entry updates in place: the pool itself + (``kernel_pool`` / ``kernel_indices``), or ``staged_state``, a dense fp32 copy of the + addressed rows (plus the snapshot rows of a plan) that is written back afterwards.""" + + kernel_pool: torch.Tensor + kernel_indices: torch.Tensor + staged_state: torch.Tensor | None + # With a mid-chunk snapshot plan: the snapshot rows of ``kernel_pool`` (the pool's own slots, + # or the rows after the sequences' in the staged copy). + snapshot_rows: torch.Tensor | None = None + + +def _stage_state_rows(ssm_pool: torch.Tensor, slot_indices: torch.Tensor) -> torch.Tensor: + """Dense fp32 copy of the addressed recurrent-state rows (``_writeback_state_rows`` rounds it + back into the pool).""" + return ssm_pool.index_select(0, slot_indices.long()).float() + + +def _writeback_state_rows( + ssm_pool: torch.Tensor, slot_indices: torch.Tensor, rows: torch.Tensor +) -> None: + """Scatter dense state rows back into the pool, rounded to the pool dtype.""" + ssm_pool.index_copy_(0, slot_indices.long(), rows.to(ssm_pool.dtype)) + + +def _stage_state_rows_with_snapshot_rows( + ssm_pool: torch.Tensor, slot_indices: torch.Tensor, num_snapshots: int +) -> torch.Tensor: + """``_stage_state_rows`` plus ``num_snapshots`` uninitialized rows after them. + + The prefill kernel writes mid-chunk snapshots into the extra rows. + """ + num_rows = slot_indices.shape[0] + staged = torch.empty( + (num_rows + num_snapshots, *ssm_pool.shape[1:]), + dtype=torch.float32, + device=ssm_pool.device, + ) + staged[:num_rows].copy_(ssm_pool.index_select(0, slot_indices.long())) + return staged + + +def _snapshot_conv_windows( + conv_input: torch.Tensor, + conv_pool: torch.Tensor, + snapshot: "KDASnapshotPlan", + token_major: bool, +) -> torch.Tensor: + """The raw convolution window ending before each snapshot point, ``[n, 3D, W - 1]`` in + ``conv_pool``'s dtype. + + ``conv_input`` is the pre-convolution projection, ``[3D, T]`` or with + ``token_major`` ``[T, 3D]``; ``conv_pool`` is ``[slots, 3D, W - 1]``. A + snapshot point is at least one KDA chunk into its sequence's step, so the + window never reaches the sequence's previous convolution state. + """ + window = conv_pool.shape[-1] + offsets = torch.arange(-window, 0, device=conv_input.device) + token_indices = (snapshot.token_indices[:, None] + offsets).reshape(-1) + if token_major: + rows = conv_input.index_select(0, token_indices) + windows = rows.view(snapshot.num_snapshots, window, -1).transpose(1, 2) + else: + columns = conv_input.index_select(1, token_indices) + windows = columns.view(-1, snapshot.num_snapshots, window).transpose(0, 1) + return windows.to(conv_pool.dtype) + + +def _swap_inforward_conv_windows( + conv_pool: torch.Tensor, + stage_conv: torch.Tensor, + snapshot: "KDASnapshotPlan", + windows: torch.Tensor, +) -> None: + """Move each snapshotted sequence's chunk-end convolution window from its live slot + (``snapshot.slots``) to its staging row and put the window ending at the snapshot point + (``windows``, :func:`_snapshot_conv_windows`) in the live slot.""" + stage_conv.index_copy_(0, snapshot.stage_rows, conv_pool.index_select(0, snapshot.slots)) + conv_pool.index_copy_(0, snapshot.slots, windows) + + +def _merge_row_weights(modules: Tuple[nn.Module, ...], pad_rows_to: int = 1) -> torch.Tensor: + """Concatenate the ``weight`` rows of ``modules`` (zero rows up to a multiple of + ``pad_rows_to``) and repoint each module's ``weight`` to its row view of the result.""" + weights = [module.weight.data for module in modules] + padding = (-sum(weight.shape[0] for weight in weights)) % pad_rows_to + if padding: + weights.append(weights[0].new_zeros((padding, weights[0].shape[1]))) + fused = torch.cat(weights, dim=0).contiguous() + offset = 0 + for module in modules: + rows = module.weight.shape[0] + module.weight.data = fused[offset : offset + rows] + offset += rows + return fused + + +def _extract_kda_extra_attrs(layer_idx: str): + """Resolve the live attention metadata and KDA module for ``layer_idx``.""" + extra_attrs = get_model_extra_attrs() + assert extra_attrs is not None, "Model extra attrs is not set" + + metadata_ref = extra_attrs.get("attention_metadata") + assert metadata_ref is not None, "Attention metadata is not set" + metadata = metadata_ref() + assert isinstance(metadata, AttentionMetadata), "Metadata must be AttentionMetadata" + + kda_layers = extra_attrs.get("kda_layers") + assert kda_layers is not None, "KDA layer registry is not set" + kda_layer_ref = kda_layers.get(layer_idx) + assert kda_layer_ref is not None, f"Cannot find KDA layer for layer {layer_idx}" + kda_layer = kda_layer_ref() + assert isinstance(kda_layer, K3KDALinearAttention), "KDA layer must be K3KDALinearAttention" + return metadata, kda_layer + + +def kda_core_inplace(hidden_states: torch.Tensor, layer_idx: str, output: torch.Tensor) -> None: + """Run the metadata-dependent KDA core and write it into ``output``.""" + metadata, kda_layer = _extract_kda_extra_attrs(layer_idx) + kda_layer._forward_impl(hidden_states, metadata, output=output) + + +maybe_bcg_kda_core_inplace = eager_on_graph(kda_core_inplace) + + +def _check_kda_shape(lin: dict) -> None: + """The KDA kernels take Kimi K3's shape only.""" + gate_lower_bound = lin.get("gate_lower_bound", None) + problems = [] + if lin["head_dim"] != 128: + problems.append(f"head_dim {lin['head_dim']} (the kernels take 128)") + if lin["short_conv_kernel_size"] != 4: + problems.append(f"short_conv_kernel_size {lin['short_conv_kernel_size']} (4)") + if not lin.get("use_full_rank_gate", True): + problems.append("a low-rank output gate (the full-rank g_proj)") + if gate_lower_bound is None or not (math.isfinite(gate_lower_bound) and gate_lower_bound < 0): + problems.append(f"gate_lower_bound {gate_lower_bound} (a finite negative bound)") + if problems: + raise ValueError("Kimi K3 KDA: unsupported linear_attn_config: " + "; ".join(problems)) + + +class K3KDALinearAttention(nn.Module): + """One Kimi K3 KDA layer (head-TP shard) on K3-owned weights.""" + + def __init__( + self, + cfg, + layer_idx: int, + mapping=None, + allreduce_strategy=AllReduceStrategy.AUTO, + aux_stream: Optional[torch.cuda.Stream] = None, + model_config: Optional[ModelConfig] = None, + reduce_output: bool = True, + ) -> None: + super().__init__() + lin = cfg.linear_attn_config + _check_kda_shape(lin) + self.hidden_size = cfg.hidden_size + self.layer_idx = layer_idx + self.head_dim = lin["head_dim"] + self.head_k_dim = self.head_dim + self.conv_size = lin["short_conv_kernel_size"] + self.gate_lower_bound = lin["gate_lower_bound"] + self.rms_norm_eps = cfg.rms_norm_eps + + if mapping is not None and mapping.tp_size > 1 and not mapping.enable_attention_dp: + self._kda_tp_size = mapping.tp_size + else: + self._kda_tp_size = 1 + self._kda_tp_rank = mapping.tp_rank if self._kda_tp_size > 1 else 0 + # reduce_output=False returns the row-sharded o_proj partials; the + # caller owns the reduction (helix_a2a reduce-scatters them). + self._o_allreduce = ( + TpAllReduce(mapping, torch.bfloat16, allreduce_strategy) + if self._kda_tp_size > 1 and reduce_output + else None + ) + # Set by the decoder layer for one forward: a pushed o_proj all-reduce may be returned + # unfinished, as a LamportPending that A_mid finishes. + self.ar_finish_deferred = False + + num_heads = lin["num_heads"] + if num_heads % self._kda_tp_size != 0: + raise ValueError( + f"KDA num_heads {num_heads} not divisible by tp_size {self._kda_tp_size}" + ) + self.num_heads = num_heads // self._kda_tp_size + projection_size = self.num_heads * self.head_dim + self.proj_size = projection_size + + projection_quant_configs = (model_config.quant_config_dict or {}) if model_config else {} + default_quant_config = ( + model_config.quant_config if model_config else None + ) or QuantConfig() + declared = { + name: projection_quant_configs.get(name, default_quant_config).quant_algo + for name in _FP8_PROJECTIONS + } + if any(algo != QuantAlgo.FP8_BLOCK_SCALES for algo in declared.values()): + raise ValueError( + "Kimi K3 KDA takes FP8 128 x 128 block-scale q/k/v/g/o projections (the " + f"checkpoint's FP8_PB_WO); declared {declared}" + ) + if self.hidden_size % 128 or projection_size % 128: + raise ValueError("KDA FP8 projection shards must be aligned to 128x128 blocks") + + bf16 = torch.bfloat16 + # This rank's shard (the loader slices the checkpoint by head), so the containers hold + # rank-local dimensions. + self.q_proj = Fp8BlockScaleLinear(self.hidden_size, projection_size) + self.k_proj = Fp8BlockScaleLinear(self.hidden_size, projection_size) + self.v_proj = Fp8BlockScaleLinear(self.hidden_size, projection_size) + self.q_conv1d = ShortConvWeight(projection_size, self.conv_size, dtype=bf16) + self.k_conv1d = ShortConvWeight(projection_size, self.conv_size, dtype=bf16) + self.v_conv1d = ShortConvWeight(projection_size, self.conv_size, dtype=bf16) + # Initialized (checkpoint loading overwrites them) with single random-init ops, which + # MetaInitMode allows: uninitialized NaN bit patterns would poison randomly built layers. + self.A_log = nn.Parameter( + torch.log(torch.empty(self.num_heads, dtype=bf16).uniform_(1, 16)) + ) + self.f_a_proj = Bf16Linear(self.hidden_size, self.head_dim, dtype=bf16) + self.f_b_proj = Bf16Linear(self.head_dim, projection_size, dtype=bf16) + self.dt_bias = nn.Parameter( + torch.empty(projection_size, dtype=bf16).uniform_(math.log(1e-3), math.log(1e-1)) + ) + self.b_proj = Bf16Linear(self.hidden_size, self.num_heads, dtype=bf16) + self.g_proj = Fp8BlockScaleLinear(self.hidden_size, projection_size) + self.o_norm = RMSNormWeight(self.head_dim, self.rms_norm_eps, dtype=bf16) + self.o_proj = Fp8BlockScaleLinear(projection_size, self.hidden_size) + # Built from q/k/v/g by finalize_decode_weights: the fused [q | k | v | g] projection and + # its output sections. + self.qkvg_proj: Optional[nn.Module] = None + self.qkvg_split_sizes: Optional[List[int]] = None + + self._dispatch = K3KDAKernelDispatch() + + # Built after checkpoint load (finalize_decode_weights): the fused BF16 [f_a | b] weight + # (rows padded to 8) and the kernel-layout constants. + self._bfa_proj_weight: Optional[torch.Tensor] = None + self._w_q_t = self._w_k_t = self._w_v_t = None + self._A_log_f32 = self._dt_bias_f32 = self._onorm_w_f32 = None + # Fork/join state for overlapping the small [f_a | b] -> f_b chain + # with the wide QKVG projection during CUDA-graph execution. + self._projection_aux_stream = aux_stream + self._projection_fork_event = torch.cuda.Event() + self._projection_join_event = torch.cuda.Event() + # Output buffer of the in-place decode kernel outside breakable CUDA graphs. + self._o_dense: Optional[torch.Tensor] = None + # [3D, W] BF16 conv weight of causal_conv1d_fwd, its [W, 3D] copy for + # fused_kda_conv_prefill, and the fp32 [D, W] weights of the verify kernels. + self._packed_conv_weight: Optional[torch.Tensor] = None + self._fused_conv_weight_t: Optional[torch.Tensor] = None + self._mtp_conv_weights: Optional[Tuple[torch.Tensor, ...]] = None + # L2 prefetch of the next KDA layer's projection weights (features.kda_l2_prefetch). + self._kda_group = _KDA_LAYERS.setdefault(id(cfg), weakref.WeakValueDictionary()) + self._kda_group[layer_idx] = self + self._l2_prefetch_ranges_cache: Optional[List[Tuple[int, int]]] = None + self._l2_prefetch_start = torch.cuda.Event() + # Set by the decoder layer for the next forward only (features.attn_res_mxfp8): the + # MXFP8 form ``(xq, xs)`` of the generation rows, which the fused QKVG projection of the + # verify path consumes instead of quantizing the BF16 rows (see qkvg_takes_mxfp8). + self.next_input_mxfp8: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + + self.register_to_config = False + self.layer_idx_str = str(layer_idx) + if model_config is not None: + kda_layers = model_config.extra_attrs.setdefault("kda_layers", {}) + suffix = 0 + while self.layer_idx_str in kda_layers: + self.layer_idx_str = f"{layer_idx}_{suffix}" + suffix += 1 + kda_layers[self.layer_idx_str] = weakref.ref(self) + self.register_to_config = True + + # ------------------------------------------------------------------ + # Post-load weight preparation. + # ------------------------------------------------------------------ + + def finalize_decode_weights(self) -> None: + """Build the fused projection weights and kernel constants after weight load. + + 1. ``qkvg_proj``: one FP8 ``[q | k | v | g]`` projection from the checkpoint codes and + scales (the swap to the FP8 weight-read GEMM transforms its scales afterwards); the + four sources become row views of it. ``_bfa_proj_weight``: BF16 ``[f_a | b]`` with + rows padded to 8; ``f_a_proj`` / ``b_proj`` become row views of it. + 2. Kernel-layout constants: transposed BF16 conv weights ``[W, D]``, fp32 ``A_log`` / + ``dt_bias`` / ``o_norm.weight`` and the packed / fp32 conv weights of the prefill and + verify kernels. + """ + if self.q_proj.weight.device.type != "cuda": + return + with torch.no_grad(): + sources = tuple(getattr(self, name) for name in _QKVG_PROJECTIONS) + if all(isinstance(module, Fp8BlockScaleLinear) for module in sources): + self.qkvg_proj = self._fuse_checkpoint_projections(sources) + self.qkvg_split_sizes = [module.out_features for module in sources] + # Without it the four sources keep their own storage next to the fused copy. + self._alias_projections_into_fused(sources, self.qkvg_proj) + self._bfa_proj_weight = _merge_row_weights((self.f_a_proj, self.b_proj), pad_rows_to=8) + self._build_decode_kernel_constants() + + @staticmethod + def _fuse_checkpoint_projections( + modules: Tuple[Fp8BlockScaleLinear, ...], + ) -> Fp8BlockScaleLinear: + """One FP8 projection over the rows of ``modules`` (block-aligned): their checkpoint codes + and FP32 scale grids, copied before any scale preparation.""" + if any(module.out_features % 128 for module in modules[:-1]): + raise ValueError("KDA FP8 fusion boundaries must be aligned to 128 rows") + first = modules[0] + fused = Fp8BlockScaleLinear( + first.in_features, + sum(module.out_features for module in modules), + device=first.weight.device, + ) + offset = 0 + scale_offset = 0 + for module in modules: + rows = module.weight.shape[0] + fused.weight.data[offset : offset + rows].copy_(module.weight.data) + offset += rows + scale_rows = module.weight_scale.shape[0] + fused.weight_scale.data[scale_offset : scale_offset + scale_rows].copy_( + module.weight_scale.data + ) + scale_offset += scale_rows + fused.input_scale.data = first.input_scale.data + fused.inv_input_scale.data = first.inv_input_scale.data + return fused + + @staticmethod + def _alias_projections_into_fused(modules: Tuple[nn.Module, ...], fused: nn.Module) -> None: + """Repoint the FP8 source projections at row views of the fused codes and scale grid.""" + offset = 0 + scale_offset = 0 + for module in modules: + rows = module.weight.shape[0] + module.weight.data = fused.weight.data[offset : offset + rows] + offset += rows + scale_rows = module.weight_scale.shape[0] + module.weight_scale.data = fused.weight_scale.data[ + scale_offset : scale_offset + scale_rows + ] + scale_offset += scale_rows + + def _build_decode_kernel_constants(self) -> None: + """Kernel-layout constants of the decode / verify / prefill kernels.""" + self._w_q_t = ( + self.q_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + self._w_k_t = ( + self.k_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + self._w_v_t = ( + self.v_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + self._A_log_f32 = self.A_log.detach().float().contiguous() + self._dt_bias_f32 = self.dt_bias.detach().float().contiguous() + self._onorm_w_f32 = self.o_norm.weight.detach().float().contiguous() + # Built eagerly too, so the first verify call never allocates (a capture-unsafe lazy + # allocation). + self._build_mtp_conv_weights() + + def _build_mtp_conv_weights(self) -> None: + """Prebuild the packed-prefill and fused-verify convolution weights. + + Building them lazily at first use would allocate at runtime; under + CUDA graph capture that would bake capture-pool pointers into the + cached tensors. + """ + conv_weights = tuple( + conv.weight.detach().squeeze(1) + for conv in (self.q_conv1d, self.k_conv1d, self.v_conv1d) + ) + self._packed_conv_weight = torch.cat(conv_weights, dim=0).to(torch.bfloat16).contiguous() + if _KDA_FUSED_CONV_PREFILL: + # Channel-contiguous taps load in the input tile's layout in the fused kernel. + self._fused_conv_weight_t = self._packed_conv_weight.t().contiguous() + self._mtp_conv_weights = tuple(weight.float().contiguous() for weight in conv_weights) + + def _require_finalized(self) -> None: + if self.qkvg_proj is None or self._bfa_proj_weight is None or self._onorm_w_f32 is None: + raise RuntimeError( + "Kimi K3 KDA runs on the fused projections and kernel constants that " + "finalize_decode_weights() builds after weight load" + ) + + # ------------------------------------------------------------------ + # MXFP8 hand-off of the decoder layer (features.attn_res_mxfp8). + # ------------------------------------------------------------------ + + def qkvg_takes_mxfp8(self, num_rows: int) -> bool: + """Whether the fused FP8 QKVG projection of ``num_rows`` rows can take their MXFP8 form + (``next_input_mxfp8``) instead of the BF16 rows: it runs the low-M kernel on them (or the + staged kernel above the low-M cut-over, :meth:`qkvg_stages_weights`).""" + takes = getattr(self.qkvg_proj, "takes_mxfp8_input", None) + if takes is None: + return False + return takes(num_rows) or self.qkvg_stages_weights(num_rows) or self.qkvg_mx_atom(num_rows) + + def qkvg_mx_atom(self, num_rows: int) -> bool: + """Whether the fused QKVG projection takes the MXFP8 rows with the scales in the mid-M + GEMM's layout (``takes_mid_mxfp8_input``).""" + takes_mid = getattr(self.qkvg_proj, "takes_mid_mxfp8_input", None) + return takes_mid is not None and bool(takes_mid(num_rows)) + + def qkvg_stages_weights(self, num_rows: int) -> bool: + """Whether the fused QKVG projection of ``num_rows`` MXFP8 rows stages its weights on chip + before its ``griddepcontrol.wait``: then the attn_res in front of it may release it at + kernel entry. + + On the MXFP8 hand-off (``next_input_mxfp8``) of a decode verify the fused QKVG low-M GEMM + keeps its whole weight slice on chip before its wait (k3_fp8w_lowm_mxfp8.staged_mx_tactic; + the same bits), also for 33 to 64 rows, and the pre-attention attn_res in front of it + releases it at kernel entry, so its CTAs stream the QKVG weights while the attn_res and the + latent-up GEMM before it still run. The L2 prefetch of the next KDA layer's weights + (:meth:`_issue_l2_prefetch`) is skipped where the next layer stages them.""" + takes = getattr(self.qkvg_proj, "takes_staged_mxfp8_input", None) + return takes is not None and takes(num_rows) + + # ------------------------------------------------------------------ + # Forward. + # ------------------------------------------------------------------ + + def forward( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """Select the eager path or breakable-CUDA-graph path. + + Returns the o_proj output (all-reduced under head TP), or with the Lamport push + (``_project_output``) a ``LamportPending`` for the decoder layer's + A_mid. + """ + if self.register_to_config and is_in_breakable_cuda_graph(): + core = hidden_states.new_empty( + (hidden_states.shape[0], self.num_heads, self.head_dim), + dtype=torch.bfloat16, + ) + maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) + self.next_input_mxfp8 = None + return self._project_output(core) + + core = self._forward_impl(hidden_states, attn_metadata) + # The MXFP8 rows of the decoder layer serve this forward only (consumed by the verify + # projection, if it ran). + self.next_input_mxfp8 = None + # The push serves the one o_proj call of a generation-only batch. + return self._project_output(core, allow_ar_push=attn_metadata.num_contexts == 0) + + def _forward_impl( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + output: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, _KdaPrequantizedCore]: + """Split the batch into its prefill, decode and verify rows and run their chains. + + ``output`` is the BCG post-o_norm, pre-o_proj core buffer. When it + is supplied, subpaths fill and return its corresponding slices; + otherwise this method returns an allocated core tensor. A + verify-only batch without ``output`` may instead return the fused + verify epilogue's MXFP8 core (see ``forward_verify_fused``). + """ + mamba_metadata = attn_metadata.mamba_metadata + num_prefills = attn_metadata.num_contexts + num_ctx_tokens = attn_metadata.num_ctx_tokens + num_tokens = attn_metadata.num_tokens + batch_size = attn_metadata.seq_lens.shape[0] + state_indices = mamba_metadata.state_indices[:batch_size] + cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] + num_decodes = batch_size - num_prefills + + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + conv_pool = layer_cache.conv # [slots, 3D, W - 1] bf16 + ssm_pool = layer_cache.temporal # [slots, H, V, K] fp32 or bf16 + generation_state_indices = getattr(mamba_metadata, "generation_state_indices", None) + if generation_state_indices is None: + generation_state_indices = state_indices[num_prefills:] + + cores: List[torch.Tensor] = [] + if num_prefills > 0: + snapshot = getattr(mamba_metadata, "kda_snapshot", None) + inforward_stage = None + if snapshot is not None and snapshot.stage_rows is not None: + inforward_stage = attn_metadata.kv_cache_manager.inforward_snapshot_stage_views( + self.layer_idx + ) + prefill_core = self.forward_prefill( + hidden_states[:num_ctx_tokens], + cu_seqlens, + mamba_metadata, + num_prefills, + conv_pool, + ssm_pool, + state_indices[:num_prefills], + layer_cache, + output=output[:num_ctx_tokens] if output is not None else None, + inforward_stage=inforward_stage, + ) + if output is None: + cores.append(prefill_core) + if num_decodes > 0: + decode_rows = num_tokens - num_ctx_tokens + if decode_rows == num_decodes: + decode_core = self.forward_decode( + hidden_states[num_ctx_tokens:num_tokens], + conv_pool, + ssm_pool, + generation_state_indices, + layer_cache, + output=output[num_ctx_tokens:num_tokens] if output is not None else None, + ) + if output is None: + cores.append(decode_core) + else: + # Speculative verification: each generation request carries + # 1 + draft_len tokens (drafts are padded to the static max, + # so T is uniform). + assert decode_rows % num_decodes == 0, ( + f"ragged generation batch: {decode_rows} tokens for {num_decodes} requests" + ) + verify_core = self.forward_verify( + hidden_states[num_ctx_tokens:num_tokens], + decode_rows // num_decodes, + layer_cache, + ssm_pool, + generation_state_indices, + output=(output[num_ctx_tokens:num_tokens] if output is not None else None), + # The MXFP8 core cannot be concatenated with a prefill + # core or stored into the BCG buffer. + allow_prequantized_output=output is None and num_prefills == 0, + context_only=getattr(attn_metadata.kv_cache_manager, "kda_context_only", False) + is True, + ) + if output is None: + cores.append(verify_core) + if output is not None: + return output + return cores[0] if len(cores) == 1 else torch.cat(cores, dim=0) + + def _project_output( + self, + core: Union[torch.Tensor, _KdaPrequantizedCore], + allow_ar_push: bool = False, + ) -> Union[torch.Tensor, LamportPending]: + """o_proj on the post-o_norm core, then the head-TP all-reduce of its partials. + + ``allow_ar_push`` (generation-only batches, outside breakable CUDA graphs): o_proj stores + its partial straight into the all-reduce's Lamport buffers where the low-M GEMM serves it + (``o_proj.next_ar_push``); the decoder layer's ``ar_finish_deferred`` then makes this + return a ``LamportPending``. + """ + num_rows = core.activation.shape[0] if isinstance(core, tuple) else core.shape[0] + ar_push = None + if allow_ar_push and self._o_allreduce is not None and hasattr(self.o_proj, "next_ar_push"): + ar_push = lamport_push_target( + self._o_allreduce, + num_rows, + self.hidden_size, + deferred=self.ar_finish_deferred, + ) + self.o_proj.next_ar_push = ar_push + + if isinstance(core, _KdaPrequantizedCore): + out = self.o_proj.forward_prequantized(core.activation, core.scale) + else: + out = self.o_proj(core.reshape(-1, self.proj_size)) + + if ar_push is not None: + self.o_proj.next_ar_push = None + if ar_push is not None and ar_push.pushed: + # The decoder layer sets ar_finish_deferred where A_mid finishes the all-reduce: then + # this returns a LamportPending. + out_shape = out.shape + out = finish_lamport_push( + self._o_allreduce, + ar_push, + deferred=self.ar_finish_deferred, + ) + if not isinstance(out, LamportPending): + out = out.view(out_shape) + elif ar_push is not None and ar_push.deferred: + # o_proj ran without the push epilogue (e.g. the MXFP8 GEMM on the prequantized + # verify output); a copy kernel pushes it and A_mid finishes. + lamport_push_tensor(self._o_allreduce, ar_push, out.reshape(ar_push.output.shape)) + out = finish_lamport_push(self._o_allreduce, ar_push, deferred=True) + elif self._o_allreduce is not None: + # Head-sharded TP: every rank ran its head shard on the same + # local batch; sum the row-sharded o_proj partials. + out = self._o_allreduce(out) + return out + + def _sync_kda_replay_conv_window(self, layer_cache, slot_indices, conv_pool) -> None: + """Seed the fused verify kernel's replay conv caches from the live committed conv pool. + + Their committed window (columns ``[0, W-1)``) must hold the live pool's ``W - 1`` raw + inputs whenever prefill or plain decode advances it. No-op without replay caches (no + speculative decoding, or a context-only server). + """ + if not bool(getattr(layer_cache, "has_kda_replay_caches", False)): + return + layer_cache.commit_conv_window(slot_indices, conv_pool) + + def _project_gate_inputs(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """``(beta, forget_gate)``: ``F.linear`` on the fused ``[f_a | b]`` weight, then + ``F.linear`` of its ``f_a`` columns on ``f_b``. ``beta`` and the forget gate are the raw + logits (the kernels apply the sigmoid / gate).""" + hd, H = self.head_dim, self.num_heads + bfa = F.linear(x, self._bfa_proj_weight) + return bfa[..., hd : hd + H], F.linear(bfa[..., :hd], self.f_b_proj.weight) + + def _project_packed_conv_input( + self, x: torch.Tensor, token_major: bool + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(conv_input, onorm_g)`` of the prefill rows ``x`` ``[1, T, hidden]``: the QKVG GEMM, + its ``[q | k | v]`` columns as a ``[T, 3D]`` view (``token_major``, for + ``fused_kda_conv_prefill``) or repacked channel-major ``[3D, T]`` + (``extract_transpose_prefill_slice``, for ``causal_conv1d_fwd``), and its output-gate + columns.""" + d = self.proj_size + qkvg = self.qkvg_proj(x) + onorm_g = qkvg[..., 3 * d : 4 * d] + qkvg_2d = qkvg.squeeze(0) + if token_major: + return qkvg_2d[:, : 3 * d], onorm_g + return extract_transpose_prefill_slice(qkvg_2d, qkvg_2d.shape[0], 0, 3 * d), onorm_g + + def forward_prefill( + self, + x2d, + cu_seqlens, + mamba_metadata, + num_prefills, + conv_pool, + ssm_pool, + slot_indices, + layer_cache=None, + output: Optional[torch.Tensor] = None, + inforward_stage: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + """Prefill ``num_prefills`` packed sequences and update their conv / recurrent rows. + + The convolution is ``fused_kda_conv_prefill`` (features.kda_fused_conv_prefill) or + ``causal_conv1d_fwd`` + ``fused_kda_post_conv``; Q/K leave it L2-normalized for + ``kda_prefill``. A single unaligned sequence is chunk-padded by the + post-convolution kernel (features.kda_prepad_prefill) or by the op. A BF16 state pool + (and every in-forward snapshot plan) is staged through an fp32 copy; a ``kda_snapshot`` + plan also stores mid-chunk states and convolution windows. + """ + if x2d.shape[0] == 0: + reset_recurrent_state_rows( + ssm_pool, slot_indices, mamba_metadata.has_initial_states[:num_prefills], conv_pool + ) + self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_pool) + return self._store_core(x2d.new_empty((0, self.num_heads, self.head_dim)), output) + self._require_finalized() + x = x2d.unsqueeze(0) # [1, T, hidden] + # Metadata owns stable int32 indices; never convert per layer. + slot_indices = mamba_metadata.state_indices[:num_prefills] + + # Mid-chunk snapshots: the recurrent state and convolution window of a sequence after its + # first tokens of this step. The prefill op writes the state (snapshot_rows / + # snapshot_chunks). With ``inforward_stage`` (the in-forward snapshots of the cache + # manager, ``snapshot.stage_rows`` set) the live slot keeps the state at the snapshot + # point and the staging row receives the chunk-end state; otherwise the snapshot goes to + # a pool slot of its own. + snapshot = getattr(mamba_metadata, "kda_snapshot", None) + if (snapshot is not None and snapshot.stage_rows is not None) != ( + inforward_stage is not None + ): + raise RuntimeError( + "Kimi K3 KDA in-forward snapshots need both the plan's staging rows and the " + "cache manager's staging views" + ) + + fused_conv = _KDA_FUSED_CONV_PREFILL + packed_conv, onorm_g = self._project_packed_conv_input(x, token_major=fused_conv) + snapshot_windows = None + if snapshot is not None: + # The window is read before the convolution below, which may overwrite its input in + # place. A live slot's window still feeds that convolution; it is swapped after it. + snapshot_windows = _snapshot_conv_windows( + packed_conv, conv_pool, snapshot, token_major=fused_conv + ) + if inforward_stage is None: + conv_pool.index_copy_(0, snapshot.slots, snapshot_windows) + + # Initial states: present for continuation chunks (chunked prefill) + # and for prefix-cache hits (block reuse), where the previous + # conv/recurrent state was onboarded into this request's slot. + has_init = mamba_metadata.has_initial_states[:num_prefills] + prepared_state = self._prepare_prefill_state( + cu_seqlens, + mamba_metadata, + num_prefills, + ssm_pool, + slot_indices, + snapshot=snapshot, + ) + if not fused_conv: + # The convolution input is a contiguous channel-major [3D, T] copy: the op writes + # its output over it. + causal_conv1d_fwd( + packed_conv, + self._packed_conv_weight, + None, + conv_pool, + mamba_metadata.query_start_loc[: num_prefills + 1], + mamba_metadata.state_indices[:num_prefills], + has_init, + True, + PAD_SLOT_ID, + ) + + # The prefill op needs a single unaligned sequence zero-padded to a chunk multiple. + # features.kda_prepad_prefill produces Q/K/V/beta at that length directly instead of + # padding copies; g needs no pad because the op rebuilds g's sentinel tail from the + # real prefix. + num_tokens = x2d.shape[0] + padded_tokens = None + single_sequence_length = getattr(mamba_metadata, "kda_single_sequence_length", None) + if ( + _KDA_PREPAD_PREFILL + and num_prefills == 1 + and single_sequence_length is not None + and single_sequence_length == num_tokens + and num_tokens % _PREFILL_CHUNK != 0 + ): + padded_tokens = num_tokens + _PREFILL_CHUNK - num_tokens % _PREFILL_CHUNK + + beta, g = self._project_gate_inputs(x) + g = g.unflatten(-1, (self.num_heads, self.head_dim)) + + # Imported here: these entries import this package's kernels, and the package imports + # this module. + from .ops.ssm.fused_kda_conv_prefill import fused_kda_conv_prefill + from .ops.ssm.fused_kda_post_conv import fused_kda_post_conv + + if fused_conv: + prepad_beta = padded_tokens is not None and beta.stride(-1) == 1 + if not prepad_beta: + beta = beta.float() + outputs = fused_kda_conv_prefill( + packed_conv, + self._fused_conv_weight_t.t(), + conv_pool, + query_start_loc=mamba_metadata.query_start_loc[: num_prefills + 1], + cache_indices=mamba_metadata.state_indices[:num_prefills], + has_initial_state=has_init, + num_heads=self.num_heads, + head_dim=self.head_dim, + padded_tokens=padded_tokens if prepad_beta else None, + beta=beta if prepad_beta else None, + normalize_qk=True, + ) + if prepad_beta: + q, k, v, beta = outputs + else: + q, k, v = outputs + elif padded_tokens is not None and beta.stride(-1) == 1 and packed_conv.is_contiguous(): + q, k, v, beta = fused_kda_post_conv( + packed_conv, + num_heads=self.num_heads, + head_dim=self.head_dim, + padded_tokens=padded_tokens, + beta=beta, + ) + else: + beta = beta.float() + q, k, v = fused_kda_post_conv( + packed_conv, + num_heads=self.num_heads, + head_dim=self.head_dim, + ) + + o = self._prefill_recurrence( + q, + k, + v, + g, + beta, + cu_seqlens, + mamba_metadata, + num_prefills, + ssm_pool, + slot_indices, + prepared_state, + snapshot=snapshot, + inforward_stage_ssm=None if inforward_stage is None else inforward_stage[0], + ) + # Fused-verify replay caches: seed the committed conv window so the + # first verify round convolves the correct history (pending drafts + # are zero for a fresh request, so the tail columns are unused). + self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_pool) + if inforward_stage is not None: + _swap_inforward_conv_windows(conv_pool, inforward_stage[1], snapshot, snapshot_windows) + + core = self._output_gate(o, onorm_g) + return self._store_core(core, output) + + def _prepare_prefill_state( + self, + cu_seqlens: torch.Tensor, + mamba_metadata, + num_prefills: int, + ssm_pool: torch.Tensor, + slot_indices: torch.Tensor, + snapshot: Optional["KDASnapshotPlan"] = None, + ) -> _KDAPrefillState: + """The state rows the prefill op updates in place, zeroed where a sequence starts fresh. + + The op takes fp32 V-first pool rows. A BF16 pool is staged through a dense fp32 copy of + the addressed rows (plus the snapshot rows of a ``snapshot`` plan), as is an fp32 pool + under an in-forward snapshot plan (whose snapshot rows are the live slots, which the + kernel's chunk-end write would overwrite); :meth:`_prefill_recurrence` writes the copy + back. + """ + has_init = mamba_metadata.has_initial_states[:num_prefills] + staged_state = None + kernel_pool, kernel_indices = ssm_pool, slot_indices + snapshot_rows = None if snapshot is None else snapshot.pool_rows + inforward = snapshot is not None and snapshot.stage_rows is not None + if ssm_pool.dtype != torch.float32 or inforward: + if snapshot is None: + staged_state = _stage_state_rows(ssm_pool, slot_indices) + else: + staged_state = _stage_state_rows_with_snapshot_rows( + ssm_pool, slot_indices, snapshot.num_snapshots + ) + snapshot_rows = snapshot.staged_rows + kernel_pool = staged_state + kernel_indices = mamba_metadata._arange_buffer[:num_prefills] + if not self._dispatch.can_use_indexed_prefill( + state_pool=kernel_pool, + state_indices=kernel_indices, + has_initial_states=has_init, + cu_seqlens=cu_seqlens, + num_sequences=num_prefills, + ): + raise RuntimeError( + "Kimi K3 KDA prefill: the kda_prefill entry does not take this state pool / " + "metadata layout (fp32 [slots, H, V, K] rows, 16-byte aligned int32 / int64 " + "indices, bool has_initial_states)" + ) + # The convolution clears fresh convolution rows itself. + reset_recurrent_state_rows(kernel_pool, kernel_indices, has_init) + return _KDAPrefillState( + kernel_pool, kernel_indices, staged_state, snapshot_rows=snapshot_rows + ) + + def _prefill_recurrence( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + mamba_metadata, + num_prefills: int, + ssm_pool: torch.Tensor, + slot_indices: torch.Tensor, + state: _KDAPrefillState, + *, + snapshot: Optional["KDASnapshotPlan"] = None, + inforward_stage_ssm: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``kda_prefill`` on the token-major L2-normalized Q/K, V, the raw gate and + beta logits, updating ``state``'s fp32 V-first ``[slots, H, D, D]`` rows in place, then + the write-back of a staged copy (and the snapshot moves).""" + o = self._dispatch.prefill_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self._A_log_f32, + dt_bias=self._dt_bias_f32, + scale=self.head_k_dim**-0.5, + lower_bound=self.gate_lower_bound, + cu_seqlens=cu_seqlens, + chunk_indices=getattr(mamba_metadata, "kda_chunk_indices", None), + chunk_size=_PREFILL_CHUNK, + state_pool=state.kernel_pool, + state_indices=state.kernel_indices, + varlen_is_aligned=getattr(mamba_metadata, "kda_varlen_is_aligned", None), + single_sequence_length=getattr(mamba_metadata, "kda_single_sequence_length", None), + snapshot_rows=state.snapshot_rows if snapshot is not None else None, + snapshot_chunks=None if snapshot is None else snapshot.chunks, + ) + + # The convolution persisted the live convolution pool in place. + staged_state = state.staged_state + if staged_state is not None: + if snapshot is None: + _writeback_state_rows(ssm_pool, slot_indices, staged_state) + else: + _writeback_state_rows(ssm_pool, slot_indices, staged_state[:num_prefills]) + if snapshot.stage_rows is not None: + # In-forward: the chunk-end states move to the staging rows; the live + # slots (``snapshot.slots``) take the states at the snapshot points. + inforward_stage_ssm.index_copy_( + 0, snapshot.stage_rows, ssm_pool.index_select(0, snapshot.slots) + ) + _writeback_state_rows(ssm_pool, snapshot.slots, staged_state[num_prefills:]) + elif snapshot is not None and snapshot.stage_rows is not None: + raise RuntimeError("Kimi K3 KDA in-forward snapshots need a staged state copy") + return o + + def forward_decode( + self, + x2d, + conv_pool, + ssm_pool, + slot_indices, + layer_cache=None, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """One token per request: the QKVG GEMM (next to ``[f_a | b]`` -> ``f_b`` on the + projection aux stream up to 128 rows in CUDA graphs), then the fused decode kernel, which + convolves, gates, runs the recurrence and the gated RMSNorm, and updates the conv and + fp32 state pools by slot. A BF16 state pool is staged through dense copies of the + addressed rows (fp32 state, conv window as is) and written back (state rounded). + """ + self._require_finalized() + d = self.proj_size + hd = self.head_dim + H = self.num_heads + B = x2d.shape[0] + + staged_state = staged_conv = None + slots_long = None + kernel_state, kernel_state_indices, conv_src = ssm_pool, slot_indices, conv_pool + if ssm_pool.dtype != torch.float32: + slots_long = slot_indices.long() + staged_state = _stage_state_rows(ssm_pool, slots_long) + staged_conv = conv_pool.index_select(0, slots_long) + kernel_state, kernel_state_indices, conv_src = staged_state, None, staged_conv + + # The decode kernel writes its core in place: the BCG core buffer, or outside breakable + # CUDA graphs a persistent buffer whose pointer stays stable across graph captures. + if output is not None: + kda_out = output.view(B, 1, H, hd) + else: + if self._o_dense is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Kimi K3 KDA decode: its output buffer is allocated by the eager forward " + "that precedes CUDA graph capture" + ) + self._o_dense = torch.empty( + max(conv_pool.shape[0], B), + 1, + H, + hd, + dtype=torch.bfloat16, + device=x2d.device, + ) + else: + assert self._o_dense.shape[0] >= B, ( + f"KDA decode output buffer holds {self._o_dense.shape[0]} rows " + f"but the decode batch is {B}; reallocating would corrupt " + f"previously captured CUDA graphs" + ) + kda_out = self._o_dense[:B] + + projection_aux_stream = ( + self._projection_aux_stream if B <= _KDA_BFA_MULTISTREAM_MAX_ROWS else None + ) + qkvg, (beta, g) = fork_join( + lambda: self.qkvg_proj(x2d), + lambda: self._project_gate_inputs(x2d), + self._projection_fork_event, + self._projection_join_event, + projection_aux_stream, + ) + x_qkvg = qkvg[:, : 3 * d] + onorm_g = qkvg[:, 3 * d : 4 * d] + + # Section views retain the live pool's slot stride, including V2 + # manager padding. The kernel uses ssm_state_indices for both pools. + o = self._dispatch.decode_kda( + x_q=x_qkvg[:, :d].unflatten(-1, (H, hd)).unsqueeze(0), + x_k=x_qkvg[:, d : 2 * d].unflatten(-1, (H, hd)).unsqueeze(0), + x_v=x_qkvg[:, 2 * d : 3 * d].unflatten(-1, (H, hd)).unsqueeze(0), + w_q_t=self._w_q_t, + w_k_t=self._w_k_t, + w_v_t=self._w_v_t, + bias_q=None, + bias_k=None, + bias_v=None, + cs_q=conv_src[:, :d], + cs_k=conv_src[:, d : 2 * d], + cs_v=conv_src[:, 2 * d :], + A_log=self._A_log_f32, + g=g.unflatten(-1, (H, hd)).unsqueeze(0), + dt_bias=self._dt_bias_f32, + beta=beta.unsqueeze(0), + state=kernel_state, + onorm_g=onorm_g.unflatten(-1, (H, hd)).unsqueeze(0), + onorm_weight=self._onorm_w_f32, + out=kda_out, + ssm_state_indices=kernel_state_indices, + scale=hd**-0.5, + onorm_eps=self.o_norm.eps, + lower_bound=self.gate_lower_bound, + use_beta_sigmoid_in_kernel=True, + verbose=False, + update_conv_cache=True, + ) + if staged_state is not None: + conv_pool.index_copy_(0, slots_long, staged_conv) + _writeback_state_rows(ssm_pool, slots_long, staged_state) + # Fused-verify replay caches (spec decoding only): keep the + # committed conv window in sync with the plain-decode advance. + self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_pool) + + # The kernel applied the output gate / RMSNorm and returns the supplied buffer. + return o.view(B, H, hd) + + def forward_verify( + self, + x2d, + num_steps, + layer_cache, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + allow_prequantized_output: bool = False, + context_only: bool = False, + ) -> Union[torch.Tensor, _KdaPrequantizedCore]: + """Speculative verification: advance each request ``num_steps`` tokens (1 golden + + ``num_steps - 1`` padded drafts) with the fused verify kernel on the cache manager's + replay caches (:meth:`forward_verify_fused`). + + On a context-only server the rows are dummies and take + ``forward_verify_context_only`` instead (the manager marks such a server with + ``kda_context_only``, passed in as ``context_only``). + """ + if bool(getattr(layer_cache, "has_kda_replay_caches", False)): + return self.forward_verify_fused( + x2d, + num_steps, + layer_cache, + ssm_pool, + slot_indices, + output=output, + allow_prequantized_output=allow_prequantized_output, + ) + if context_only or getattr(layer_cache, "context_only", False): + return self.forward_verify_context_only(x2d, output) + raise RuntimeError( + "Kimi K3 KDA verification runs the fused verify kernel on the cache manager's KDA " + "replay caches, which this layer cache does not have" + ) + + def forward_verify_context_only( + self, x2d: torch.Tensor, output: Optional[torch.Tensor] = None + ) -> torch.Tensor: + """Verification rows on a context-only server. + + Such a server (disaggregated context server) runs no real generation + request; its generation rows are dummies (warmup, attention-DP + padding) whose outputs are discarded, and its cache manager keeps no + replay caches for them. The recurrence is skipped and the live pools + are left untouched; the output projection (``_project_output``) still + runs on a zero attention core, so the layer's projection and collectives + match those of a real verification step. + """ + core = x2d.new_zeros(x2d.shape[0], self.num_heads, self.head_dim) + return self._store_core(core, output) + + def _project_verify_inputs(self, x: torch.Tensor, num_rows: int, gate_from_fa: bool) -> tuple: + """The verify projections: the QKVG GEMM on the calling stream, the gate projection on + the projection aux stream. + + Returns ``(q, k, v, forget_gate, beta, onorm_g, f_a)``, the q / k / v / output-gate + sections as strided views of the QKVG output (the verify kernel reads them in place) and + beta as one of the ``[f_a | b]`` output. ``f_a`` is the BF16 ``f_a`` view where + ``gate_from_fa`` and the f_a-mode fab kernel served these rows (``forget_gate`` is then + None and the verify kernel computes it), else None. + + The QKVG GEMM takes the decoder layer's MXFP8 rows (``next_input_mxfp8``, + features.attn_res_mxfp8) where they were handed over, staging its weights. The gate is + ``k3_kda_fab_fa`` (f_a mode), ``k3_kda_fab_proj`` (beta and the forget gate) where they + accept the shapes, else ``F.linear`` ``[f_a | b]`` -> ``f_b``. + """ + fused_qkvg = self.qkvg_proj + mxfp8_input, self.next_input_mxfp8 = self.next_input_mxfp8, None + + def _project_qkvg() -> torch.Tensor: + if mxfp8_input is not None: + return fused_qkvg.forward_mxfp8(x, *mxfp8_input, stage_weights=True) + return fused_qkvg(x) + + bfa_weight = self._bfa_proj_weight + fb_weight = self.f_b_proj.weight + fab_fa = gate_from_fa and _use_fab_fa(x, bfa_weight, fb_weight, self.num_heads) + fused_fab = not fab_fa and _use_fused_fab(x, bfa_weight, fb_weight, self.num_heads) + + def _project_gates() -> tuple: + if fab_fa: + from .ops.ssm.k3_kda_fab_fa import k3_kda_fab_fa + + # (beta, f_a): the verify kernel derives the forget gate from f_a. + # W_fb is warmed in L2 for the verify kernel's per-head load. + bfa = k3_kda_fab_fa( + x.reshape(-1, x.shape[-1]), bfa_weight, self.num_heads, fb_weight + ).view(*x.shape[:-1], bfa_weight.shape[0]) + return bfa[..., self.head_dim : self.head_dim + self.num_heads], bfa[ + ..., : self.head_dim + ] + if fused_fab: + from .ops.ssm.k3_kda_fab_proj import k3_kda_fab_proj + + beta, forget_gate = k3_kda_fab_proj( + x.reshape(-1, x.shape[-1]), bfa_weight, fb_weight, self.num_heads + ) + return ( + beta.view(*x.shape[:-1], self.num_heads), + forget_gate.view(*x.shape[:-1], fb_weight.shape[0]), + ) + return self._project_gate_inputs(x) + + if fab_fa or fused_fab: + max_rows = _KDA_FAB_FUSED_MAX_ROWS + else: + # Large batches: the unfused chain still overlaps the (all-SM) QKVG GEMM better + # than it runs after it. + max_rows = num_rows + projection_aux_stream = self._projection_aux_stream if 0 < num_rows <= max_rows else None + qkvg, (beta, gate) = fork_join( + _project_qkvg, + _project_gates, + self._projection_fork_event, + self._projection_join_event, + projection_aux_stream, + ) + + d = self.proj_size + q_proj, k_proj, v_proj = qkvg[..., : 3 * d].split(d, dim=-1) + onorm_g = qkvg[..., 3 * d : 4 * d] + if fab_fa: + return q_proj, k_proj, v_proj, None, beta, onorm_g, gate + return q_proj, k_proj, v_proj, gate, beta, onorm_g, None + + def _next_kda_layer(self) -> Optional["K3KDALinearAttention"]: + """The KDA mixer of the next KDA layer of this model, if any.""" + later = [idx for idx in self._kda_group.keys() if idx > self.layer_idx] + return self._kda_group.get(min(later)) if later else None + + def l2_prefetch_ranges(self) -> List[Tuple[int, int]]: + """``(address, bytes)`` of the weights the fused verify path streams from HBM: the QKVG + and o_proj weights and scales, ``[f_a | b]`` and ``f_b``.""" + if self._l2_prefetch_ranges_cache is None: + from .kernels.k3_l2_prefetch import MAX_RANGES, tensor_range + + tensors = [] + for module in (self.qkvg_proj, self.o_proj): + if module is not None: + tensors += [ + getattr(module, "weight", None), + getattr(module, "weight_scale", None), + ] + tensors += [self._bfa_proj_weight, self.f_b_proj.weight] + ranges = [] + for tensor in tensors: + if isinstance(tensor, torch.Tensor) and tensor.is_cuda and tensor.numel(): + try: + ranges.append(tensor_range(tensor)) + except ValueError: + continue + self._l2_prefetch_ranges_cache = ranges[:MAX_RANGES] + return self._l2_prefetch_ranges_cache + + def _issue_l2_prefetch(self, num_rows: int) -> None: + """Prefetch the next KDA layer's weights into L2 from the projection aux stream, after the + projections of this layer (see ``_KDA_L2_PREFETCH``).""" + if ( + not _KDA_L2_PREFETCH + or num_rows > _KDA_L2_PREFETCH_MAX_ROWS + or self._projection_aux_stream is None + or not do_multi_stream() + ): + return + nxt = self._next_kda_layer() + if nxt is not None and nxt.qkvg_stages_weights(num_rows): + # The next layer streams its QKVG weights on chip itself. + return + ranges = nxt.l2_prefetch_ranges() if nxt is not None else [] + if not ranges: + return + from .ops.ssm.k3_l2_prefetch import k3_l2_prefetch + + # Recorded after the [f_a | b] join: the prefetch starts once QKVG and [f_a | b] are done + # (both read their weights from HBM), i.e. under the recurrence. The aux stream is joined + # back by its next user (see _KDA_L2_PREFETCH). + self._l2_prefetch_start.record() + with torch.cuda.stream(self._projection_aux_stream): + self._l2_prefetch_start.wait() + k3_l2_prefetch(ranges) + + def forward_verify_fused( + self, + x2d, + num_steps, + layer_cache, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + allow_prequantized_output: bool = False, + ) -> Union[torch.Tensor, _KdaPrequantizedCore]: + """Fused multi-token verify (``mtp_ops.py``: the v1 / v2 / two-CTA kernels). + + The runtime packs ``num_steps`` new rows contiguously per request, so the kernel derives + the row offsets from the request index and reads each request's pending accepted count + from the slot-indexed replay pool (``packed_token_layout``). One launch replays the + previous round's accepted drafts from the per-slot replay caches, then processes the new + tokens, committing the recurrent state and conv windows in place after the golden token + and caching the new drafts, and applies the gated RMSNorm in its epilogue: + + * with ``fa`` / ``w_fb``: the forget gate ``bf16(f_a @ W_fb^T)`` computed in the + kernel (BF16 core); + * ``quantize_output``: where ``allow_prequantized_output`` and ``o_proj`` takes a + prequantized activation for this row count (SM107), the core quantized to MXFP8 for + ``o_proj.forward_prequantized``; + * otherwise the BF16 core, a programmatic dependent of the QKVG GEMM. + """ + self._require_finalized() + num_generations = x2d.shape[0] // num_steps + num_spec = num_steps - 1 + H = self.num_heads + K = self.head_k_dim + x = x2d.view(num_generations, num_steps, -1) # [B, T, hidden] + T_total = num_generations * num_steps + + # The MXFP8 core cannot be concatenated with a prefill core or stored into the BCG buffer + # (``allow_prequantized_output``). + mxfp8_output = ( + allow_prequantized_output + and output is None + and _o_proj_takes_mxfp8(self.o_proj, T_total) + and callable(getattr(self.o_proj, "forward_prequantized", None)) + ) + q_proj, k_proj, v_proj, forget_gate, beta_proj, onorm_g, f_a = self._project_verify_inputs( + x, T_total, gate_from_fa=not mxfp8_output + ) + self._issue_l2_prefetch(T_total) + x_q = q_proj.view(1, T_total, H, K) + x_k = k_proj.view(1, T_total, H, K) + x_v = v_proj.view(1, T_total, H, self.head_dim) + # Raw gate / beta: the kernel applies dt_bias, A_log, the + # lower-bound sigmoid gate, and the beta sigmoid itself. + beta = beta_proj.view(1, T_total, H) + w_q, w_k, w_v = self._mtp_conv_weights + onorm_g_view = onorm_g.view(1, T_total, H, self.head_dim) + + if f_a is not None: + out = self._dispatch.mtp_verify_fused_fb( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=layer_cache.kda_conv_q, + cs_k=layer_cache.kda_conv_k, + cs_v=layer_cache.kda_conv_v, + fa=f_a.view(1, T_total, K), + w_fb=self.f_b_proj.weight.detach(), + beta=beta, + A_log=self.A_log.detach(), + dt_bias=self.dt_bias.detach(), + recurrent_state=ssm_pool, + k_cache=layer_cache.kda_k_cache, + g_cache=layer_cache.kda_g_cache, + v_cache=layer_cache.kda_v_cache, + beta_cache=layer_cache.kda_beta_cache, + ssm_state_indices=slot_indices, + num_spec=num_spec, + num_accepted_tokens=layer_cache.prev_num_accepted_tokens, + lower_bound=self.gate_lower_bound, + scale=self.head_k_dim**-0.5, + onorm_g=onorm_g_view, + onorm_weight=self._onorm_w_f32, + onorm_eps=self.o_norm.eps, + use_pdl=True, + ) + return self._store_core(out, output) + + out = self._dispatch.mtp_verify( + mxfp8_output=mxfp8_output, + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=layer_cache.kda_conv_q, + cs_k=layer_cache.kda_conv_k, + cs_v=layer_cache.kda_conv_v, + g=forget_gate.view(1, T_total, H, K), + beta=beta, + # .detach(): the CuTe DSL DLPack bridge rejects grad-tracking tensors. + A_log=self.A_log.detach(), + dt_bias=self.dt_bias.detach(), + recurrent_state=ssm_pool, + k_cache=layer_cache.kda_k_cache, + g_cache=layer_cache.kda_g_cache, + v_cache=layer_cache.kda_v_cache, + beta_cache=layer_cache.kda_beta_cache, + ssm_state_indices=slot_indices, + cu_seqlens=None, + num_spec=num_spec, + num_accepted_tokens=layer_cache.prev_num_accepted_tokens, + packed_token_layout=True, + lower_bound=self.gate_lower_bound, + scale=self.head_k_dim**-0.5, + onorm_g=onorm_g_view, + onorm_weight=self._onorm_w_f32, + onorm_eps=self.o_norm.eps, + fuse_output_norm=True, + # A BF16-output verify launches as a programmatic dependent of the QKVG projection: + # its state, replay and weight loads run before griddepcontrol.wait. + use_pdl=not mxfp8_output, + ) + if mxfp8_output: + if not isinstance(out, tuple): + raise RuntimeError("MXFP8-output MTP KDA must return activation and scale") + return _KdaPrequantizedCore(*out) + return self._store_core(out, output) + + def _output_gate(self, o: torch.Tensor, onorm_g: torch.Tensor) -> torch.Tensor: + """The gated output RMSNorm ``rmsnorm(o) * sigmoid(g)`` (``rms_norm_gated_token_major``) + of the prefill core.""" + return rms_norm_gated_token_major( + o.reshape(-1, self.head_dim), + onorm_g.reshape(-1, self.num_heads, self.head_dim), + self.o_norm.weight, + self.o_norm.eps, + gate_activation="sigmoid", + ) + + def _store_core(self, core: torch.Tensor, output: Optional[torch.Tensor]) -> torch.Tensor: + """Return the normalized core ``[tokens, H, D]``, copied into the BCG buffer ``output`` + where one is given.""" + core = core.reshape(-1, self.num_heads, self.head_dim) + if output is None: + return core + output.copy_(core) + return output + + +# Granularity of the branch's mid-chunk snapshot offsets (the CuTe DSL prefill kernel's chunk). +KDA_PREFILL_CHUNK_SIZE = 64 + + +@dataclass +class KDASnapshotPlan: + """Mid-chunk recurrent-state snapshots requested for one context batch. + + Snapshot ``j`` (ordered by context sequence) makes the KDA layers store + its sequence's recurrent state and convolution window after the + sequence's first ``offset_j`` tokens of this step (a multiple of + ``KDA_PREFILL_CHUNK_SIZE``, strictly inside the step) into pool slot + ``slots[j]``. Per-sequence vectors hold ``-1`` for sequences without a + snapshot. All tensors live on the device. + """ + + num_snapshots: int + # int64 [n]: pool slot that receives each snapshot. + slots: torch.Tensor + # int64 [n]: packed token index of the first token after each snapshot point. + token_indices: torch.Tensor + # int32 [num_contexts]: KDA chunk at whose start the snapshot is taken. + chunks: torch.Tensor + # int32 [num_contexts]: snapshot row in the pool itself (fp32 pools). + pool_rows: torch.Tensor + # int32 [num_contexts]: snapshot row in a staged copy holding the + # sequences' rows first and the snapshot rows after them (bf16 pools). + staged_rows: torch.Tensor + # In-forward snapshots: int64 [n] rows of the cache manager's staging tensors + # (``inforward_snapshot_stage_views``) that receive each sequence's chunk-end state; + # ``slots`` are then the sequences' live slots, which keep the state at the snapshot point. + # None: the snapshots go to pool slots of their own. + stage_rows: Optional[torch.Tensor] = None + + +def build_kda_snapshot_plan( + snapshots: List[Tuple[int, int, int]], + context_lengths: List[int], + device: torch.device, +) -> Optional[KDASnapshotPlan]: + """Build the device plan for ``(context index, slot, offset)`` snapshots.""" + if not snapshots: + return None + num_contexts = len(context_lengths) + starts = [0] * num_contexts + for i in range(1, num_contexts): + starts[i] = starts[i - 1] + context_lengths[i - 1] + chunks = [-1] * num_contexts + pool_rows = [-1] * num_contexts + staged_rows = [-1] * num_contexts + slots, token_indices = [], [] + for j, (seq, slot, offset) in enumerate(sorted(snapshots)): + if not 0 <= seq < num_contexts or chunks[seq] != -1: + raise ValueError(f"Invalid KDA snapshot sequence index {seq}") + if offset % KDA_PREFILL_CHUNK_SIZE or not 0 < offset < context_lengths[seq]: + raise ValueError( + f"KDA snapshot offset {offset} must be a multiple of " + f"{KDA_PREFILL_CHUNK_SIZE} inside the {context_lengths[seq]}-token step" + ) + chunks[seq] = offset // KDA_PREFILL_CHUNK_SIZE + pool_rows[seq] = slot + staged_rows[seq] = num_contexts + j + slots.append(slot) + token_indices.append(starts[seq] + offset) + num_snapshots = len(slots) + pin = device.type == "cuda" and prefer_pinned() + snapshot_vectors = torch.tensor([*slots, *token_indices], dtype=torch.int64, pin_memory=pin).to( + device, non_blocking=True + ) + # The kernel reads the per-sequence vectors as 16-byte aligned int32 arrays. + stride = (num_contexts + 3) // 4 * 4 + padding = [-1] * (stride - num_contexts) + sequence_vectors = torch.tensor( + [*chunks, *padding, *pool_rows, *padding, *staged_rows, *padding], + dtype=torch.int32, + pin_memory=pin, + ).to(device, non_blocking=True) + return KDASnapshotPlan( + num_snapshots=num_snapshots, + slots=snapshot_vectors[:num_snapshots], + token_indices=snapshot_vectors[num_snapshots:], + chunks=sequence_vectors[:num_contexts], + pool_rows=sequence_vectors[stride : stride + num_contexts], + staged_rows=sequence_vectors[2 * stride : 2 * stride + num_contexts], + ) + + +@torch.compile(dynamic=True, fullgraph=True) +def _prepare_kda_chunk_indices( + sequence_lengths: torch.Tensor, + chunk_indices: torch.Tensor, +) -> torch.Tensor: + """Fill FLA-compatible chunk indices without reading CUDA counts on CPU.""" + chunk_counts = (sequence_lengths + KDA_PREFILL_CHUNK_SIZE - 1).div( + KDA_PREFILL_CHUNK_SIZE, rounding_mode="floor" + ) + num_chunks = chunk_indices.shape[0] + # Supplying output_size keeps repeat_interleave from synchronizing to read + # sum(chunk_counts) from the device. + sequence_indices = torch.repeat_interleave( + torch.arange(sequence_lengths.shape[0], dtype=torch.long, device=sequence_lengths.device), + chunk_counts, + output_size=num_chunks, + ) + chunk_starts = torch.cumsum(chunk_counts, dim=0) - chunk_counts + local_chunk_indices = ( + torch.arange(num_chunks, dtype=torch.long, device=sequence_lengths.device) + - chunk_starts[sequence_indices] + ) + chunk_indices[:, 0].copy_(sequence_indices) + chunk_indices[:, 1].copy_(local_chunk_indices) + return chunk_indices + + +class K3MambaMetadata(Mamba2Metadata): + """Reuse stable int32 slots/starts and int64 offsets from Mamba metadata.""" + + # The CuTe DSL prefill kernel keys its compile cache on chunk alignment (64); the warmup + # primes both sides of it. A 64-aligned length is 16-aligned too, so main's FlashInfer + # prefill (alignment 16; the fallback) gets both of its variants from the same two passes. + prefill_chunk_alignment = KDA_PREFILL_CHUNK_SIZE + + def __init__(self, max_batch_size: int, chunk_size: int, max_num_tokens: int) -> None: + super().__init__(max_batch_size, chunk_size, max_num_tokens) + self.kda_max_seqlen = 0 + # Align both rows to 16 bytes. One transfer refreshes the metadata shared + # by all layers; device storage stays fixed for CUDA graph replay. + width = ((max_batch_size + 1 + 3) // 4) * 4 + self._kda_schedule = torch.zeros((2, width), dtype=torch.int32, device="cuda") + self.kda_seq_order = self._kda_schedule[0, :0] + self.kda_cu_chunks = self._kda_schedule[1, :1] + + # KDA's CuTe verify kernel requires a 16-byte-aligned DLPack pointer. + # Keep the backing allocation stable so metadata preparation remains + # safe for CUDA graph capture and replay. + self._generation_state_indices = torch.empty( + max_batch_size, dtype=torch.int32, device="cuda" + ) + self.generation_state_indices: torch.Tensor | None = None + + # Rubin agentic branch: an int64 mirror of the state indices, refreshed once per step. + self._state_indices_long = torch.zeros(max_batch_size, dtype=torch.long, device="cuda") + self.state_indices_long = self._state_indices_long[:0] + # The CuTe DSL prefill's varlen chunk table ([num_chunks, 2]: sequence, chunk in it), + # filled on the device once per step instead of in each KDA layer. + self._kda_chunk_indices = torch.empty((0, 2), dtype=torch.long, device="cuda") + self.kda_chunk_indices: torch.Tensor | None = None + self.kda_varlen_is_aligned: bool | None = None + self.kda_single_sequence_length: int | None = None + self.kda_snapshot: KDASnapshotPlan | None = None + + def prepare(self, attn_metadata: AttentionMetadata) -> None: + context_lengths = attn_metadata.seq_lens[: attn_metadata.num_contexts].tolist() + assert all(length > 0 for length in context_lengths), ( + f"Expected positive context lengths, got {context_lengths}" + ) + + # Mamba2Metadata.prepare builds this step's in-forward plan (_prepare_inforward_split). + self.kda_snapshot = None + super().prepare(attn_metadata) + + self.generation_state_indices = None + batch_size = attn_metadata.seq_lens.shape[0] + if step_inputs.is_open() and self._state_indices_aliased_ptr is None: + # ``Mamba2Metadata.prepare`` wrote the int32 indices from ``state_indices_cpu`` (it + # did not alias a cache manager buffer): the pack uploads the same values, widened on + # the host, at the graph replay. + step_inputs.copy_( + self._state_indices_long[:batch_size], self.state_indices_cpu[:batch_size] + ) + else: + self._state_indices_long[:batch_size].copy_(self.state_indices[:batch_size]) + self.state_indices_long = self._state_indices_long[:batch_size] + + generation_state_indices = self.state_indices[attn_metadata.num_contexts : batch_size] + if ( + getattr(attn_metadata.kv_cache_manager, "use_kda_replay_update", False) + and generation_state_indices.numel() > 0 + and generation_state_indices.data_ptr() % 16 + ): + aligned_state_indices = self._generation_state_indices[ + : generation_state_indices.shape[0] + ] + aligned_state_indices.copy_(generation_state_indices) + self.generation_state_indices = aligned_state_indices + + self.kda_varlen_is_aligned = ( + all(length % KDA_PREFILL_CHUNK_SIZE == 0 for length in context_lengths) + if context_lengths + else None + ) + self.kda_single_sequence_length = context_lengths[0] if len(context_lengths) == 1 else None + + self.kda_max_seqlen = max(context_lengths, default=0) + n = len(context_lengths) + self.kda_seq_order = self._kda_schedule[0, :n] + self.kda_cu_chunks = self._kda_schedule[1, : n + 1] + if not n: + self.kda_chunk_indices = None + return + + order = sorted(range(n), key=lambda i: (-context_lengths[i], i)) + chunks = list(accumulate(((length + 15) // 16 for length in context_lengths), initial=0)) + width = self._kda_schedule.shape[1] + # Use a fresh pinned source so another prepare cannot overwrite an + # in-flight transfer. The pinned allocator retains it until DMA ends. + host_schedule = torch.tensor( + [order + [0] * (width - n), chunks + [0] * (width - n - 1)], + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._kda_schedule.copy_(host_schedule, non_blocking=True) + + # The CuTe DSL prefill's varlen chunk table (the branch's, 64-token chunks). + num_chunks = sum( + (length + KDA_PREFILL_CHUNK_SIZE - 1) // KDA_PREFILL_CHUNK_SIZE + for length in context_lengths + ) + if self._kda_chunk_indices.shape[0] < num_chunks: + self._kda_chunk_indices = torch.empty((num_chunks, 2), dtype=torch.long, device="cuda") + self.kda_chunk_indices = self._kda_chunk_indices[:num_chunks] + _prepare_kda_chunk_indices(attn_metadata.seq_lens_cuda[:n], self.kda_chunk_indices) + + def _prepare_inforward_split( + self, attn_metadata: AttentionMetadata, kv_cache_manager, plan, num_contexts, batch_size + ) -> None: + """This step's in-forward snapshots (``MambaHybridCacheManagerV2`` with block reuse) as a + ``kda_snapshot`` plan with staging rows. + + A planned request's chunk runs whole: the CuTe DSL prefill stores the recurrent state at + the snapshot point P (``num_before`` tokens into the chunk, a multiple of the KDA chunk), + and the layer leaves the state and convolution window at P in the request's live slot + and the chunk-end ones in its staging row, as ``finish_inforward_snapshots`` expects + (``K3KDALinearAttention.forward_prefill``). + """ + if batch_size != num_contexts: + raise RuntimeError( + "In-forward KDA snapshot planned for a batch with generation requests" + ) + width = kv_cache_manager.conv_state_shape[-1] + context_lengths = attn_metadata.seq_lens[:num_contexts].tolist() + request_ids = attn_metadata.request_ids[:num_contexts] + snapshots, stage_rows = [], [] + for seq, (rid, length) in enumerate(zip(request_ids, context_lengths)): + entry = plan.get(rid) + if entry is None: + continue + before = entry.num_before + if ( + before + entry.num_after != length + or before < width + or before % KDA_PREFILL_CHUNK_SIZE + ): + raise RuntimeError( + f"In-forward KDA snapshot: request {rid} chunk of {length} tokens does " + f"not match the plan ({before} + {entry.num_after}, window {width})" + ) + snapshots.append((seq, kv_cache_manager.state_index(rid), before)) + stage_rows.append(entry.row) + if not snapshots: + return + device = self._state_indices_long.device + rows = torch.tensor(stage_rows, dtype=torch.int64, pin_memory=prefer_pinned()) + self.kda_snapshot = replace( + build_kda_snapshot_plan(snapshots, context_lengths, device), + stage_rows=rows.to(device, non_blocking=True), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dsl_hints.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dsl_hints.py new file mode 100644 index 000000000000..bb81218e271a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dsl_hints.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Instruction-scheduling hints of the internal CuTe DSL build, for the K3 Rubin kernels. + +The internal ``nvidia-cutlass-dsl`` build provides four scheduling hints under ``cute.nvgpu``; the +public package provides none of them. Kernel code calls the wrappers below while it is traced. +When the installed DSL provides all four hints (``HAS_SCHEDULING_HINTS``, fixed at import), each +wrapper emits its hint at the call site; otherwise it emits nothing: + +- ``cfence()`` emits ``cute.nvgpu.cfence()``, a compiler code-motion fence. +- ``warp_switch()`` emits ``cute.nvgpu.warp_switch()``, a warp-switch hint. +- ``sched_res_busy_xu64()`` emits ``cute.nvgpu.sched_res_busy_xu64()`` and + ``reset_sched_res_busy_xu64()`` emits ``cute.nvgpu.reset_sched_res_busy_xu64()``, the MUFU (XU) + resource hints. + +The hints only shape the SASS schedule. The orderings the kernels rely on come from mbarriers, +named barriers, tcgen05 / proxy fences and ``.sync`` instructions, so a kernel computes the same +result with or without them. +""" + +import cutlass.cute as cute + +_CFENCE = getattr(cute.nvgpu, "cfence", None) +_WARP_SWITCH = getattr(cute.nvgpu, "warp_switch", None) +_SCHED_RES_BUSY_XU64 = getattr(cute.nvgpu, "sched_res_busy_xu64", None) +_RESET_SCHED_RES_BUSY_XU64 = getattr(cute.nvgpu, "reset_sched_res_busy_xu64", None) +HAS_SCHEDULING_HINTS = None not in ( + _CFENCE, + _WARP_SWITCH, + _SCHED_RES_BUSY_XU64, + _RESET_SCHED_RES_BUSY_XU64, +) + + +def cfence() -> None: + """Emit ``cute.nvgpu.cfence()`` with the internal build's hints (traced in place).""" + if HAS_SCHEDULING_HINTS: + _CFENCE() + + +def warp_switch() -> None: + """Emit ``cute.nvgpu.warp_switch()`` with the internal build's hints (traced in place).""" + if HAS_SCHEDULING_HINTS: + _WARP_SWITCH() + + +def sched_res_busy_xu64() -> None: + """Emit ``cute.nvgpu.sched_res_busy_xu64()`` with the internal build's hints.""" + if HAS_SCHEDULING_HINTS: + _SCHED_RES_BUSY_XU64() + + +def reset_sched_res_busy_xu64() -> None: + """Emit ``cute.nvgpu.reset_sched_res_busy_xu64()`` with the internal build's hints.""" + if HAS_SCHEDULING_HINTS: + _RESET_SCHED_RES_BUSY_XU64() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_draft_ar.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_draft_ar.py new file mode 100644 index 000000000000..cd36128a8143 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_draft_ar.py @@ -0,0 +1,619 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark drafter all-reduce sites on a pushed MNNVL all-reduce (SM100 / SM103 / SM107). + +The drafter's o_proj and down_proj outputs are TP partial sums that the MNNVL all-reduce reduces +together with the residual add and the following RMSNorm (``AllReduceFusionOp.RESIDUAL_RMS_NORM``). +Here the producer GEMM stores its BF16 output straight into the Lamport slots of the all-reduce +(:mod:`dspark_draft_gemm_push`, or ``k3_lamport_ar.run_lamport_push``), and +:class:`DraftArNormKernel` runs the rest of the call: + +* every thread owns eight consecutive values of a row; the rows are split over a cluster of CTAs + exactly as the C++ launcher splits them (:func:`grid_config`, ``adjustGridConfig``); +* it polls the slots of its eight values (one-shot transport: every rank's slot of the row; + two-shot: the row's owner sums its scatter slots and multicasts the row, the others poll the + broadcast row), sums the ranks in FP32 in rank order and rounds to BF16; +* adds the residual in BF16 (``add.rn.bf16x2``) and stores it as the new residual; +* RMSNorm: BF16 squares summed in element order, warp butterflies (a partial last warp adds zero + for its missing lanes), the warp-0 butterfly over the warp sums, the CTA sums of the row added in + CTA order, ``rsqrt.approx(div.rn(sum, hidden) + eps)``, then the output as the one-shot kernel + (``bf16((x * r) * gamma)``) or as ``rmsNormLamport`` after the two-shot all-reduce + (``bf16((gamma * x) * r)``), whichever C++ path the call stands in for. + +The result is bitwise ``mnnvl_fusion_allreduce(RESIDUAL_RMS_NORM)`` of the pushed input. + +Launch and ordering (the rules of ``k3_lamport_ar.K3LamportFinishKernel``): the kernel releases its +dependents first (the next GEMM streams its weight while the all-reduce is in flight), reads the +workspace's rotation state and polls without waiting for its producer; that is safe because the +producer releases it only after its own ``griddepcontrol.wait``. The residual and gamma it reads are +older than the producer. CTA (0, 0) alone waits for the producer before it advances the rotation +state, after every CTA has counted itself; the CTAs re-arm the previous call's buffer. + +:class:`DraftEmbedKernel` is the drafter's embedding site: the vocab-sharded table lookup followed by +the all-reduce of a tensor whose rows are zero except on the rank that owns the token. The owner +stores its row (``-0`` sent as ``+0``) into slot ``(token, owner)`` of every rank's one-shot buffer +and every rank polls that one slot, so the result is the all-reduce's, bit for bit, in one launch +and one hop. +""" + +from __future__ import annotations + +import functools +import struct + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass.cutlass_dsl import T, dsl_user_op + +from .k3_lamport_ar import ( + NUM_BUFFERS, + clean4, + clear_dirty_stages, + finish_two_shot_v4, + griddepcontrol_launch_dependents, + griddepcontrol_wait, + lamport_state, + ld_acquire_u32, + ld_v4, + ld_volatile_v4, + multimem_st_v4, + pack_bf16x2, + poll_reduce_v4, + red_release_add, + reduce_chunk, + sanitize_bf16x2, + st_u32, + st_v4, + stage_addr, +) +from .k3_latent_ar_norm import ( + _asm, + _bf16x2_f32, + _cluster_sync_release_acquire, + _fadd, + _fmul, + _fv, + _iv, + _lds_f32, + _sq_bf16x2, + _st_cluster_f32, + _sts_f32, +) + +#: Values per thread (one 16-byte access). +ELTS = 8 +#: Word of the CTA arrival counter in the workspace's state flags. +_FLAG_ARRIVALS = 8 +#: Threads per CTA of the embedding site. +EMBED_THREADS = 128 + + +def grid_config(num_tokens: int, hidden: int, sm_count: int) -> tuple[int, int]: + """``(threads per CTA, CTAs per row)`` of the C++ one-shot kernel and of ``rmsNormLamport`` for + BF16 rows (``adjustGridConfig`` with eight values per thread), or ``(0, 0)`` where it + takes more than one load per thread or no cluster.""" + threads = -(-hidden // ELTS) + cluster = 8 + while threads % cluster and cluster > 1: + cluster //= 2 + max_cluster = cluster + block = -(-threads // cluster) + while block < 128 and cluster >= 2: + block *= 2 + cluster //= 2 + while num_tokens * cluster > sm_count and cluster > 1 and block <= 512: + block *= 2 + cluster //= 2 + while cluster < max_cluster: + cand = cluster * 2 + cand_block = -(-threads // cand) + if cand_block < 64 or num_tokens * cand > sm_count: + break + cluster, block = cand, cand_block + if block > 1024 or block * cluster * ELTS != hidden: + return 0, 0 + return block, cluster + + +@dsl_user_op +def _add_bf16x2(a, b, *, loc=None, ip=None): + """Elementwise BF16 sum of two pairs, rounded to nearest (``__hadd2``).""" + return Int32( + _asm( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "add.rn.bf16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_add(value, offset: int, lanes: int, lane, *, loc=None, ip=None): + """One butterfly step of the C++ warp sums over the first ``lanes`` lanes of a warp: + ``value + (partner < lanes ? partner's value : 0)``.""" + mask = (1 << lanes) - 1 if lanes < 32 else 0xFFFFFFFF + return Float32( + _asm( + T.f32(), + [_fv(value, loc=loc, ip=ip), _iv(lane, loc=loc, ip=ip)], + "{\n.reg .pred p;\n.reg .b32 t;\n.reg .f32 o;\n" + f"shfl.sync.bfly.b32 o, $1, {offset}, 0x1f, {mask};\n" + f"xor.b32 t, $2, {offset};\nsetp.lt.u32 p, t, {lanes};\n" + "selp.f32 o, o, 0f00000000, p;\nadd.rn.f32 $0, $1, o;\n}", + "=f,f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _rms_scale(total, eps, inv_hidden_hex: str, *, loc=None, ip=None): + """``rsqrtf(total / hidden + eps)`` as the C++ kernels compute it (IEEE division, then add).""" + return Float32( + _asm( + T.f32(), + [_fv(total, loc=loc, ip=ip), _fv(eps, loc=loc, ip=ip)], + "{\n.reg .f32 q;\n" + f"div.rn.f32 q, $1, 0f{inv_hidden_hex};\n" + "add.rn.f32 q, q, $2;\nrsqrt.approx.f32 $0, q;\n}", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +class DraftArNormKernel: + """``norm(allreduce(pushed) + residual)`` and the new residual for ``M`` rows of ``hidden`` BF16 + values on ``world`` ranks (see the module docstring).""" + + def __init__( + self, + world: int, + hidden: int, + threads: int, + cluster: int, + one_shot: bool, + cpp_one_shot: bool, + use_pdl: bool, + ): + if threads * cluster * ELTS != hidden: + raise ValueError(f"{cluster} x {threads} threads do not cover {hidden} values") + self.world = world + self.hidden = hidden + self.threads = threads + self.cluster = cluster + self.one_shot = one_shot + self.cpp_one_shot = cpp_one_shot + self.use_pdl = use_pdl + self.warps = (threads + 31) // 32 + self.last_lanes = threads - 32 * (self.warps - 1) + self.chunk = reduce_chunk(world, cpp_one_shot) + self.hidden_hex = struct.pack(">f", float(hidden)).hex() + + @cute.jit + def __call__( + self, + uc_local: Int64, + mc_base: Int64, + flags: Int64, + res_in: Int64, + gamma: Int64, + out: Int64, + res_out: Int64, + m: Int32, + rank: Int32, + eps: Float32, + stream: cuda.CUstream, + ): + @cute.struct + class SharedStorage: + warp_sum: cute.struct.MemRange[cutlass.Float32, 32] + block_sum: cute.struct.MemRange[cutlass.Float32, 1] + cta_sum: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, 8], 16] + + self.shared_storage = SharedStorage + self.kernel(uc_local, mc_base, flags, res_in, gamma, out, res_out, m, rank, eps).launch( + grid=(m, self.cluster, 1), + block=(self.threads, 1, 1), + cluster=(1, self.cluster, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + uc_local: Int64, + mc_base: Int64, + flags: Int64, + res_in: Int64, + gamma: Int64, + out: Int64, + res_out: Int64, + m: Int32, + rank: Int32, + eps: Float32, + ): + W, H, threads = self.world, self.hidden, self.threads + tidx, _, _ = cute.arch.thread_idx() + tok, cy, _ = cute.arch.block_idx() + warp = tidx // 32 + lane = tidx % 32 + storage = utils.SmemAllocator().allocate(self.shared_storage) + warp_sum = storage.warp_sum.data_ptr().toint() + block_sum = storage.block_sum.data_ptr().toint() + cta_sum = storage.cta_sum.data_ptr().toint() + + if cutlass.const_expr(self.use_pdl): + griddepcontrol_launch_dependents() + cur, dirty, nbytes, dstages, c0, c1, c2, c3 = lamport_state(flags) + # Count this CTA once every thread has read the state. + cute.arch.barrier() + if tidx == 0: + red_release_add(flags + Int64(_FLAG_ARRIVALS * 4), Int32(1)) + + col = (cy * threads + tidx) * ELTS + row_off = Int64(tok * H + col) * 2 + rw = ld_v4(res_in + row_off) + gw = ld_v4(gamma + Int64(col) * 2) + + # Re-arm the previous call's buffer while the rows travel. + clear_dirty_stages( + uc_local, + dirty, + nbytes, + dstages, + c0, + c1, + c2, + c3, + (tok * self.cluster + cy) * threads + tidx, + m * self.cluster * threads, + ) + + red = [Int32(0)] * 4 + if cutlass.const_expr(self.one_shot): + buf = stage_addr(uc_local, cur, Int32(0), 1, nbytes) + acc = poll_reduce_v4(buf + Int64(tok * W * H + col) * 2, H * 2, W, self.chunk) + red = [pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)] + else: + w0, w1, w2, w3 = finish_two_shot_v4( + uc_local, mc_base, cur, nbytes, tok, col * 2, rank, Int32(1), W, self.chunk, H + ) + red = [w0, w1, w2, w3] + + y = [_add_bf16x2(red[j], rw[j]) for j in range(4)] + st_v4(res_out + row_off, y[0], y[1], y[2], y[3]) + + # RMSNorm: this thread's BF16 squares in element order. + thread_sum = Float32(0.0) + for j in cutlass.range_constexpr(4): + lo, hi = _bf16x2_f32(_sq_bf16x2(y[j])) + thread_sum = _fadd(thread_sum, lo) + thread_sum = _fadd(thread_sum, hi) + v = thread_sum + if cutlass.const_expr(self.last_lanes < 32): + if warp == self.warps - 1: + for off in (16, 8, 4, 2, 1): + v = _shfl_add(v, off, self.last_lanes, lane) + else: + for off in (16, 8, 4, 2, 1): + v = _shfl_add(v, off, 32, lane) + else: + for off in (16, 8, 4, 2, 1): + v = _shfl_add(v, off, 32, lane) + if lane == 0: + _sts_f32(warp_sum + warp * 4, v) + cute.arch.barrier() + if warp == 0: + b = Float32(0.0) + if lane < self.warps: + b = _lds_f32(warp_sum + lane * 4) + for off in (16, 8, 4, 2, 1): + b = _shfl_add(b, off, 32, lane) + if lane == 0: + _sts_f32(block_sum, b) + cute.arch.barrier() + total = _lds_f32(block_sum) + if cutlass.const_expr(self.cluster > 1): + # The row's CTA sums: CTA c's into every peer's slot c, then added in CTA order. + if tidx < self.cluster: + _st_cluster_f32(cta_sum + cy * 4, tidx, total) + _cluster_sync_release_acquire() + total = Float32(0.0) + for c in cutlass.range_constexpr(self.cluster): + total = _fadd(total, _lds_f32(cta_sum + c * 4)) + scale = _rms_scale(total, eps, self.hidden_hex) + + ow = [] + for j in cutlass.range_constexpr(4): + x_lo, x_hi = _bf16x2_f32(y[j]) + g_lo, g_hi = _bf16x2_f32(gw[j]) + if cutlass.const_expr(self.cpp_one_shot): + ow.append( + pack_bf16x2(_fmul(_fmul(x_lo, scale), g_lo), _fmul(_fmul(x_hi, scale), g_hi)) + ) + else: + ow.append( + pack_bf16x2(_fmul(_fmul(g_lo, x_lo), scale), _fmul(_fmul(g_hi, x_hi), scale)) + ) + st_v4(out + row_off, ow[0], ow[1], ow[2], ow[3]) + + # CTA (0, 0): the next call's state, once the producer (which reads the state after its own + # wait) has completed and every CTA has read the state. + if tok == 0: + if cy == 0: + if tidx == 0: + if cutlass.const_expr(self.use_pdl): + griddepcontrol_wait() + arrivals = flags + Int64(_FLAG_ARRIVALS * 4) + total_ctas = m * self.cluster + arrived = ld_acquire_u32(arrivals) + while arrived < total_ctas: + arrived = ld_acquire_u32(arrivals) + nxt = cur + 1 + if nxt == NUM_BUFFERS: + nxt = Int32(0) + row_bytes = H * 2 + if cutlass.const_expr(self.one_shot): + st_v4(flags, nxt, cur, nbytes, Int32(1)) + st_v4(flags + Int64(16), m * (W * row_bytes), Int32(0), Int32(0), Int32(0)) + else: + st_v4(flags, nxt, cur, nbytes, Int32(2)) + st_v4( + flags + Int64(16), + ((m + (W - 1)) // W) * (W * row_bytes), + m * row_bytes, + Int32(0), + Int32(0), + ) + st_u32(arrivals, Int32(0)) + + +class DraftEmbedKernel: + """``all_reduce(vocab-sharded lookup)`` of ``M`` token ids: the owner rank's table row of each + token, ``-0`` as ``+0``, zero for ids outside ``[0, vocab)`` (see the module docstring).""" + + def __init__(self, world: int, hidden: int, anchor64: bool, use_pdl: bool): + if hidden % (ELTS * EMBED_THREADS): + raise ValueError(f"hidden={hidden} must be a multiple of {ELTS * EMBED_THREADS}") + self.world = world + self.hidden = hidden + self.ids64 = anchor64 + self.use_pdl = use_pdl + self.cols = hidden // (ELTS * EMBED_THREADS) + + @cute.jit + def __call__( + self, + ids: Int64, + table: Int64, + uc_local: Int64, + mc_base: Int64, + flags: Int64, + out: Int64, + m: Int32, + rank: Int32, + slice_width: Int32, + vocab: Int32, + stream: cuda.CUstream, + ): + self.kernel(ids, table, uc_local, mc_base, flags, out, m, rank, slice_width, vocab).launch( + grid=(m, self.cols, 1), + block=(EMBED_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + ids: Int64, + table: Int64, + uc_local: Int64, + mc_base: Int64, + flags: Int64, + out: Int64, + m: Int32, + rank: Int32, + slice_width: Int32, + vocab: Int32, + ): + W, H = self.world, self.hidden + tidx, _, _ = cute.arch.thread_idx() + tok, cy, _ = cute.arch.block_idx() + if cutlass.const_expr(self.use_pdl): + griddepcontrol_wait() + griddepcontrol_launch_dependents() + cur, dirty, nbytes, dstages, c0, c1, c2, c3 = lamport_state(flags) + cute.arch.barrier() + if tidx == 0: + red_release_add(flags + Int64(_FLAG_ARRIVALS * 4), Int32(1)) + + col = (cy * EMBED_THREADS + tidx) * ELTS + token = _ld_id(ids + Int64(tok) * (8 if self.ids64 else 4)) + owner = Int32(-1) + if token >= 0: + if token < vocab: + owner = token // slice_width + buf_off = Int64(cur) * Int64(nbytes) + slot = Int64(((tok * W + owner) * H + col) * 2) + if owner == rank: + w = ld_v4(table + Int64((token - rank * slice_width) * H + col) * 2) + multimem_st_v4( + mc_base + buf_off + slot, + sanitize_bf16x2(w[0]), + sanitize_bf16x2(w[1]), + sanitize_bf16x2(w[2]), + sanitize_bf16x2(w[3]), + ) + clear_dirty_stages( + uc_local, + dirty, + nbytes, + dstages, + c0, + c1, + c2, + c3, + (tok * self.cols + cy) * EMBED_THREADS + tidx, + m * self.cols * EMBED_THREADS, + ) + r0 = Int32(0) + r1 = Int32(0) + r2 = Int32(0) + r3 = Int32(0) + if owner >= 0: + w = ld_volatile_v4(uc_local + buf_off + slot) + while clean4(w[0], w[1], w[2], w[3]) == 0: + w = ld_volatile_v4(uc_local + buf_off + slot) + r0 = w[0] + r1 = w[1] + r2 = w[2] + r3 = w[3] + st_v4(out + Int64(tok * H + col) * 2, r0, r1, r2, r3) + + if tok == 0: + if cy == 0: + if tidx == 0: + arrivals = flags + Int64(_FLAG_ARRIVALS * 4) + arrived = ld_acquire_u32(arrivals) + while arrived < m * self.cols: + arrived = ld_acquire_u32(arrivals) + nxt = cur + 1 + if nxt == NUM_BUFFERS: + nxt = Int32(0) + st_v4(flags, nxt, cur, nbytes, Int32(1)) + st_v4(flags + Int64(16), m * (W * H * 2), Int32(0), Int32(0), Int32(0)) + st_u32(arrivals, Int32(0)) + + +@dsl_user_op +def _ld_id(addr, *, loc=None, ip=None): + """A token id (an int32, or the low word of an int64) written by an earlier grid.""" + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def _stream() -> cuda.CUstream: + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + +@functools.cache +def _compiled_ar_norm(world, hidden, threads, cluster, one_shot, cpp_one_shot, use_pdl): + kernel = DraftArNormKernel(world, hidden, threads, cluster, one_shot, cpp_one_shot, use_pdl) + a = Int64(1 << 20) + return cute.compile(kernel, a, a, a, a, a, a, a, Int32(1), Int32(0), Float32(1e-6), _stream()) + + +@functools.cache +def _compiled_embed(world, hidden, ids64, use_pdl): + kernel = DraftEmbedKernel(world, hidden, ids64, use_pdl) + a = Int64(1 << 20) + i = Int32(1) + return cute.compile(kernel, a, a, a, a, a, a, i, i, i, i, _stream()) + + +@functools.cache +def _sm_count(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def ar_norm_config(num_tokens: int, hidden: int, device: torch.device) -> tuple[int, int]: + """``(threads, cluster)`` of :class:`DraftArNormKernel` for this call, ``(0, 0)`` if none.""" + return grid_config(num_tokens, hidden, _sm_count(device.index)) + + +def run_draft_ar_norm( + workspace, + one_shot: bool, + cpp_one_shot: bool, + residual: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + residual_out: torch.Tensor, + use_pdl: bool = True, +) -> None: + """Finish the all-reduce every rank pushed into ``workspace`` (a ``k3_lamport_ar.LamportWorkspace``) + with the ``one_shot`` or two-shot transport, add ``residual`` and RMSNorm with ``gamma``: + ``out`` / ``residual_out`` (contiguous BF16 ``[M, hidden]``) as the C++ fused all-reduce + (``cpp_one_shot``: its one-shot kernel, else two-shot + ``rmsNormLamport``) returns them.""" + m, hidden = out.shape + threads, cluster = ar_norm_config(m, hidden, out.device) + if not threads: + raise ValueError(f"no single-load launch shape for {m} x {hidden}") + _compiled_ar_norm(workspace.world, hidden, threads, cluster, one_shot, cpp_one_shot, use_pdl)( + Int64(workspace.uc_local), + Int64(workspace.mc_base), + Int64(workspace.flags.data_ptr()), + Int64(residual.data_ptr()), + Int64(gamma.data_ptr()), + Int64(out.data_ptr()), + Int64(residual_out.data_ptr()), + Int32(m), + Int32(workspace.rank), + Float32(eps), + _stream(), + ) + + +def run_draft_embed( + workspace, + ids: torch.Tensor, + table: torch.Tensor, + slice_width: int, + vocab: int, + out: torch.Tensor, + use_pdl: bool = True, +) -> None: + """``out`` (contiguous BF16 ``[M, hidden]``) = the all-reduce of the vocab-sharded lookup of + ``ids`` (contiguous int32 / int64 ``[M]``) in this rank's ``table`` rows + ``[rank * slice_width, ...)``, through the one-shot buffer of ``workspace``.""" + m, hidden = out.shape + _compiled_embed(workspace.world, hidden, ids.dtype == torch.int64, use_pdl)( + Int64(ids.data_ptr()), + Int64(table.data_ptr()), + Int64(workspace.uc_local), + Int64(workspace.mc_base), + Int64(workspace.flags.data_ptr()), + Int64(out.data_ptr()), + Int32(m), + Int32(workspace.rank), + Int32(slice_width), + Int32(vocab), + _stream(), + ) + + +__all__ = [ + "DraftArNormKernel", + "DraftEmbedKernel", + "ar_norm_config", + "grid_config", + "run_draft_ar_norm", + "run_draft_embed", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_draft_gemm_push.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_draft_gemm_push.py new file mode 100644 index 000000000000..17a85079a67f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_draft_gemm_push.py @@ -0,0 +1,692 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The low-M BF16 split-K GEMM (:mod:`low_m_bf16_splitk`) with an all-reduce push epilogue. + +``out[M, N] = A[M, K] @ W[N, K]^T`` for a TP partial sum that an MNNVL all-reduce reduces next: +instead of storing its BF16 tile, the owner CTA of each cluster (after the split-K reduction) stages +the tile token-major in shared memory and stores it as 16-byte vectors into the Lamport slots of the +call (``k3_lamport_ar.push_slot`` / ``push_v4``: every rank's slot through the multicast address for +the one-shot transport, the token owner's scatter slot for the two-shot one; ``-0`` sent as ``+0``). +A kernel that finishes the call (``dspark_draft_ar.DraftArNormKernel``, +``k3_lamport_ar.run_lamport_finish``) then gives the all-reduce of the GEMM output, bit for bit +equal to this GEMM's plain output followed by the all-reduce. + +PDL: the weight warp streams before ``griddepcontrol.wait`` as in the plain kernel, but the +dependents are released by the activation warp right after its wait (not by the weight warp): the +finish kernel reads the workspace's rotation state as soon as it is resident, which is safe only +once every earlier grid has completed. The epilogue reads the state after its own wait. + +Rows: up to :data:`MAX_M` = 64 with ``mma_n`` up to 64 (the vendored kernel's own policy stops at +32 rows / ``mma_n`` 32). :func:`run_gemm` runs the same kernel without the push (plain store). +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils.blackwell_helpers as sm100_utils +import torch as _torch +from cutlass import Int32, Int64 +from cutlass.cute import experimental as cute_ext +from cutlass.cutlass_dsl import T, dsl_user_op + +from . import k3_lamport_ar +from .k3_latent_ar_norm import _asm, _iv, _unpack +from .low_m_bf16_splitk import ( + _AB_BUFFER_ALIGN_BYTES, + _MAILBOX_ALIGN_BYTES, + _MBARRIER_BYTES, + _SMEM_CAPACITY_BYTES, + _TMEM_POINTER_BYTES, + OWNER_RANK, + SplitKDenseGemmKernel, + SplitKTactic, + _align_up, + _from_dlpack_dynamic, + _smem_bytes, + _store_shared_remote_v4, + validate_tactic, +) + +#: Alignment of the token-major staging tile. +_STAGE_ALIGN_BYTES = 16 + + +class _NoopKernel: + """An empty ``cute.compile`` kernel (see :func:`ensure_cute_runtime`).""" + + @cute.jit + def __call__(self, stream: _cuda.CUstream): + self.kernel().launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) + + @cute.kernel + def kernel(self): + pass + + +@functools.cache +def ensure_cute_runtime() -> None: + """Load and run one ``cute.compile`` module in this process. With the image's cutlass-dsl, the + first ``cute.experimental`` module a process loads fails with ``cudaErrorInvalidValue`` unless + a ``cute.compile`` module was loaded before; this makes the split-K kernels safe to be first.""" + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + cute.compile(_NoopKernel(), stream)(stream) + + +def smem_bytes(tactic: SplitKTactic) -> int: + """Shared memory of the push kernel: the plain kernel's plus the staging tile.""" + return ( + _align_up(_smem_bytes(tactic, tactic.ab_stages), _STAGE_ALIGN_BYTES) + + tactic.mma_m * tactic.mma_n * 2 + ) + + +#: Largest row count and the UMMA N tiles (rows per CTA) of the push kernel. +MAX_M = 64 +_MMA_N = (8, 16, 32, 64) + + +def validate_push_tactic(tactic: SplitKTactic, m: int, n: int, k: int) -> None: + """Reject a tactic the push kernel cannot run for ``x [m, k] @ w[n, k]^T``.""" + if m > 32 or tactic.mma_n > 32: + # The vendored rules, with the row bound and UMMA N widened. + if not 1 <= m <= MAX_M or tactic.mma_n not in _MMA_N: + raise ValueError(f"unsupported rows {m} / mma_n {tactic.mma_n}") + validate_tactic(dataclasses.replace(tactic, mma_n=32), 1, n, k) + if _smem_bytes(tactic, tactic.ab_stages) > _SMEM_CAPACITY_BYTES: + raise ValueError(f"tactic {tactic} exceeds shared memory") + else: + validate_tactic(tactic, m, n, k) + if smem_bytes(tactic) > _SMEM_CAPACITY_BYTES: + raise ValueError(f"tactic {tactic} with the staging tile exceeds shared memory") + if tactic.mma_m % 8: + raise ValueError("the push epilogue stores 8-column vectors") + + +@dsl_user_op +def _lds_v4(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm + + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +def _lamport_cur(flags): + """``(current buffer, bytes per buffer)`` of the workspace's rotation state.""" + f = k3_lamport_ar.ld_volatile_v4(flags) + return f[0], f[2] + + +def _push_vector(peer_ptrs, mc_base, cur, nbytes, tok, rank, col, words, world, hidden, one_shot): + """Store eight BF16 values (``-0`` as ``+0``) of row ``tok`` at column ``col`` into this + rank's slot of the call in buffer ``cur`` (``k3_lamport_ar.push_slot`` / ``push_v4``).""" + w = [k3_lamport_ar.sanitize_bf16x2(x) for x in words] + row_bytes = hidden * 2 + if one_shot: + addr = ( + mc_base + + Int64(cur) * Int64(nbytes) + + Int64(tok * world + rank) * row_bytes + + Int64(col) * 2 + ) + k3_lamport_ar.multimem_st_v4(addr, w[0], w[1], w[2], w[3]) + else: + owner = k3_lamport_ar.ld_nc_u64(peer_ptrs + Int64(tok % world) * 8) + addr = ( + owner + + Int64(cur * 2) * Int64(nbytes // 2) + + Int64((tok // world) * world + rank) * row_bytes + + Int64(col) * 2 + ) + k3_lamport_ar.st_v4(addr, w[0], w[1], w[2], w[3]) + + +class SplitKArPushGemmKernel(SplitKDenseGemmKernel): + """:class:`SplitKDenseGemmKernel` whose owner CTAs push the BF16 tile into the Lamport slots of + an all-reduce of ``world`` ranks over ``hidden`` (= N) columns (see the module docstring). + ``store_out`` also stores the tile into ``out`` as the plain kernel does; without ``push`` the + kernel only stores (no workspace access).""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + world: int, + hidden: int, + one_shot: bool, + store_out: bool, + push: bool = True, + ) -> None: + super().__init__(tactic=tactic, use_pdl=use_pdl, has_bias=False) + self.world = world + self.hidden = hidden + self.one_shot = one_shot + self.store_out = store_out + self.push = push + self.vecs_per_token = tactic.mma_m // 8 + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + m: Int32, + stream: _cuda.CUstream, + ): + self.kernel(a, b, c, peer_ptrs, mc_base, flags, rank, m).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(_SMEM_CAPACITY_BYTES), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + m: Int32, + ): + """The plain kernel's warp dispatch plus the staging tile and the push arguments.""" + stages = self.num_ab_stage + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + cutlass.utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + cutlass.utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + 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()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + mailbox = sA + bar_reduce = bar_mma_epilog + # The owner's BF16 tile, token-major: (hidden column, token) -> column + token * cta_m. + stage = cute_ext.allocate( + mC.element_type, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=_STAGE_ALIGN_BYTES, + ) + + 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): + 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() + else: + cute.arch.barrier() + + 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() + + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.push_epilog_warp( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + cutlass.utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + stage, + bidx, + n_idx, + peer_ptrs, + mc_base, + flags, + rank, + m, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + """The plain kernel's TMA producer; the activation warp releases the dependents right + after its ``griddepcontrol.wait`` and the weight warp never does.""" + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + 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, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + 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=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def push_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, + stage: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + m: Int32, + ): + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + 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) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + 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, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate(self.acc_dtype, cute.AddressSpace.rmem, rmem_layout, alignment=32) + rD = cute_ext.allocate(c_dtype, cute.AddressSpace.rmem, rmem_layout, alignment=32) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + rD.store(rAcc.load().to(c_dtype)) + if cutlass.const_expr(self.store_out): + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + if cutlass.const_expr(self.push): + cute_ext.partition_and_copy(thr_t2r, rD, stage) + # The four epilogue warps only. + cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) + # This epilogue thread's 16-byte vectors of the tile, token-major. + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + cur, nbytes = _lamport_cur(flags) + stage_base = stage.iterator.toint() + vecs = self.cta_n * self.vecs_per_token + for it in cutlass.range_constexpr( + (vecs + self.epilog_threads - 1) // self.epilog_threads + ): + v = epi_tid + it * self.epilog_threads + t = v // self.vecs_per_token + col = (v % self.vecs_per_token) * 8 + tok = n_idx * self.cta_n + t + if v < vecs: + if tok < m: + _push_vector( + peer_ptrs, + mc_base, + cur, + nbytes, + tok, + rank, + bidx * self.cta_m + col, + _lds_v4(stage_base + (t * self.cta_m + col) * 2), + self.world, + self.hidden, + self.one_shot, + ) + + +@cute.experimental.jit +def _bmm_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + m: Int32, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + peer_ptrs, + mc_base, + flags, + rank, + m, + stream, + ) + + +def _make_layout_tensor(shape, dtype, leading_dim): + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +@functools.cache +def _compiled(dtype, tactic, use_pdl, world, hidden, one_shot, store_out, push=True): + kernel = SplitKArPushGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + world=world, + hidden=hidden, + one_shot=one_shot, + store_out=store_out, + push=push, + ) + # The swapped operands of ``x [M, K] @ w[N, K]^T``: A = w (batch, N, K) K-major, + # B = x^T (batch, K, M) K-major, C = out^T (batch, N, M) N-major. + m, n, k = 8, 64, 128 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading), leading) + for shape, leading in (((1, n, k), 2), ((1, k, m), 1), ((1, n, m), 1)) + ) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + a = Int64(1 << 20) + return cute_ext.compile(_bmm_push, kernel, *tensors, a, a, a, Int32(0), Int32(1), stream) + + +def _swapped(x, weight, out): + return ( + _from_dlpack_dynamic(weight.unsqueeze(0), 2), + _from_dlpack_dynamic(x.unsqueeze(0).transpose(-2, -1), 1), + _from_dlpack_dynamic(out.unsqueeze(0).transpose(-2, -1), 1), + ) + + +def run_gemm( + x: _torch.Tensor, + weight: _torch.Tensor, + out: _torch.Tensor, + tactic: SplitKTactic, + *, + pdl: bool = True, +) -> _torch.Tensor: + """``out = x [M, K] @ weight[N, K]^T`` (BF16, contiguous) with this kernel's store-only + variant: the same values the push variant pushes, for ``M`` up to :data:`MAX_M`. Returns + ``out``.""" + m, k = x.shape + n = weight.shape[0] + validate_push_tactic(tactic, m, n, k) + ensure_cute_runtime() + _compiled(x.dtype, tactic, pdl, 2, n, True, True, False)( + *_swapped(x, weight, out), + Int64(0), + Int64(0), + Int64(0), + Int32(0), + Int32(m), + _cuda.CUstream(_torch.cuda.current_stream(x.device).cuda_stream), + ) + return out + + +def run_gemm_push( + x: _torch.Tensor, + weight: _torch.Tensor, + out: _torch.Tensor, + target, + tactic: SplitKTactic, + *, + pdl: bool = True, + store_out: bool = False, +) -> None: + """Push ``x [M, K] @ weight[N, K]^T`` (BF16) into ``target`` (an + ``ops.LamportPushTarget`` of an all-reduce of ``[M, N]``) and set ``target.pushed``; with + ``store_out`` also write it into ``out`` (``[M, N]``, contiguous; its shape always sets the + launch).""" + m, k = x.shape + n = weight.shape[0] + validate_push_tactic(tactic, m, n, k) + ensure_cute_runtime() + ws = target.workspace + compiled = _compiled(x.dtype, tactic, pdl, ws.world, n, target.one_shot, store_out) + compiled( + *_swapped(x, weight, out), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(ws.rank), + Int32(m), + _cuda.CUstream(_torch.cuda.current_stream(x.device).cuda_stream), + ) + target.pushed = True + + +__all__ = [ + "MAX_M", + "SplitKArPushGemmKernel", + "ensure_cute_runtime", + "run_gemm", + "run_gemm_push", + "smem_bytes", + "validate_push_tactic", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_markov_greedy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_markov_greedy.py new file mode 100644 index 000000000000..234e63d4317b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_markov_greedy.py @@ -0,0 +1,974 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Adapted from vLLM (https://github.com/vllm-project/vllm) under the Apache-2.0 license. +# Original file: vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/draft/markov.py +# Original copyright: Copyright contributors to the vLLM project +"""DSpark greedy Markov block chain over a vocab-sharded LM head, one launch (SM100 / SM103 / SM107). + +Every TP rank holds its vocab shard of the block's base logits ``[B, K, shard]`` (BF16) and the +matching rows of ``markov_w2``. For each block position ``i`` the chain adds +``markov_w2 @ markov_w1[prev_i]`` to the position's logits and takes the global argmax, which is +draft token ``i`` and ``prev_{i+1}``; ``prev_0`` is the anchor, and an anchor outside +``[0, vocab)`` contributes a zero bias. + +Grid: one CTA of 512 threads per 256 shard rows, all resident at once (each CTA waits for every +other CTA's keys). Per CTA: + +* before ``griddepcontrol.wait``: its 256 ``markov_w2`` rows are copied into shared memory + (``cp.async``, L2 evict-first) and the exchange set of the launch before the previous one is + cleared; +* per position: warp ``w`` multiplies rows ``[16 w, 16 w + 16)`` by the requests' ``markov_w1`` rows + with ``mma.sync.m16n8k16`` (BF16 in, FP32 accumulated over the Markov rank in order), rounds the + bias to BF16 (``F.linear``'s output), adds the BF16 base logit in FP32 (the reference adds the + BF16 bias to the FP32-cast logits, without a further rounding) and reduces each request to a + 64-bit key ``ordered(value) << 32 | ~index``. The CTA's key per request is stored into the + exchange of every rank (``multimem.st``); every CTA then polls all ``tp x CTAs`` keys of the + position and takes their maximum. Larger keys win: the larger value, then the lower global + index, NaN above everything and ``-0 == +0``, which is ``torch.argmax`` on the gathered logits. + +Exchange: three sets of ``[position, entry = rank x CTAs + CTA, request]`` u64 keys, 0 where +nothing was written (a key never has a zero 32-bit word). Launch ``L`` (from a per-rank ordinal +that CTA 0 advances after its last poll) uses set ``L % 3`` and clears set ``(L + 2) % 3``, which +launch ``L - 1`` used: every CTA of this rank finished reading it before launch ``L - 1`` +completed, and peers write it again only in launch ``L + 2``, which needs this rank's keys of +launch ``L + 1``. The kernel never triggers its dependents early, so the next launch starts only +after this one has completed. + +Options (per request width, :data:`OPTIONS`): ``staged`` -- per position, warp ``r`` copies request +``r``'s ``markov_w1`` row into shared memory (one coalesced 16-byte load per lane) and every warp +takes its MMA B fragments with ``ldmatrix`` instead of 32 scattered 4-byte loads per lane; +``gate`` -- before the sweep over all keys, lane ``s`` of warp 0 waits for the chunk holding the key +of rank ``s``'s last CTA, so fewer polls are outstanding while the keys are still in flight. +""" + +from __future__ import annotations + +import functools +from typing import Optional + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Shard rows per CTA: 16 warps x one 16-row MMA tile. +ROWS = 256 +THREADS = 512 +WARPS = THREADS // 32 +#: Largest block size and request count a launch serves (the exchange set is sized for them). +MAX_STEPS = 16 +MAX_WIDTH = 16 +#: Exchange sets in rotation. +NUM_SETS = 3 +#: Int32 words of the per-rank state: launch ordinal, then the bytes each set's last launch used. +STATE_WORDS = 4 +#: ``(staged, gate)`` per request width. Rubin R200, TP4 on one node, K = 7, 640 keys per request +#: and position (40960-row shards), CUDA graph, us per call, plain / staged / gate / both: +#: n1 32.7 / 32.2 / 28.7 / 28.2, n2 43.7 / 41.7 / 32.3 / 30.3, n4 41.8 / 36.5 / 34.7 / 29.5, +#: n8 42.7 / 31.9 / 40.0 / 30.1, n16 57.5 / 36.4 / 56.1 / 35.3 (scripts/draftfuse/bench_markov.py). +OPTIONS = {1: (True, True), 2: (True, True), 4: (True, True), 8: (True, True), 16: (True, True)} + + +def set_bytes(entries: int) -> int: + """Bytes of one exchange set for ``entries`` (TP ranks x CTAs per rank).""" + return MAX_STEPS * entries * MAX_WIDTH * 8 + + +def ctas_for(shard: int) -> int: + return (shard + ROWS - 1) // ROWS + + +def width_for(num_reqs: int) -> int: + """Requests per key row: the next power of two.""" + return 1 << max(num_reqs - 1, 0).bit_length() + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _lv(value, *, loc=None, ip=None): + return Int64(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return Float32(value).ir_value(loc=loc, ip=ip) + + +def _unpack(res, ty, cls, n, *, loc=None, ip=None): + return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n)) + + +@dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.wait;", "~{memory}", loc=loc, ip=ip) + + +@dsl_user_op +def _cp_async_evict_first(smem_addr, gmem_addr, *, loc=None, ip=None): + """16-byte ``cp.async`` of static weights with an L2 evict-first policy.""" + _asm( + None, + [_iv(smem_addr, loc=loc, ip=ip), _lv(gmem_addr, loc=loc, ip=ip)], + "{\n.reg .b64 pol;\ncreatepolicy.fractional.L2::evict_first.b64 pol, 1.0;\n" + "cp.async.cg.shared.global.L2::cache_hint [$0], [$1], 16, pol;\n}", + "r,l,~{memory}", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cp_async_commit(*, loc=None, ip=None): + _asm(None, [], "cp.async.commit_group;", "~{memory}", loc=loc, ip=ip) + + +@dsl_user_op +def _cp_async_wait_all(*, loc=None, ip=None): + _asm(None, [], "cp.async.wait_group 0;", "~{memory}", loc=loc, ip=ip) + + +@dsl_user_op +def _ld_volatile_v4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.volatile.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l,~{memory}", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _st_zero_v4(addr, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)], + "st.global.v4.b32 [$0], {0, 0, 0, 0};", + "l,~{memory}", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r,~{memory}", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_anchor(addr, ok, *, loc=None, ip=None): + """The anchor token at ``addr`` (an int32, or the low word of an int64) if ``ok``, else -1.""" + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip), _iv(ok, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.s32 p, $2, 0;\nmov.b32 $0, -1;\n@p ld.global.b32 $0, [$1];\n}", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_bf16_bits(addr, ok, *, loc=None, ip=None): + """The BF16 at ``addr`` zero-extended to 32 bits if ``ok``, else 0 (a logit the previous + grid wrote: read after ``griddepcontrol.wait``).""" + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip), _iv(ok, loc=loc, ip=ip)], + "{\n.reg .pred p;\n.reg .b16 h;\nsetp.ne.s32 p, $2, 0;\nmov.b16 h, 0;\n" + "@p ld.global.b16 h, [$1];\ncvt.u32.u16 $0, h;\n}", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_nc_b32(addr, *, loc=None, ip=None): + """A 32-bit word of a static weight (read-only path).""" + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_nc_v4_masked(addr, mask, *, loc=None, ip=None): + """16 bytes of a static weight (read-only path), each word ANDed with ``mask``.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip), _iv(mask, loc=loc, ip=ip)], + "{\n.reg .b32 a, b, c, d;\nld.global.nc.v4.b32 {a, b, c, d}, [$4];\n" + "and.b32 $0, a, $5;\nand.b32 $1, b, $5;\nand.b32 $2, c, $5;\nand.b32 $3, d, $5;\n}", + "=r,=r,=r,=r,l,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _st_shared_v4(addr, w0, w1, w2, w3, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)] + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r,~{memory}", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ldsm_x4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _mma_bf16(a, b0, b1, c, *, loc=None, ip=None): + """``d[16x8] = a[16x16] @ b[16x8] + c`` (BF16 inputs, FP32 accumulate).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(x, loc=loc, ip=ip) for x in (*a, b0, b1)] + [_fv(x, loc=loc, ip=ip) for x in c], + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _logit_key(acc, base_bits, index, ok, *, loc=None, ip=None): + """Key of ``bf16(acc) + base`` (FP32 add, round to nearest) at vocab ``index``, or 0 if not + ``ok``. High word: the value's bits mapped to an unsigned order (``-0`` as ``+0``, every NaN + above +inf); low word: ``~index`` (the lower index wins a tie). Neither word is ever 0.""" + return Int64( + _asm( + T.i64(), + [ + _fv(acc, loc=loc, ip=ip), + _iv(base_bits, loc=loc, ip=ip), + _iv(index, loc=loc, ip=ip), + _iv(ok, loc=loc, ip=ip), + ], + "{\n.reg .pred pz, pn, pnan, pok;\n.reg .b16 hb;\n.reg .b32 bb, xb, b, h, l, m;\n" + ".reg .f32 bf, xf, v;\n" + "cvt.rn.bf16.f32 hb, $1;\ncvt.u32.u16 bb, hb;\nshl.b32 bb, bb, 16;\nmov.b32 bf, bb;\n" + "shl.b32 xb, $2, 16;\nmov.b32 xf, xb;\nadd.rn.f32 v, xf, bf;\nmov.b32 b, v;\n" + "setp.eq.b32 pz, b, 0x80000000;\nselp.b32 b, 0, b, pz;\n" + "setp.lt.s32 pn, b, 0;\nnot.b32 h, b;\nor.b32 m, b, 0x80000000;\nselp.b32 h, h, m, pn;\n" + "and.b32 m, b, 0x7fffffff;\nsetp.gt.u32 pnan, m, 0x7f800000;\n" + "selp.b32 h, 0xffffffff, h, pnan;\n" + "not.b32 l, $3;\nsetp.ne.s32 pok, $4, 0;\nselp.b32 h, h, 0, pok;\nselp.b32 l, l, 0, pok;\n" + "mov.b64 $0, {l, h};\n}", + "=l,f,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _max_u64(a, b, *, loc=None, ip=None): + return Int64( + _asm( + T.i64(), + [_lv(a, loc=loc, ip=ip), _lv(b, loc=loc, ip=ip)], + "max.u64 $0, $1, $2;", + "=l,l,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shuffle_max_u64(key, offset: int, *, loc=None, ip=None): + """``max(key, key of lane ^ offset)`` (unsigned); every lane of the warp participates.""" + return Int64( + _asm( + T.i64(), + [_lv(key, loc=loc, ip=ip)], + "{\n.reg .b32 lo, hi, olo, ohi;\n.reg .b64 o;\nmov.b64 {lo, hi}, $1;\n" + f"shfl.sync.bfly.b32 olo, lo, {offset}, 0x1f, 0xffffffff;\n" + f"shfl.sync.bfly.b32 ohi, hi, {offset}, 0x1f, 0xffffffff;\n" + "mov.b64 o, {olo, ohi};\nmax.u64 $0, $1, o;\n}", + "=l,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _pack_key(lo, hi, ok, *, loc=None, ip=None): + """``hi << 32 | lo`` if ``ok``, else 0.""" + return Int64( + _asm( + T.i64(), + [_iv(lo, loc=loc, ip=ip), _iv(hi, loc=loc, ip=ip), _iv(ok, loc=loc, ip=ip)], + "{\n.reg .pred p;\n.reg .b32 l, h;\nsetp.ne.s32 p, $3, 0;\n" + "selp.b32 l, $1, 0, p;\nselp.b32 h, $2, 0, p;\nmov.b64 $0, {l, h};\n}", + "=l,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _key_low(key, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_lv(key, loc=loc, ip=ip)], + "{\n.reg .b32 hi;\nmov.b64 {$0, hi}, $1;\n}", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_shared_b64(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "st.shared.b64 [$0], $1;", + "r,l,~{memory}", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_shared_b64(addr, *, loc=None, ip=None): + return Int64( + _asm( + T.i64(), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.b64 $0, [$1];", + "=l,r,~{memory}", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_shared_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r,~{memory}", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_shared_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.b32 $0, [$1];", + "=r,r,~{memory}", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _publish_key(addr, key, multicast: bool, *, loc=None, ip=None): + """Store the 8-byte ``key`` at ``addr``: the same offset of every rank's exchange through the + multicast mapping, or the local exchange.""" + store = ( + "multimem.st.relaxed.sys.global.v2.f32 [$0], {lo, hi};" + if multicast + else "st.global.v2.b32 [$0], {lo, hi};" + ) + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _lv(key, loc=loc, ip=ip)], + "{\n.reg .b32 lo, hi;\nmov.b64 {lo, hi}, $1;\n" + store + "\n}", + "l,l,~{memory}", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _chunk_ready(w0, w1, w2, w3, second, *, loc=None, ip=None): + """1 if the 16-byte chunk's first key and (where ``second``) its second key are published.""" + return Int32( + _asm( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3, second)], + "{\n.reg .pred p, q;\n" + "setp.ne.b32 p, $1, 0;\nsetp.ne.and.b32 p, $2, 0, p;\n" + "setp.ne.b32 q, $3, 0;\nsetp.ne.and.b32 q, $4, 0, q;\n" + "setp.eq.or.s32 q, $5, 0, q;\nand.pred p, p, q;\nselp.u32 $0, 1, 0, p;\n}", + "=r,r,r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +class DSparkMarkovGreedyKernel: + """One launch of the chain for ``width`` request slots (a power of two up to 16, the batch + padded), ``steps`` block positions, ``tp`` ranks of ``ctas`` CTAs and Markov rank + ``rank_dim``. See the module docstring.""" + + def __init__( + self, + width: int, + steps: int, + tp: int, + ctas: int, + rank_dim: int, + multicast: bool, + anchor64: bool, + use_pdl: bool, + staged: bool = False, + gate: bool = False, + ): + if width not in (1, 2, 4, 8, 16) or not 1 <= steps <= MAX_STEPS: + raise ValueError(f"unsupported width {width} / steps {steps}") + if rank_dim % 16 or rank_dim > 256: + raise ValueError(f"unsupported Markov rank {rank_dim}") + self.width = width + self.steps = steps + self.tp = tp + self.ctas = ctas + self.rank_dim = rank_dim + self.multicast = multicast + self.anchor64 = anchor64 + self.use_pdl = use_pdl + self.staged = staged + self.gate = gate + self.ntiles = (width + 7) // 8 + self.ksteps = rank_dim // 16 + # Shared-memory row pitch of the staged markov_w2 rows: 16 bytes of padding put the eight + # rows of an ldmatrix 8x8 tile in distinct banks. + self.pitch = rank_dim * 2 + 16 + self.row_chunks = rank_dim * 2 // 16 + self.entries = tp * ctas + self.keys = self.entries * width + self.chunks = (self.keys + 1) // 2 + self.per_thread = (self.chunks + THREADS - 1) // THREADS + if self.per_thread > 31: + raise ValueError(f"{self.keys} keys per position exceed the poll capacity") + self.set_bytes = set_bytes(self.entries) + + @cute.jit + def __call__( + self, + w2: Int64, + w1: Int64, + base: Int64, + stride_req: Int32, + stride_step: Int32, + anchors: Int64, + tokens: Int64, + stride_tok: Int32, + state: Int64, + box_uc: Int64, + box_mc: Int64, + num_reqs: Int32, + shard: Int32, + vocab: Int32, + rank: Int32, + stream: cuda.CUstream, + ): + @cute.struct + class SharedStorage: + stage: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, ROWS * self.pitch], 128] + # Per position: this CTA's key per (warp, request), every rank's per (warp, request), + # and the winning token per request. + best: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, WARPS * MAX_WIDTH], 16] + polled: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, WARPS * MAX_WIDTH], 16] + winners: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, MAX_WIDTH], 16] + # ``staged``: the requests' markov_w1 rows of the position. + w1s: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, MAX_WIDTH * self.pitch], 128] + + self.shared_storage = SharedStorage + self.kernel( + w2, + w1, + base, + stride_req, + stride_step, + anchors, + tokens, + stride_tok, + state, + box_uc, + box_mc, + num_reqs, + shard, + vocab, + rank, + ).launch( + grid=(self.ctas, 1, 1), + block=(THREADS, 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + w2: Int64, + w1: Int64, + base: Int64, + stride_req: Int32, + stride_step: Int32, + anchors: Int64, + tokens: Int64, + stride_tok: Int32, + state: Int64, + box_uc: Int64, + box_mc: Int64, + num_reqs: Int32, + shard: Int32, + vocab: Int32, + rank: Int32, + ): + width, steps = self.width, self.steps + tid, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + warp = tid // 32 + lane = tid % 32 + g = lane // 4 + q = lane % 4 + storage = utils.SmemAllocator().allocate(self.shared_storage) + stage = storage.stage.data_ptr().toint() + best = storage.best.data_ptr().toint() + polled = storage.polled.data_ptr().toint() + winners = storage.winners.data_ptr().toint() + w1s = storage.w1s.data_ptr().toint() + first = cta * ROWS + row_bytes = self.rank_dim * 2 + + # Static weights before the wait: this CTA's markov_w2 rows. + for it in cutlass.range_constexpr(ROWS * self.row_chunks // THREADS): + c = tid + it * THREADS + r = c // self.row_chunks + j = c % self.row_chunks + if first + r < shard: + _cp_async_evict_first( + stage + r * self.pitch + j * 16, + w2 + Int64(first + r) * row_bytes + Int64(j * 16), + ) + _cp_async_commit() + if cutlass.const_expr(self.staged): + # Rows of requests without a valid previous token stay zero. + for it in cutlass.range_constexpr( + (MAX_WIDTH * self.pitch // 16 + THREADS - 1) // THREADS + ): + c = tid + it * THREADS + if c < MAX_WIDTH * self.pitch // 16: + _st_shared_v4(w1s + c * 16, Int32(0), Int32(0), Int32(0), Int32(0)) + + # Launch ordinal and the bytes each set's last launch used. The previous launch has + # completed: it never triggers its dependents early. + st = _ld_volatile_v4(state) + launch = st[0] + cur_set = launch % NUM_SETS + clear_set = (launch + 2) % NUM_SETS + clear_bytes = st[1] + if clear_set == 1: + clear_bytes = st[2] + if clear_set == 2: + clear_bytes = st[3] + clear_base = box_uc + Int64(clear_set) * self.set_bytes + for v in cutlass.range( + cta * THREADS + tid, clear_bytes // 16, self.ctas * THREADS, unroll=1 + ): + _st_zero_v4(clear_base + Int64(v) * 16) + + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + + # Lane (g, q): B fragments of request 8 t + g, accumulators of requests 8 t + 2 q + {0, 1} + # at rows (g, g + 8) of its warp's tile. + # Each request's previous token lives in ``winners``: the anchor (-1 for padding requests), + # then each position's winner. + anchor_bytes = 8 if self.anchor64 else 4 + if tid < width: + ok = Int32(0) + if tid < num_reqs: + ok = Int32(1) + _st_shared_b32(winners + tid * 4, _ld_anchor(anchors + Int64(tid) * anchor_bytes, ok)) + rows = [first + warp * 16 + g, first + warp * 16 + g + 8] + row_ok = [] + for r in cutlass.range_constexpr(2): + ok = Int32(0) + if rows[r] < shard: + ok = Int32(1) + row_ok.append(ok) + load_ok = [] + for t in cutlass.range_constexpr(self.ntiles): + for r in cutlass.range_constexpr(2): + for n in cutlass.range_constexpr(2): + ok = Int32(0) + if 8 * t + 2 * q + n < num_reqs: + ok = row_ok[r] + load_ok.append(ok) + + base_cur = self.load_base(base, stride_req, stride_step, rows, q, load_ok, 0) + _cp_async_wait_all() + cute.arch.barrier() + + tile = ( + stage + (warp * 16 + lane % 8 + 8 * ((lane // 8) % 2)) * self.pitch + 16 * (lane // 16) + ) + index0 = rank * shard + own = rank * self.ctas + cta + pos_bytes = self.keys * 8 + set_base_uc = box_uc + Int64(cur_set) * self.set_bytes + set_base_mc = box_mc + Int64(cur_set) * self.set_bytes + for i in cutlass.range_constexpr(steps): + # Bias of the warp's 16 rows for every request tile. + accs = [[Float32(0.0)] * 4 for _ in range(self.ntiles)] + if cutlass.const_expr(self.staged): + # Warp r stages request r's markov_w1 row; B fragments by ldmatrix (matrix j of + # x4 load u: k-step 2 u + j // 2, half j % 2). + if warp < width: + p = _ld_shared_b32(winners + warp * 4) + valid = Int32(0) + if warp < num_reqs: + if p >= 0: + if p < vocab: + valid = Int32(1) + if lane < row_bytes // 16: + words = _ld_nc_v4_masked( + w1 + Int64(p * valid) * row_bytes + Int64(16 * lane), Int32(0) - valid + ) + _st_shared_v4( + w1s + warp * self.pitch + 16 * lane, + words[0], + words[1], + words[2], + words[3], + ) + cute.arch.barrier() + for u in cutlass.range_constexpr(self.ksteps // 2): + a0 = _ldsm_x4(tile + 64 * u) + a1 = _ldsm_x4(tile + 64 * u + 32) + for t in cutlass.range_constexpr(self.ntiles): + frag = _ldsm_x4( + w1s + (8 * t + lane % 8) * self.pitch + 16 * (lane // 8) + 64 * u + ) + accs[t] = list(_mma_bf16(a0, frag[0], frag[1], accs[t])) + accs[t] = list(_mma_bf16(a1, frag[2], frag[3], accs[t])) + else: + srcs = [] + masks = [] + for t in cutlass.range_constexpr(self.ntiles): + p = _ld_shared_b32(winners + ((8 * t + g) % width) * 4) + valid = Int32(0) + if 8 * t + g < num_reqs: + if p >= 0: + if p < vocab: + valid = Int32(1) + srcs.append(w1 + Int64(p * valid) * row_bytes + Int64(4 * q)) + masks.append(Int32(0) - valid) + for s in cutlass.range_constexpr(self.ksteps): + a = _ldsm_x4(tile + 32 * s) + for t in cutlass.range_constexpr(self.ntiles): + b0 = _ld_nc_b32(srcs[t] + Int64(32 * s)) & masks[t] + b1 = _ld_nc_b32(srcs[t] + Int64(32 * s + 16)) & masks[t] + accs[t] = list(_mma_bf16(a, b0, b1, accs[t])) + + # Keys of this CTA per request: lane max, then warp max, then CTA max. + for t in cutlass.range_constexpr(self.ntiles): + for n in cutlass.range_constexpr(2): + key = Int64(0) + for r in cutlass.range_constexpr(2): + k = _logit_key( + accs[t][2 * r + n], + base_cur[(t * 2 + r) * 2 + n], + index0 + rows[r], + row_ok[r], + ) + key = _max_u64(key, k) + for off in (4, 8, 16): + key = _shuffle_max_u64(key, off) + if g == 0: + _st_shared_b64(best + (warp * MAX_WIDTH + 8 * t + 2 * q + n) * 8, key) + cute.arch.barrier() + pos_off = Int64(i * pos_bytes) + if tid < width: + mine = _ld_shared_b64(best + tid * 8) + for w in cutlass.range_constexpr(1, WARPS): + mine = _max_u64(mine, _ld_shared_b64(best + (w * MAX_WIDTH + tid) * 8)) + _publish_key( + set_base_mc + pos_off + Int64((own * width + tid) * 8), mine, self.multicast + ) + + # The next position's base logits travel while the keys do. + base_nxt = base_cur + if cutlass.const_expr(i + 1 < steps): + base_nxt = self.load_base(base, stride_req, stride_step, rows, q, load_ok, i + 1) + + # Every rank's keys of this position. + region = set_base_uc + pos_off + if cutlass.const_expr(self.gate): + if tid < self.tp: + gc = ((tid + 1) * self.ctas - 1) * width // 2 + second = Int32(1) + if 2 * gc + 1 >= self.keys: + second = Int32(0) + gw = _ld_volatile_v4(region + Int64(gc) * 16) + while _chunk_ready(gw[0], gw[1], gw[2], gw[3], second) == 0: + gw = _ld_volatile_v4(region + Int64(gc) * 16) + cute.arch.barrier() + got = cute.make_rmem_tensor((self.per_thread * 4,), Int32) + pending = Int32(0) + for j in cutlass.range_constexpr(self.per_thread): + if tid + j * THREADS < self.chunks: + pending = pending | Int32(1 << j) + while pending != 0: + for j in cutlass.range_constexpr(self.per_thread): + # Threads past the end re-load the last chunk. + c = tid + j * THREADS + if c > self.chunks - 1: + c = Int32(self.chunks - 1) + w = _ld_volatile_v4(region + Int64(c) * 16) + for e in cutlass.range_constexpr(4): + got[j * 4 + e] = w[e] + second = Int32(1) + if 2 * c + 1 >= self.keys: + second = Int32(0) + ready = _chunk_ready(w[0], w[1], w[2], w[3], second) + pending = pending & (Int32(-1) ^ (Int32(1 << j) * ready)) + lo_key = Int64(0) + hi_key = Int64(0) + for j in cutlass.range_constexpr(self.per_thread): + c = tid + j * THREADS + ok0 = Int32(0) + if c < self.chunks: + ok0 = Int32(1) + ok1 = Int32(0) + if 2 * c + 1 < self.keys: + ok1 = Int32(1) + lo_key = _max_u64(lo_key, _pack_key(got[j * 4], got[j * 4 + 1], ok0)) + hi_key = _max_u64(hi_key, _pack_key(got[j * 4 + 2], got[j * 4 + 3], ok1)) + # Thread tid's chunks hold requests (2 tid) % width and (2 tid + 1) % width. + if cutlass.const_expr(width == 1): + lo_key = _max_u64(lo_key, hi_key) + for off in (1, 2, 4, 8, 16): + lo_key = _shuffle_max_u64(lo_key, off) + if lane == 0: + _st_shared_b64(polled + warp * MAX_WIDTH * 8, lo_key) + else: + for off in (1, 2, 4, 8, 16): + if cutlass.const_expr(off >= width // 2): + lo_key = _shuffle_max_u64(lo_key, off) + hi_key = _shuffle_max_u64(hi_key, off) + if lane < width // 2: + _st_shared_b64(polled + (warp * MAX_WIDTH + 2 * lane) * 8, lo_key) + _st_shared_b64(polled + (warp * MAX_WIDTH + 2 * lane + 1) * 8, hi_key) + cute.arch.barrier() + if tid < width: + final = _ld_shared_b64(polled + tid * 8) + for w in cutlass.range_constexpr(1, WARPS): + final = _max_u64(final, _ld_shared_b64(polled + (w * MAX_WIDTH + tid) * 8)) + token = Int32(-1) ^ _key_low(final) + _st_shared_b32(winners + tid * 4, token) + if cta == 0: + if tid < num_reqs: + _st_global_b32(tokens + (Int64(tid) * stride_tok + i) * 4, token) + cute.arch.barrier() + base_cur = base_nxt + + if cta == 0: + if tid == 0: + _st_global_b32(state, launch + 1) + _st_global_b32(state + Int64(4 + 4 * cur_set), Int32(steps * self.keys * 8)) + + def load_base(self, base, stride_req, stride_step, rows, q, load_ok, i: int): + """The lane's base logits of position ``i`` (BF16 bits), ``[tile][row][request]`` + (traced inline into the kernel).""" + vals = [] + for t in range(self.ntiles): + for r in range(2): + for n in range(2): + req = 8 * t + 2 * q + n + addr = ( + base + + (Int64(req * load_ok[(t * 2 + r) * 2 + n]) * stride_req) * 2 + + Int64(i * stride_step) * 2 + + Int64(rows[r]) * 2 + ) + vals.append(_ld_bf16_bits(addr, load_ok[(t * 2 + r) * 2 + n])) + return vals + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def _stream() -> cuda.CUstream: + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + +@functools.cache +def _compiled(width, steps, tp, ctas, rank_dim, multicast, anchor64, use_pdl, staged, gate): + kernel = DSparkMarkovGreedyKernel( + width, steps, tp, ctas, rank_dim, multicast, anchor64, use_pdl, staged, gate + ) + a = Int64(1 << 20) + i = Int32(1) + return cute.compile(kernel, a, a, a, i, i, a, a, i, a, a, a, i, i, i, i, _stream()) + + +def smem_bytes(rank_dim: int) -> int: + """Shared memory one CTA needs (an upper bound).""" + pitch = rank_dim * 2 + 16 + return (ROWS + MAX_WIDTH) * pitch + 2 * WARPS * MAX_WIDTH * 8 + MAX_WIDTH * 4 + 1024 + + +def run_dspark_markov_greedy( + local_block: torch.Tensor, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: torch.Tensor, + tokens: torch.Tensor, + *, + vocab: int, + tp: int, + rank: int, + state: torch.Tensor, + box_uc: int, + box_mc: int, + multicast: bool, + use_pdl: bool = True, + staged: Optional[bool] = None, + gate: Optional[bool] = None, +) -> None: + """Launch the chain on the current stream; writes ``tokens`` (int32 ``[B, K]``). + + ``local_block``: ``[B, K, shard]`` BF16 base logits with a unit vocab stride; + ``anchors``: ``[B]`` int32 / int64; ``markov_w1``: ``[vocab, R]`` BF16 contiguous; + ``w2_shard``: this rank's ``[shard, R]`` rows of ``markov_w2`` (contiguous); ``state``: the + rank's int32 ``[STATE_WORDS]`` (zero-initialized); ``box_uc`` / ``box_mc``: the local and + multicast addresses of the exchange (``NUM_SETS`` sets of :func:`set_bytes`, zero-initialized). + Every rank makes the same sequence of calls. ``staged`` / ``gate`` override :data:`OPTIONS`.""" + num_reqs, steps, shard = local_block.shape + rank_dim = markov_w1.shape[1] + width = width_for(num_reqs) + default_staged, default_gate = OPTIONS[width] + _compiled( + width, + steps, + tp, + ctas_for(shard), + rank_dim, + multicast, + anchors.dtype == torch.int64, + use_pdl, + default_staged if staged is None else staged, + default_gate if gate is None else gate, + )( + Int64(w2_shard.data_ptr()), + Int64(markov_w1.data_ptr()), + Int64(local_block.data_ptr()), + Int32(local_block.stride(0)), + Int32(local_block.stride(1)), + Int64(anchors.data_ptr()), + Int64(tokens.data_ptr()), + Int32(tokens.stride(0)), + Int64(state.data_ptr()), + Int64(box_uc), + Int64(box_mc), + Int32(num_reqs), + Int32(shard), + Int32(vocab), + Int32(rank), + _stream(), + ) + + +__all__ = [ + "MAX_STEPS", + "MAX_WIDTH", + "NUM_SETS", + "OPTIONS", + "ROWS", + "STATE_WORDS", + "DSparkMarkovGreedyKernel", + "ctas_for", + "run_dspark_markov_greedy", + "set_bytes", + "smem_bytes", + "width_for", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_mla_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_mla_rope.py new file mode 100644 index 000000000000..4152ecddf555 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/dspark_mla_rope.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Kimi K3 MLA drafter's fused RMSNorm + RoPE kernels (``TLLM_DSPARK_TAIL_V2``). + +Per row they compute what ``DSparkRMSNormRoPEKernel`` (``attention/kernels/blackwell/ +dspark_rmsnorm_rope.py``) computes, with every load of a row issued before its first store, so the +values are bitwise those of the launches they replace: + +* :class:`DSparkMlaCtxNormRoPEKernel`: the context latents of every drafter layer in one launch; +* :class:`DSparkMlaBlockRoPEKernel`: one layer's query RoPE and block-latent RMSNorm + RoPE, the + latents stored straight into the paged pool. + +Launchers: ``ops/spec/_dspark_mla_rope.py``. +""" + +import cutlass +import cutlass.cute as cute + +try: + from cuda.bindings import driver as cuda +except ImportError: + from cuda import cuda + + +class DSparkMlaCtxNormRoPEKernel: + """MLA context latents of every drafter layer in one launch: RMSNorm + RoPE. + + Row ``r`` of the output is layer ``r // N``, token ``r % N``: the + ``kv_lora_rank`` latent of ``x[token, layer]`` RMS-normalized and scaled by + ``weight[layer]``, then the ``rope_dim`` tail rotated at + ``freqs[positions[token]]`` (adjacent pairs, no norm, no weight). + + One warp per row, and per row the arithmetic of + ``DSparkRMSNormRoPEKernel(kv_lora_rank + rope_dim, rope_dim, 1, eps, True, + True, False, norm_dim=kv_lora_rank)``: lane ``l`` owns columns ``l + 32 i``, + accumulates their squares in ``i`` order, the warp reduction and ``rsqrt`` + are the same calls, and each output is the same chain of FP32 products. The + results are therefore bitwise those of the per-layer launches. What differs: + every load of a row is issued before its first store (the per-layer kernel + interleaves loads with stores the compiler cannot reorder, so each of its + 16 weight loads waits for the previous store), the layers share one launch, + and the RoPE table is indexed by position in the kernel instead of through a + gathered copy. + """ + + num_warps = 4 + num_threads = num_warps * cute.arch.WARP_SIZE + + def __init__(self, num_layers: int, kv_lora_rank: int, rope_dim: int, eps: float): + warp = cute.arch.WARP_SIZE + if kv_lora_rank % warp != 0 or (rope_dim // 2) % warp != 0 or rope_dim % 2 != 0: + raise ValueError( + f"kv_lora_rank ({kv_lora_rank}) and rope_dim / 2 ({rope_dim // 2}) must be " + f"multiples of {warp}" + ) + self.num_layers = num_layers + self.norm_dim = kv_lora_rank + self.head_dim = kv_lora_rank + rope_dim + self.eps = eps + self.norm_items = kv_lora_rank // warp + self.pair_items = rope_dim // 2 // warp + + @cute.jit + def __call__( + self, + x: cute.Tensor, + weight: cute.Tensor, + positions: cute.Tensor, + freqs: cute.Tensor, + output: cute.Tensor, + stream: cuda.CUstream, + ): + rows = x.shape[0] * self.num_layers + self.kernel(x, weight, positions, freqs, output).launch( + grid=[cute.ceil_div(rows, self.num_warps), 1, 1], + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x: cute.Tensor, + weight: cute.Tensor, + positions: cute.Tensor, + freqs: cute.Tensor, + output: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + lane = tidx % cute.arch.WARP_SIZE + row = bidx * self.num_warps + tidx // cute.arch.WARP_SIZE + num_tokens = x.shape[0] + if row < num_tokens * self.num_layers: + layer = row // num_tokens + token = row - layer * num_tokens + col0 = layer * self.head_dim + values = cute.make_rmem_tensor((self.norm_items,), cutlass.Float32) + scales = cute.make_rmem_tensor((self.norm_items,), cutlass.Float32) + reals = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + imags = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + coss = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + sins = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + position = positions[token] + for item in cutlass.range_constexpr(self.norm_items): + dim = lane + item * cute.arch.WARP_SIZE + values[item] = cutlass.Float32(x[token, col0 + dim]) + scales[item] = cutlass.Float32(weight[layer, dim]) + for item in cutlass.range_constexpr(self.pair_items): + pair = lane + item * cute.arch.WARP_SIZE + reals[item] = cutlass.Float32(x[token, col0 + self.norm_dim + pair * 2]) + imags[item] = cutlass.Float32(x[token, col0 + self.norm_dim + pair * 2 + 1]) + coss[item] = cutlass.Float32(freqs[position, pair, 0]) + sins[item] = cutlass.Float32(freqs[position, pair, 1]) + + sum_sq = cutlass.Float32(0.0) + for item in cutlass.range_constexpr(self.norm_items): + value = values[item] + sum_sq += value * value + sum_sq = cute.arch.warp_reduction_sum(sum_sq) + inverse_rms = cute.math.rsqrt(sum_sq / self.norm_dim + self.eps) + + for item in cutlass.range_constexpr(self.norm_items): + dim = lane + item * cute.arch.WARP_SIZE + value = values[item] * inverse_rms + value *= scales[item] + output[row, dim] = value.to(output.element_type) + for item in cutlass.range_constexpr(self.pair_items): + pair = lane + item * cute.arch.WARP_SIZE + real = reals[item] + imag = imags[item] + cos = coss[item] + sin = sins[item] + output[row, self.norm_dim + pair * 2] = (real * cos - imag * sin).to( + output.element_type + ) + output[row, self.norm_dim + pair * 2 + 1] = (imag * cos + real * sin).to( + output.element_type + ) + + +class DSparkMlaBlockRoPEKernel: + """One drafter layer's block RoPE work in one launch, all loads up front. + + Rows ``[0, T * H)`` rotate the ``rope_dim`` tail of the 192-wide query row + ``q[r]`` at ``freqs[positions[r // H]]`` into ``query[r, kv_lora_rank:]`` + (the rope half of the 576-wide absorbed query; the latent half is left to + the absorption GEMM). Rows ``[T * H, T * H + T)`` RMS-normalize the latent of + the block token ``t`` (``kv[t, :kv_lora_rank]``, scaled by ``weight``), + rotate its ``k_pe`` at ``freqs[positions[t]]`` and store the 576-wide row + straight into the paged pool at ``pool[pages[t], offsets[t]]``. + + Per row the arithmetic is that of the separate launches it replaces -- + ``cute_dsl_dspark_rope_into`` for the query rows and + ``cute_dsl_dspark_rmsnorm_rope(norm_dim=kv_lora_rank)`` followed by the + pool scatter for the latent rows -- so the values written are bitwise the + same. + """ + + num_warps = 4 + num_threads = num_warps * cute.arch.WARP_SIZE + + def __init__( + self, num_heads: int, qk_nope_dim: int, kv_lora_rank: int, rope_dim: int, eps: float + ): + warp = cute.arch.WARP_SIZE + if kv_lora_rank % warp != 0 or (rope_dim // 2) % warp != 0 or rope_dim % 2 != 0: + raise ValueError( + f"kv_lora_rank ({kv_lora_rank}) and rope_dim / 2 ({rope_dim // 2}) must be " + f"multiples of {warp}" + ) + self.num_heads = num_heads + self.qk_nope_dim = qk_nope_dim + self.norm_dim = kv_lora_rank + self.eps = eps + self.norm_items = kv_lora_rank // warp + self.pair_items = rope_dim // 2 // warp + + @cute.jit + def __call__( + self, + q: cute.Tensor, + query: cute.Tensor, + kv: cute.Tensor, + weight: cute.Tensor, + positions: cute.Tensor, + freqs: cute.Tensor, + pool: cute.Tensor, + pages: cute.Tensor, + offsets: cute.Tensor, + stream: cuda.CUstream, + ): + rows = q.shape[0] + kv.shape[0] + self.kernel(q, query, kv, weight, positions, freqs, pool, pages, offsets).launch( + grid=[cute.ceil_div(rows, self.num_warps), 1, 1], + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + q: cute.Tensor, + query: cute.Tensor, + kv: cute.Tensor, + weight: cute.Tensor, + positions: cute.Tensor, + freqs: cute.Tensor, + pool: cute.Tensor, + pages: cute.Tensor, + offsets: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + lane = tidx % cute.arch.WARP_SIZE + row = bidx * self.num_warps + tidx // cute.arch.WARP_SIZE + num_q_rows = q.shape[0] + reals = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + imags = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + coss = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + sins = cute.make_rmem_tensor((self.pair_items,), cutlass.Float32) + if row < num_q_rows: + position = positions[row // self.num_heads] + for item in cutlass.range_constexpr(self.pair_items): + pair = lane + item * cute.arch.WARP_SIZE + reals[item] = cutlass.Float32(q[row, self.qk_nope_dim + pair * 2]) + imags[item] = cutlass.Float32(q[row, self.qk_nope_dim + pair * 2 + 1]) + coss[item] = cutlass.Float32(freqs[position, pair, 0]) + sins[item] = cutlass.Float32(freqs[position, pair, 1]) + for item in cutlass.range_constexpr(self.pair_items): + pair = lane + item * cute.arch.WARP_SIZE + real = reals[item] + imag = imags[item] + cos = coss[item] + sin = sins[item] + query[row, self.norm_dim + pair * 2] = (real * cos - imag * sin).to( + query.element_type + ) + query[row, self.norm_dim + pair * 2 + 1] = (imag * cos + real * sin).to( + query.element_type + ) + elif row < num_q_rows + kv.shape[0]: + token = row - num_q_rows + values = cute.make_rmem_tensor((self.norm_items,), cutlass.Float32) + scales = cute.make_rmem_tensor((self.norm_items,), cutlass.Float32) + position = positions[token] + page = pages[token] + offset = offsets[token] + for item in cutlass.range_constexpr(self.norm_items): + dim = lane + item * cute.arch.WARP_SIZE + values[item] = cutlass.Float32(kv[token, dim]) + scales[item] = cutlass.Float32(weight[dim]) + for item in cutlass.range_constexpr(self.pair_items): + pair = lane + item * cute.arch.WARP_SIZE + reals[item] = cutlass.Float32(kv[token, self.norm_dim + pair * 2]) + imags[item] = cutlass.Float32(kv[token, self.norm_dim + pair * 2 + 1]) + coss[item] = cutlass.Float32(freqs[position, pair, 0]) + sins[item] = cutlass.Float32(freqs[position, pair, 1]) + + sum_sq = cutlass.Float32(0.0) + for item in cutlass.range_constexpr(self.norm_items): + value = values[item] + sum_sq += value * value + sum_sq = cute.arch.warp_reduction_sum(sum_sq) + inverse_rms = cute.math.rsqrt(sum_sq / self.norm_dim + self.eps) + + for item in cutlass.range_constexpr(self.norm_items): + dim = lane + item * cute.arch.WARP_SIZE + value = values[item] * inverse_rms + value *= scales[item] + pool[page, offset, dim] = value.to(pool.element_type) + for item in cutlass.range_constexpr(self.pair_items): + pair = lane + item * cute.arch.WARP_SIZE + real = reals[item] + imag = imags[item] + cos = coss[item] + sin = sins[item] + pool[page, offset, self.norm_dim + pair * 2] = (real * cos - imag * sin).to( + pool.element_type + ) + pool[page, offset, self.norm_dim + pair * 2 + 1] = (imag * cos + real * sin).to( + pool.element_type + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/embedding_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/embedding_gather.py new file mode 100644 index 000000000000..753aa64bedbe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/embedding_gather.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The token-embedding gather as a vectorized row copy (``features.embedding_gather``). + +``F.embedding(ids, weight)`` for 1-D ids and a contiguous table: one program per (token, 1024 +hidden columns) copies the row with 16-byte accesses instead of the scalar-load index-select +kernel. A copy, so bitwise ``nn.Embedding``; the launch shape depends only on the token count, +so it is CUDA-graph capturable. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _embedding_gather_kernel( + ids_ptr, weight_ptr, out_ptr, vocab, hidden, BLOCK_H: tl.constexpr, VEC: tl.constexpr +): + row = tl.program_id(0) + offs = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + offs = tl.max_contiguous(tl.multiple_of(offs, VEC), VEC) + mask = offs < hidden + token = tl.load(ids_ptr + row).to(tl.int64) + valid = (token >= 0) & (token < vocab) + src = weight_ptr + tl.where(valid, token, 0) * hidden + vals = tl.load(src + offs, mask=mask & valid, other=0.0) + tl.store(out_ptr + row.to(tl.int64) * hidden + offs, vals, mask=mask) + + +def embedding_gather(ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``F.embedding(ids, weight)`` for 1-D ids and a contiguous table, as a + vectorized row copy (ids outside the table give a zero row).""" + assert ids.dim() == 1 and weight.dim() == 2 and weight.is_contiguous() + vocab, hidden = weight.shape + out = torch.empty((ids.shape[0], hidden), dtype=weight.dtype, device=weight.device) + if ids.shape[0] == 0: + return out + vec = 16 // weight.element_size() + assert hidden % vec == 0 + block_h = 1024 + _embedding_gather_kernel[(ids.shape[0], triton.cdiv(hidden, block_h))]( + ids, weight, out, vocab, hidden, BLOCK_H=block_h, VEC=vec, num_warps=4 + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_context_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_context_gather.py new file mode 100644 index 000000000000..0802a1c039d9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_context_gather.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 FP4 MLA cached-prefix gather with MXFP8 activation and unit-scale FP8 RoPE outputs. + +:func:`gather_mxfp8` reads one cached-prefix partition of dense FP4 MLA V2 storage (page 128, +320-byte rows: 512 latent + 64 RoPE + 64 RoPE-residual FP4 values, 40 swizzled E4M3 scales per +row) and writes, per token, the BF16-rounded latent quantized to MXFP8 (E4M3 ``[M, 512]`` with +K32 R128c4 UE8M0 scales, 1x128 quantization, padded to 128 rows) for the FP8-input kv_b_proj +GEMM, and the BF16-rounded RoPE part plus its residual as unit-scale E4M3 ``[M, 64]``. The bytes +are those of the dense BF16 gather followed by ``fp8_quantize_1x128_packed_ue8m0`` and +``static_quantize_e4m3_per_tensor`` with a unit scale. +""" + +from collections.abc import Callable + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import llvm +from cutlass.base_dsl.dsl import BaseDSL +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + + +def _compile_cutedsl(*args: object, **kwargs: object) -> Callable[..., object]: + """Compile with the PyIR frontend the FP4 MLA kernels require.""" + with BaseDSL.enable_pyir(): + return cute.compile(*args, **kwargs) + + +def _current_cu_stream() -> cuda.CUstream: + """The caller's stream, for launch ordering and CUDA graph capture.""" + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + +@dsl_user_op +def _u32_as_f32(value, *, loc=None, ip=None): + return cutlass.Float32(llvm.bitcast(T.f32(), cutlass.Uint32(value).ir_value(loc=loc, ip=ip))) + + +@dsl_user_op +def _round_up_ue8m0(value, *, loc=None, ip=None): + """Convert positive FP32 to UE8M0 with round-toward-positive-infinity.""" + packed = llvm.inline_asm( + T.i16(), + [cutlass.Float32(value).ir_value(loc=loc, ip=ip)], + "cvt.rp.satfinite.ue8m0x2.f32 $0, 0f00000000, $1;", + "=h,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Uint8( + llvm.inline_asm( + T.i32(), + [packed], + "mov.b32 $0, { $1, 0 };", + "=r,h", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _rcp(value, *, loc=None, ip=None): + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Float32(value).ir_value(loc=loc, ip=ip)], + "rcp.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _div_full(value, denominator, *, loc=None, ip=None): + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + cutlass.Float32(value).ir_value(loc=loc, ip=ip), + cutlass.Float32(denominator).ir_value(loc=loc, ip=ip), + ], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _sf_offset(row, col, columns: cutlass.Constexpr): + return ( + col % 4 + + col // 4 * 512 + + row % 32 * 16 + + row % 128 // 32 * 4 + + row // 128 * (128 * columns) + ) + + +@cute.jit +def _decode_fp4(code): + mag = code & 7 + bits = ((mag // 2 + 126) << 23) | ((mag & 1) << 22) + if mag < 2: + bits = mag * 0x3F000000 + bits = bits.to(cutlass.Uint32) | ((code.to(cutlass.Uint32) & 8) << 28) + return _u32_as_f32(bits) + + +@cute.jit +def _fp8(value): + if value > 448.0: + value = cutlass.Float32(448.0) + if value < -448.0: + value = cutlass.Float32(-448.0) + return value.to(cutlass.Float8E4M3FN) + + +@cute.kernel +def _gather( + a_ptr: cute.Pointer, + asc_ptr: cute.Pointer, + pe_ptr: cute.Pointer, + kv_ptr: cute.Pointer, + sf_ptr: cute.Pointer, + pages_ptr: cute.Pointer, + cu_ptr: cute.Pointer, + offsets_ptr: cute.Pointer, + global_ptr: cute.Pointer, + m: cutlass.Int64, + max_chunk: cutlass.Int64, + table_stride: cutlass.Int64, + num_pages: cutlass.Int64, + page_size: cutlass.Constexpr, + kv_s0: cutlass.Int64, + kv_s2: cutlass.Int64, + kv_s4: cutlass.Constexpr, + sf_s0: cutlass.Int64, + warps: cutlass.Constexpr, +): + tid, _, _ = cute.arch.thread_idx() + block, batch, _ = cute.arch.block_idx() + lane = tid % 32 + blocks = (max_chunk + warps - 1) // warps + # Flat pointer views retain arbitrary layer/page strides without copying. + huge = cute.make_layout((cutlass.Int64(1 << 40),)) + asc = cute.make_tensor(asc_ptr, huge) + pages = cute.make_tensor(pages_ptr, huge) + cu = cute.make_tensor(cu_ptr, huge) + offsets = cute.make_tensor(offsets_ptr, huge) + global_scale = cute.make_tensor(global_ptr, cute.make_layout((1,))) + if block == blocks: + if batch == 0: + padded_m = (m + 127) // 128 * 128 + for i in range(tid, (padded_m - m) * 16, warps * 32): + asc[_sf_offset(m + i // 16, i % 16, 16)] = cutlass.Uint8(0) + else: + local = cutlass.Int64(block * warps + tid // 32) + begin, end = cu[batch], cu[batch + 1] + position = offsets[batch] + local + logical_page = position // page_size + valid = ( + (local < end - begin) + & (local < max_chunk) + & (position >= 0) + & (logical_page < table_stride) + ) + if valid: + physical = cutlass.Int64(pages[batch * table_stride + logical_page]) + if (physical >= 0) & (physical < num_pages): + row = position % page_size + outrow = begin + local + base = physical * kv_s0 + row * kv_s2 + # Keep large cache offsets in 64-bit pointer arithmetic. Scalar + # tensor indices must remain within one row/page, below 2 GiB. + kv_row = cute.make_tensor(kv_ptr + base, cute.make_layout((320,))) + sf_page = cute.make_tensor( + sf_ptr + physical * sf_s0, cute.make_layout((page_size * 40,)) + ) + global_value = global_scale[0] + scale = sf_page[_sf_offset(row, lane, 40)].to(cutlass.Float32) + input_bytes = cute.make_rmem_tensor((8,), cutlass.Uint8) + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Uint8, num_bits_per_copy=64 + ) + packed_input = cute.make_tensor( + kv_ptr + cute.assume(base + lane * 8, 8), cute.make_layout((8,)) + ) + cute.copy(load_atom, packed_input, input_bytes) + values = cute.make_rmem_tensor((16,), cutlass.Float32) + amax = cutlass.Float32(1.0e-10) + for i in cutlass.range_constexpr(16): + packed = input_bytes[i // 2].to(cutlass.Int32) + code = (packed >> ((i % 2) * 4)) & 15 + value = ( + _div_full(_decode_fp4(code) * scale, global_value) + .to(cutlass.BFloat16) + .to(cutlass.Float32) + ) + values[i] = value + magnitude = value + if value < 0.0: + magnitude = -value + amax = cute.arch.fmax(amax, magnitude) + for delta in cutlass.range_constexpr(3): + other = cute.arch.shuffle_sync_bfly(amax, 4 >> delta) + amax = cute.arch.fmax(amax, other) + quant_sf = _round_up_ue8m0(amax * _rcp(cutlass.Float32(448.0))) + inv_scale = _rcp(_u32_as_f32(quant_sf.to(cutlass.Uint32) << 23)) + packed_out = cute.make_rmem_tensor((16,), cutlass.Float8E4M3FN) + for i in cutlass.range_constexpr(16): + packed_out[i] = _fp8(values[i] * inv_scale) + output = cute.make_tensor( + a_ptr + cute.assume(outrow * 512 + lane * 16, 16), cute.make_layout((16,)) + ) + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float8E4M3FN, num_bits_per_copy=128 + ) + cute.copy(copy_atom, packed_out, output) + if lane % 8 == 0: + dst = _sf_offset(outrow, lane // 8 * 4, 16) + for i in cutlass.range_constexpr(4): + asc[dst + i] = quant_sf + main_scale = sf_page[_sf_offset(row, 32 + lane // 8, 40)].to(cutlass.Float32) + residual_scale = sf_page[_sf_offset(row, 36 + lane // 8, 40)].to(cutlass.Float32) + main_byte = kv_row[256 + lane].to(cutlass.Int32) + residual_byte = kv_row[288 + lane].to(cutlass.Int32) + rope_out = cute.make_rmem_tensor((2,), cutlass.Float8E4M3FN) + for i in cutlass.range_constexpr(2): + main_code = (main_byte >> (i * 4)) & 15 + residual_code = (residual_byte >> (i * 4)) & 15 + main = _div_full(_decode_fp4(main_code) * main_scale, global_value) + residual = _div_full(_decode_fp4(residual_code) * residual_scale, global_value) + value = (main + residual).to(cutlass.BFloat16).to(cutlass.Float32) + rope_out[i] = _fp8(value) + rope = cute.make_tensor( + pe_ptr + cute.assume(outrow * 64 + lane * 2, 2), cute.make_layout((2,)) + ) + rope_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float8E4M3FN, num_bits_per_copy=16 + ) + cute.copy(rope_copy, rope_out, rope) + + +@cute.jit +def _launch( + a: cute.Pointer, + asc: cute.Pointer, + pe: cute.Pointer, + kv: cute.Pointer, + sf: cute.Pointer, + pages: cute.Pointer, + cu: cute.Pointer, + offsets: cute.Pointer, + global_scale: cute.Pointer, + m: cutlass.Int64, + max_chunk: cutlass.Int64, + table_stride: cutlass.Int64, + num_pages: cutlass.Int64, + page_size: cutlass.Constexpr, + num_contexts: cutlass.Int64, + kv_s0: cutlass.Int64, + kv_s2: cutlass.Int64, + kv_s4: cutlass.Constexpr, + sf_s0: cutlass.Int64, + warps: cutlass.Constexpr, + stream: cuda.CUstream, +): + _gather( + a, + asc, + pe, + kv, + sf, + pages, + cu, + offsets, + global_scale, + m, + max_chunk, + table_stride, + num_pages, + page_size, + kv_s0, + kv_s2, + kv_s4, + sf_s0, + warps, + ).launch( + grid=((max_chunk + warps - 1) // warps + 1, num_contexts, 1), + block=(warps * 32, 1, 1), + stream=stream, + ) + + +_COMPILED = {} + + +def gather_mxfp8( + a, + asc, + pe, + kv, + sf, + pages, + cu, + offsets, + global_scale, + max_chunk, + page_size=128, + warps=4, + *, + page_table_stride=None, + num_contexts=None, +): + """Gather canonical page-128 FP4 storage; prefix lengths and page strides are dynamic.""" + if ( + page_size != 128 + or kv.stride(4) != 1 + or kv.data_ptr() % 8 + or kv.stride(0) % 8 + or kv.stride(2) % 8 + ): + raise ValueError( + "MXFP8 gather requires canonical page-128 FP4 storage with aligned contiguous packed rows." + ) + if warps not in (4, 8): + raise ValueError("gather warps must be 4 or 8") + if num_contexts is None: + num_contexts = offsets.numel() + if page_table_stride is None: + page_table_stride = pages.numel() // num_contexts + tensors = (a, asc, pe, kv, sf, pages, cu, offsets, global_scale) + dtypes = ( + cutlass.Float8E4M3FN, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Int32, + cutlass.Int64, + cutlass.Int64, + cutlass.Float32, + ) + alignments = (16, 16, 16, 8, 1, 4, 8, 8, 4) + pointers = tuple( + make_ptr(dtype, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=alignment) + for tensor, dtype, alignment in zip(tensors, dtypes, alignments) + ) + args = ( + *pointers, + a.shape[0], + max_chunk, + page_table_stride, + kv.shape[0], + page_size, + num_contexts, + kv.stride(0), + kv.stride(2), + kv.stride(4), + sf.stride(0), + warps, + _current_cu_stream(), + ) + key = (a.device.index, warps) + if key not in _COMPILED: + _COMPILED[key] = _compile_cutedsl(_launch, *args, options="--opt-level 2") + _COMPILED[key](*args[:13], *args[14:17], args[18], args[-1]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_kv_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_kv_append.py new file mode 100644 index 000000000000..6c90538d32a0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_kv_append.py @@ -0,0 +1,2046 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KV-only FP4 MLA generation cache update of K3's fused query. + +:func:`fp4_mla_kv_append_generation` is the FP4 MLA backend's fused generation cache update +(RoPE on the new K tail, FP4 tiles, HP ring, V scales and packed V) without the query rows: it +prepares the packed-Q buffers of a query of ``num_q_heads`` heads on ``metadata.fp4_mla_state`` +and leaves their rows to the caller (``k3_mla_q_fused`` writes them on another stream). The Triton +kernel is the backend's ``_fp4_mla_generation_fused_qk_rope_cache_update_kernel`` with a +``WRITE_Q_ROWS`` switch; launched with it off and a grid of the KV programs only, it writes the +cache bytes of the full update. +""" + +from typing import Any, Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.fp4_mla.cache_update import _prepare_fp4_mla_q_buffers +from tensorrt_llm._torch.attention.backends.fp4_mla.config import ( + _FP4_MLA_CUTEDSL_BACKEND, + _FP4_MLA_K_RESIDUAL_BACKENDS, + FP4_BLOCK_SIZE, + FP4_MLA_K_RESIDUAL_DIM, + FP4_MLA_Q1_PREFIX_BLOCK_DIM, + FP4_MLA_Q_PREFIX_BLOCK_DIM, + FP4_MLA_Q_PREFIX_DIM, + FP4_MLA_Q_SF_GROUPS, + FP4_MLA_SCALE_ROW_GROUP, + _fp4_mla_attention_backend, + _fp4_mla_cutedsl_fused_v_transpose_enabled, + _fp4_mla_packed_rows_per_seq, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.fp4_mla_kernels import ( + _fp4_mla_generation_q1_kv_tail8_update, + _fp4_mla_generation_q1_kv_tile_update, + _fp4_mla_q_sf_offset, + _fp4_mla_q_sf_offset_h128, + _fp4_mla_swizzled_sf_offset, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.layout import ( + _get_fp4_mla_global_scale, + _get_fp4_mla_hp_pool_layout, + _get_fp4_mla_kv_cache_tensors, + _get_fp4_mla_q_global_scale, + _validate_fp4_mla_cache_shape, + _validate_fp4_mla_hp_generation_width, + _validate_fp4_mla_kv_storage_shape, + get_fp4_mla_v_scale_pool_view, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.metadata import ( + _fp4_mla_generation_hp_page_ids, + _fp4_mla_generation_num_blocks_device, + _fp4_mla_generation_page_ids, + _fp4_mla_uniform_generation_lengths, + _materialize_fp4_mla_device_page_table_for_forward, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.packed_plan import ( + _fp4_mla_helix_plan_bounds, + _fp4_mla_packed_plan_geometry, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.v_cache import ( + _get_cutedsl_persistent_v_packed_cache, + _get_fp4_mla_v_packed_pool_base, + _maybe_update_triton_v_packed_cache, + _repack_cutedsl_v_packed_cache, +) + +# --------------------------------------------------------------------------- +# Launch policy and element-wise FP4 helpers of the FP4 MLA generation update (the backend's +# ``fp4_mla.config`` / ``fp4_mla.fp4_mla_kernels``), used by this module's kernel. +# --------------------------------------------------------------------------- + +_FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH = 4 +_FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH = 16 +_FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH = 32 +_FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD = 256 +_FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD = 512 +_FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD = 640 +_FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD = 768 + + +def _ceil_div(lhs: int, rhs: int) -> int: + return (lhs + rhs - 1) // rhs + + +def _fp4_mla_q1_kv_blocks_per_program(num_gen: int, v_head_dim: int) -> int: + """Select the Q1 KV work per program without changing the launch boundary.""" + large_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH * FP4_BLOCK_SIZE + if num_gen >= _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD and v_head_dim % large_batch_block_dim == 0: + return _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH + medium_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH * FP4_BLOCK_SIZE + if ( + num_gen >= _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD + and v_head_dim % medium_batch_block_dim == 0 + ): + return _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH + small_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH * FP4_BLOCK_SIZE + if v_head_dim % small_batch_block_dim == 0: + return _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH + return 1 + + +def _fp4_mla_q1_prefix_blocks_per_program( + num_gen: int, + q1_kv_blocks_per_program: int, +) -> int: + """Select rolled Q-prefix work only when the batch keeps it efficient.""" + max_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM + if ( + num_gen >= _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD + and q1_kv_blocks_per_program == _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH + ): + if num_gen >= _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD: + return min(4, max_prefix_blocks) + return min(2, max_prefix_blocks) + return 1 + + +@triton.jit +def _fp4_mla_pack_e2m1x2_rne(even, odd): + """Match cvt.rn.satfinite.e2m1x2.f32 used by arcquantFP4.cu.""" + return tl.inline_asm_elementwise( + """ + { + .reg .b8 packed0, packed1, packed2, packed3; + cvt.rn.satfinite.e2m1x2.f32 packed0, $5, $1; + cvt.rn.satfinite.e2m1x2.f32 packed1, $6, $2; + cvt.rn.satfinite.e2m1x2.f32 packed2, $7, $3; + cvt.rn.satfinite.e2m1x2.f32 packed3, $8, $4; + mov.b32 $0, {packed0, packed1, packed2, packed3}; + } + """, + constraints="=r,f,f,f,f,f,f,f,f", + args=[even.to(tl.float32), odd.to(tl.float32)], + dtype=tl.uint8, + is_pure=True, + pack=4, + ) + + +@triton.jit +def _fp4_mla_pack_e2m1x8_rne(even0, even1, even2, even3, odd0, odd1, odd2, odd3): + """Pack four consecutive E2M1x2 bytes into one word.""" + return tl.inline_asm_elementwise( + """ + { + .reg .b8 packed0, packed1, packed2, packed3; + cvt.rn.satfinite.e2m1x2.f32 packed0, $5, $1; + cvt.rn.satfinite.e2m1x2.f32 packed1, $6, $2; + cvt.rn.satfinite.e2m1x2.f32 packed2, $7, $3; + cvt.rn.satfinite.e2m1x2.f32 packed3, $8, $4; + mov.b32 $0, {packed0, packed1, packed2, packed3}; + } + """, + constraints="=r,f,f,f,f,f,f,f,f", + args=[even0, even1, even2, even3, odd0, odd1, odd2, odd3], + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fp4_mla_unpack_e2m1x2(packed): + """Decode packed E2M1 through the native GR100 conversion path.""" + return tl.inline_asm_elementwise( + """ + { + .reg .b8 packed0, packed1, packed2, packed3; + .reg .b16 low, high; + .reg .b32 half2; + mov.b32 {packed0, packed1, packed2, packed3}, $2; + cvt.rn.f16x2.e2m1x2 half2, packed0; + mov.b32 {low, high}, half2; + cvt.f32.f16 $0, low; + cvt.f32.f16 $1, high; + } + """, + constraints="=f,=f,r", + args=[packed.to(tl.uint32)], + dtype=(tl.float32, tl.float32), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fp4_mla_rcp_approx_ftz(value): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", + constraints="=f,f", + args=[value], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fp4_mla_rope_fp32(even, odd, cos, sin): + return tl.inline_asm_elementwise( + """ + { + .reg .f32 even_cos, odd_sin, odd_cos, even_sin; + mul.rn.f32 even_cos, $2, $4; + mul.rn.f32 odd_sin, $3, $5; + sub.rn.f32 $0, even_cos, odd_sin; + mul.rn.f32 odd_cos, $3, $4; + mul.rn.f32 even_sin, $2, $5; + add.rn.f32 $1, odd_cos, even_sin; + } + """, + constraints="=f,=f,f,f,f,f", + args=[even, odd, cos, sin], + dtype=(tl.float32, tl.float32), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fp4_mla_quantize_q_groups(values_even, values_odd): + amax = tl.max( + tl.maximum(tl.abs(values_even), tl.abs(values_odd)), + axis=2, + ) + scale = tl.maximum(amax / 6.0, 0.001953125) + scale_fp8 = scale.to(tl.float8e4nv) + inv_scale = _fp4_mla_rcp_approx_ftz(scale_fp8.to(tl.float32)) + packed = _fp4_mla_pack_e2m1x2_rne( + values_even * inv_scale[:, :, None], + values_odd * inv_scale[:, :, None], + ) + return packed, scale_fp8 + + +@triton.jit +def _fp4_mla_quantize_q_groups_words( + values_even, + values_odd, + BLOCK_Q_HEADS: tl.constexpr, + GROUPS: tl.constexpr, +): + """Quantize FP4 groups while retaining native packed-word ownership.""" + tl.static_assert(values_even.shape[2] == 8) + tl.static_assert(values_odd.shape[2] == 8) + amax = tl.max( + tl.maximum(tl.abs(values_even), tl.abs(values_odd)), + axis=2, + ) + scale = tl.maximum(amax / 6.0, 0.001953125) + scale_fp8 = scale.to(tl.float8e4nv) + inv_scale = _fp4_mla_rcp_approx_ftz(scale_fp8.to(tl.float32)) + scaled_even = tl.reshape( + values_even * inv_scale[:, :, None], + (BLOCK_Q_HEADS, GROUPS, 2, 2, 2), + ) + scaled_odd = tl.reshape( + values_odd * inv_scale[:, :, None], + (BLOCK_Q_HEADS, GROUPS, 2, 2, 2), + ) + even_02, even_13 = tl.split(scaled_even) + odd_02, odd_13 = tl.split(scaled_odd) + even0, even2 = tl.split(even_02) + even1, even3 = tl.split(even_13) + odd0, odd2 = tl.split(odd_02) + odd1, odd3 = tl.split(odd_13) + packed_words = _fp4_mla_pack_e2m1x8_rne( + even0, + even1, + even2, + even3, + odd0, + odd1, + odd2, + odd3, + ) + return packed_words, scale_fp8 + + +@triton.jit +def _fp4_mla_quantize_q_residual_groups(values_even, values_odd): + packed, scale_fp8 = _fp4_mla_quantize_q_groups(values_even, values_odd) + rounded_scale = scale_fp8.to(tl.float32) + main_even, main_odd = _fp4_mla_unpack_e2m1x2(packed) + residual_even = values_even - main_even * rounded_scale[:, :, None] + residual_odd = values_odd - main_odd * rounded_scale[:, :, None] + residual_packed, residual_scale_fp8 = _fp4_mla_quantize_q_groups(residual_even, residual_odd) + return packed, scale_fp8, residual_packed, residual_scale_fp8 + + +@triton.jit +def _fp4_e2m1_to_f32(nibble): + magnitude = nibble & 0x7 + value = tl.where( + magnitude == 0, + 0.0, + tl.where( + magnitude == 1, + 0.5, + tl.where( + magnitude == 2, + 1.0, + tl.where( + magnitude == 3, + 1.5, + tl.where( + magnitude == 4, + 2.0, + tl.where(magnitude == 5, 3.0, tl.where(magnitude == 6, 4.0, 6.0)), + ), + ), + ), + ), + ) + sign = (nibble & 0x8) != 0 + return tl.where(sign, -value, value) + + +@triton.jit +def _fp4_mla_floor_half_up(abs_value, multiplier, bias): + return tl.inline_asm_elementwise( + """ + { + .reg .f32 rounded; + fma.rz.f32 rounded, $1, $2, $3; + cvt.rzi.s32.f32 $0, rounded; + } + """, + constraints="=r,f,f,f", + args=[abs_value, multiplier, bias], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fp4_e2m1_quantize(x): + abs_x = tl.abs(x) + bounded = tl.where(abs_x < 6.0, abs_x, 6.0) + small = bounded < 2.0 + multiplier = tl.where(small, 2.0, 1.0) + bias = tl.where(small, 0.5, 2.5) + magnitude = _fp4_mla_floor_half_up(bounded, multiplier, bias) + magnitude = tl.minimum(magnitude, 6) + (~(abs_x < 5.0)).to(tl.int32) + sign = tl.where(x < 0.0, 8, 0) + return (magnitude | sign).to(tl.uint8) + + +def fp4_mla_kv_append_generation( + mqa: Any, + latent_cache: torch.Tensor, + metadata: Any, + num_q_heads: int, +) -> None: + """Append this step's latent rows of the FP4 MLA layer behind ``mqa`` (its + ``TrtllmAttention``) to the FP4 cache and HP ring, and prepare the packed-Q buffers of + ``num_q_heads`` heads on ``metadata`` for the decode, without writing the query rows.""" + assert mqa.kv_lora_rank is not None + assert mqa.qk_rope_head_dim is not None + if not mqa.uses_fp4_mla_attention: + raise RuntimeError("The FP4 MLA KV append requires FP4 MLA attention.") + mqa._ensure_rope_table_size(metadata.max_seq_len) + if latent_cache.shape[-1] != mqa.kv_lora_rank + mqa.qk_rope_head_dim: + raise RuntimeError( + "FP4 MLA latent_cache last dimension must be " + f"{mqa.kv_lora_rank + mqa.qk_rope_head_dim}, got {latent_cache.shape[-1]}." + ) + if ( + mqa.rotary_cos_sin is None + or latent_cache.dtype != torch.bfloat16 + or mqa.rotary_cos_sin.dtype != torch.float32 + ): + raise RuntimeError( + "FP4 MLA generation requires a BF16 latent cache and an FP32 rotary table." + ) + hp_pool_updated = _scatter_generation_kv_only( + metadata, + latent_cache, + mqa.layer_idx, + token_offset=getattr(metadata, "num_ctx_tokens", 0), + local_layer=mqa.get_fp4_mla_local_layer_idx(metadata), + v_head_dim=mqa.kv_lora_rank, + rotary_cos_sin=mqa.rotary_cos_sin, + helix_position_offsets=metadata.helix_position_offsets, + helix_is_inactive_rank=metadata.helix_is_inactive_rank, + num_q_heads=num_q_heads, + ) + if not hp_pool_updated: + raise RuntimeError("FP4 MLA generation KV append did not update the HP pool.") + + +def _scatter_generation_kv_only( + metadata: Any, + latent_cache: torch.Tensor, + layer_idx: int, + *, + token_offset: int, + local_layer: int, + v_head_dim: int, + rotary_cos_sin: torch.Tensor, + helix_position_offsets: torch.Tensor | None, + helix_is_inactive_rank: torch.Tensor | None, + num_q_heads: int, +) -> bool: + """The generation phase of the backend's ``scatter_fp4_mla_kv_cache`` with the query rows + left to the caller; returns whether it updated the HP pool.""" + state = metadata.fp4_mla_state + state.generation_cache_scattered = False + state.prequantized_q = None + state.prequantized_q_sf = None + state.q_batch_capacity = None + state.prequantized_q_rows_per_seq = None + if latent_cache.numel() == 0: + raise ValueError("FP4 MLA cache scatter requires at least one latent token.") + + latent_cache = latent_cache.reshape(latent_cache.shape[0], -1).contiguous() + num_tokens = latent_cache.shape[0] + head_dim = latent_cache.shape[-1] + if head_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA KV head_dim must be divisible by {FP4_BLOCK_SIZE}, got {head_dim}." + ) + indices_len = state.batch_indices.shape[0] + positions_len = state.positions.shape[0] + if token_offset + num_tokens > indices_len or token_offset + num_tokens > positions_len: + raise RuntimeError( + f"FP4 MLA scatter would read batch_indices[{token_offset}:" + f"{token_offset + num_tokens}] / positions[{token_offset}:" + f"{token_offset + num_tokens}], but only {indices_len} / " + f"{positions_len} entries are available: latent_cache must hold the generation " + "tokens only." + ) + + _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) + + backend = _fp4_mla_attention_backend() + global_scale = _get_fp4_mla_global_scale(metadata, latent_cache.device) + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + storage_head_dim = _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=backend, + ) + sf_per_token = storage_head_dim // FP4_BLOCK_SIZE + + kv_cache_manager = getattr(metadata, "kv_cache_manager", None) + if getattr(kv_cache_manager, "fp4_mla_context_only_cache", False) is True: + raise RuntimeError("FP4 MLA context-only cache cannot execute generation.") + if getattr(state, "v_scale_pool", None) is None: + raise RuntimeError("FP4 MLA scatter requires the auxiliary V scale pool.") + if metadata.page_size % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA scatter requires page_size divisible by " + f"{FP4_BLOCK_SIZE}, got {metadata.page_size}." + ) + if v_head_dim > head_dim: + raise ValueError(f"FP4 MLA v_head_dim={v_head_dim} cannot exceed head_dim={head_dim}.") + if head_dim - v_head_dim != FP4_MLA_K_RESIDUAL_DIM: + raise ValueError( + "FP4 MLA K residual quantization requires the K-only tail to match " + f"the {FP4_MLA_K_RESIDUAL_DIM}-channel residual, got " + f"head_dim={head_dim} v_head_dim={v_head_dim}." + ) + if v_head_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA v_head_dim must be divisible by {FP4_BLOCK_SIZE}, got {v_head_dim}." + ) + + sf_cache = sf_cache.view(torch.float8_e4m3fn) + v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=v_head_dim) + num_dim_blocks = triton.cdiv(head_dim, FP4_BLOCK_SIZE) + sf_per_page = metadata.page_size // FP4_BLOCK_SIZE + + if (helix_position_offsets is None) != (helix_is_inactive_rank is None): + raise ValueError( + "FP4 MLA Helix position-offset and inactive-rank metadata must be provided together." + ) + if rotary_cos_sin is None: + raise ValueError("FP4 MLA generation with external query rows requires rotary_cos_sin.") + num_gen_for_q = metadata.num_seqs - metadata.num_contexts + q_rows_per_seq = 0 + if num_gen_for_q > 0 and num_tokens % num_gen_for_q == 0: + q_rows_per_seq = _fp4_mla_packed_rows_per_seq(num_q_heads, num_tokens // num_gen_for_q) + q_fp4_out, q_sf_out, q_batch_capacity = _prepare_fp4_mla_q_buffers( + metadata, + num_tokens, + num_q_heads, + latent_cache.device, + rows_per_seq=q_rows_per_seq, + ) + state.prequantized_q = q_fp4_out + state.prequantized_q_sf = q_sf_out + state.q_batch_capacity = q_batch_capacity + state.prequantized_q_rows_per_seq = q_rows_per_seq + hp_pool_updated = True + + cutedsl_backend = backend == _FP4_MLA_CUTEDSL_BACKEND + fused_v_transpose = cutedsl_backend and _fp4_mla_cutedsl_fused_v_transpose_enabled() + persistent_v_packed = None + v_packed_base = None + v_page_offset = 0 + direct_v_packed_write = False + if cutedsl_backend and not fused_v_transpose: + persistent_v_packed = _get_cutedsl_persistent_v_packed_cache( + metadata, + local_layer, + kv_cache, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + ) + # Reused/imported pages may not carry this process-local sidecar. Write + # packed V directly only when the cache pages are exclusively owned. + num_gen = metadata.num_seqs - metadata.num_contexts + block_reuse = getattr(kv_cache_manager, "enable_block_reuse", True) + direct_v_packed_write = ( + num_gen > 0 + and num_tokens >= num_gen + and num_tokens % num_gen == 0 + and block_reuse is False + ) + if direct_v_packed_write: + v_packed_base = _get_fp4_mla_v_packed_pool_base(metadata) + get_v_page_offset = getattr(kv_cache_manager, "get_mla_v_packed_page_offset", None) + v_page_offset = ( + int(get_v_page_offset(local_layer)) + if callable(get_v_page_offset) + else local_layer * kv_cache.shape[0] + ) + expected_row_width = metadata.page_size // 2 + required_base_rows = (v_page_offset + kv_cache.shape[0]) * v_head_dim + if ( + not isinstance(v_packed_base, torch.Tensor) + or v_packed_base.dtype != torch.uint8 + or v_packed_base.device != kv_cache.device + or v_packed_base.ndim != 2 + or v_packed_base.shape[0] < required_base_rows + or v_packed_base.shape[1] != expected_row_width + or not v_packed_base.is_contiguous() + ): + raise RuntimeError( + "FP4 MLA direct V-packed cache update requires the " + "stable full-pool base to be a contiguous uint8 tensor " + f"with at least {required_base_rows} rows and " + f"{expected_row_width} columns on {kv_cache.device}." + ) + expected_layer_ptr = v_packed_base.data_ptr() + ( + v_page_offset * v_head_dim * expected_row_width + ) + if expected_layer_ptr != persistent_v_packed.data_ptr(): + raise RuntimeError( + "FP4 MLA V-packed layer view does not match its stable " + "full-pool base and page offset." + ) + + generation_state = _scatter_generation_kv_only_2d( + metadata, + latent_cache, + kv_cache, + sf_cache, + v_sf, + global_scale, + token_offset=token_offset, + local_layer=local_layer, + v_head_dim=v_head_dim, + head_dim=head_dim, + num_tokens=num_tokens, + num_dim_blocks=num_dim_blocks, + sf_per_token=sf_per_token, + sf_per_page=sf_per_page, + rotary_cos_sin=rotary_cos_sin, + num_q_heads=num_q_heads, + q_fp4_out=q_fp4_out, + q_sf_out=q_sf_out, + v_packed_base=v_packed_base, + v_page_offset=v_page_offset, + helix_position_offsets=helix_position_offsets, + helix_is_inactive_rank=helix_is_inactive_rank, + ) + v_pack_page_ids = _fp4_mla_generation_page_ids( + metadata, metadata.num_seqs - metadata.num_contexts + ) + v_pack_num_valid = None + if getattr(metadata, "is_cuda_graph", False): + # Frozen launch grids cannot follow the per-replay page count; + # the repack kernels stride over this device-side count instead. + v_pack_num_valid = _fp4_mla_generation_num_blocks_device(metadata) + if cutedsl_backend and not fused_v_transpose and not direct_v_packed_write: + cutedsl_repack_page_indptr = None + cutedsl_repack_kv_lens = None + cutedsl_repack_generation_lens = None + cutedsl_repack_max_touched_pages = 1 + if generation_state is not None: + kv_lens_gen, gen_lens_gen, generation_page_ids = generation_state + num_gen = kv_lens_gen.numel() + cutedsl_v_pack_page_ids = generation_page_ids + cutedsl_repack_page_indptr = state.paged_kv_indptr_decode + cutedsl_repack_kv_lens = kv_lens_gen + cutedsl_repack_generation_lens = gen_lens_gen + cutedsl_repack_max_touched_pages = _ceil_div( + num_tokens // num_gen + metadata.page_size - 1, + metadata.page_size, + ) + else: + cutedsl_v_pack_page_ids = v_pack_page_ids + _repack_cutedsl_v_packed_cache( + persistent_v_packed, + kv_cache, + cutedsl_v_pack_page_ids, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + page_indptr=cutedsl_repack_page_indptr, + kv_lens=cutedsl_repack_kv_lens, + generation_lens=cutedsl_repack_generation_lens, + max_touched_pages=cutedsl_repack_max_touched_pages, + ) + _maybe_update_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + v_pack_page_ids, + num_queries=num_tokens, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + local_layer=local_layer, + v_sf=v_sf[local_layer], + num_valid_pages=v_pack_num_valid, + ) + state.generation_cache_scattered = hp_pool_updated + return hp_pool_updated + + +def _scatter_generation_kv_only_2d( + metadata: Any, + latent_cache: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + *, + token_offset: int, + local_layer: int, + v_head_dim: int, + head_dim: int, + num_tokens: int, + num_dim_blocks: int, + sf_per_token: int, + sf_per_page: int, + rotary_cos_sin: torch.Tensor, + num_q_heads: int, + q_fp4_out: torch.Tensor, + q_sf_out: torch.Tensor, + v_packed_base: Optional[torch.Tensor], + v_page_offset: int, + helix_position_offsets: Optional[torch.Tensor], + helix_is_inactive_rank: Optional[torch.Tensor], +) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """The generation KV update of a query of ``num_q_heads`` heads whose rows the caller writes.""" + num_contexts = metadata.num_contexts + num_seqs = metadata.num_seqs + num_gen = num_seqs - num_contexts + if num_gen <= 0: + return + if num_tokens < num_gen: + raise RuntimeError( + f"FP4 MLA 2D generation scatter needs at least {num_gen} generation " + f"tokens, got {num_tokens}." + ) + if num_tokens % num_gen != 0: + raise NotImplementedError( + "FP4 MLA no-dequant generation scatter requires a uniform linear MTP " + f"generation length, got {num_tokens} tokens for {num_gen} sequences." + ) + # The prompt_lens/kv_lens runtime aliases can lag at the decode anchor + # (seq_lens == 1) under CUDA graph / one-engine MTP while each generation + # sequence really appends num_tokens // num_gen tokens this step. Recover the + # true per-sequence lengths for the no-dequant kernel below (a no-op when the + # aliases already match). As the first FP4 MLA kernel of the step, the same + # launch also writes the packed split plan the decode layers reuse. + plan_geometry = None + if num_q_heads > 0: + plan_geometry = _fp4_mla_packed_plan_geometry( + metadata, + num_heads=num_q_heads, + query_len=num_tokens // num_gen, + num_requests=num_gen, + device=latent_cache.device, + helix_kv_bounds=_fp4_mla_helix_plan_bounds( + metadata, + num_tokens, + use_helix=helix_position_offsets is not None or helix_is_inactive_rank is not None, + ), + ) + kv_lens_gen, gen_lens_gen = _fp4_mla_uniform_generation_lengths( + metadata, num_tokens, num_gen, plan_geometry=plan_geometry + ) + _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens_gen) + + pool = getattr(metadata.fp4_mla_state, "hp_pool", None) + if pool is None: + raise RuntimeError("FP4 MLA 2D generation scatter requires the HP KV pool.") + try: + hp_pool_size, hp_head_dim = _get_fp4_mla_hp_pool_layout(metadata, pool) + except ValueError as error: + raise RuntimeError(str(error)) from error + if hp_head_dim < head_dim: + raise RuntimeError( + f"FP4 MLA 2D generation scatter needs at least {head_dim} HP channels, got " + f"{hp_head_dim}." + ) + hp_page_ids = _fp4_mla_generation_hp_page_ids(metadata, num_gen) + if not isinstance(hp_page_ids, torch.Tensor): + raise RuntimeError("FP4 MLA generation requires HP page metadata.") + num_hp_pages = pool.shape[0] + + max_gen_len = num_tokens // num_gen + use_helix = helix_position_offsets is not None or helix_is_inactive_rank is not None + use_helix_local_slots = use_helix and bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + if use_helix: + if helix_position_offsets is None or helix_is_inactive_rank is None: + raise RuntimeError( + "FP4 MLA Helix requires both position-offset and inactive-rank metadata." + ) + if max_gen_len != 1 and not use_helix_local_slots: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + if ( + helix_position_offsets.dtype != torch.int32 + or helix_position_offsets.device != latent_cache.device + or helix_position_offsets.ndim != 1 + or helix_position_offsets.numel() < num_tokens + or not helix_position_offsets.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix position offsets must be a contiguous same-device " + "int32 tensor covering every generation token." + ) + if ( + helix_is_inactive_rank.dtype != torch.bool + or helix_is_inactive_rank.device != latent_cache.device + or helix_is_inactive_rank.ndim != 1 + or helix_is_inactive_rank.numel() < num_gen + or not helix_is_inactive_rank.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix inactive-rank metadata must be a contiguous " + "same-device bool tensor covering every generation sequence." + ) + if use_helix_local_slots: + helix_local_slots = getattr(metadata, "helix_local_slots", None) + if ( + not isinstance(helix_local_slots, torch.Tensor) + or helix_local_slots.dtype != torch.int32 + or helix_local_slots.device != latent_cache.device + or helix_local_slots.ndim != 1 + or helix_local_slots.numel() < num_tokens + or not helix_local_slots.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix local slots must be a contiguous " + "same-device int32 tensor covering every generation token." + ) + else: + helix_local_slots = helix_position_offsets + else: + helix_position_offsets = kv_lens_gen + helix_local_slots = kv_lens_gen + helix_is_inactive_rank = gen_lens_gen + _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) + page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) + rope_dim = head_dim - v_head_dim + block_q_heads = 32 + q1_kv_blocks_per_program = 1 + # Grouped Q1 kernels reuse K's packed codes for V. Keep one dimension + # block per program until their warp-specialized V quantizer is split out. + if max_gen_len == 1: + q1_kv_blocks_per_program = _fp4_mla_q1_kv_blocks_per_program(num_gen, v_head_dim) + q_prefix_block_dim = ( + FP4_MLA_Q1_PREFIX_BLOCK_DIM if max_gen_len == 1 else FP4_MLA_Q_PREFIX_BLOCK_DIM + ) + q_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // q_prefix_block_dim + q_prefix_blocks_per_program = _fp4_mla_q1_prefix_blocks_per_program( + num_gen, + q1_kv_blocks_per_program, + ) + q_work_blocks = q_prefix_blocks // q_prefix_blocks_per_program + 1 # noqa: F841 + if latent_cache.dtype != torch.bfloat16: + raise TypeError( + "Fused FP4 MLA Q/K RoPE and cache storage requires BF16 latent KV, " + f"got {latent_cache.dtype}." + ) + if rope_dim <= 0 or rope_dim % 2 != 0: + raise ValueError( + "Fused FP4 MLA K RoPE requires a positive even K tail, " + f"got head_dim={head_dim} v_head_dim={v_head_dim}." + ) + if ( + rotary_cos_sin is None + or rotary_cos_sin.device != latent_cache.device + or rotary_cos_sin.dtype != torch.float32 + or rotary_cos_sin.numel() % (rope_dim * 2) != 0 + ): + raise ValueError( + "Fused FP4 MLA K RoPE requires a same-device FP32 rotary " + f"table with rows of {rope_dim * 2} values." + ) + if num_q_heads <= 0: + raise ValueError( + f"FP4 MLA generation needs a positive query head count, got {num_q_heads}." + ) + q_head_blocks = _ceil_div(num_q_heads, block_q_heads) + q_rows_per_seq = _fp4_mla_packed_rows_per_seq(num_q_heads, max_gen_len) + if q_rows_per_seq and q_fp4_out.shape[0] < num_gen * q_rows_per_seq: + raise RuntimeError( + "FP4 MLA packed Q buffer is smaller than the generation batch: " + f"{q_fp4_out.shape[0]} rows < {num_gen} x {q_rows_per_seq}." + ) + max_gen_tiles = _ceil_div(max_gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) + rotary_table = rotary_cos_sin + q_global_scale = _get_fp4_mla_q_global_scale(metadata, latent_cache.device) + # Without query rows the kernel reads none of the three; any device tensor fills their slots. + q_pe_input = latent_cache + q_rope_output = latent_cache + q_full_input = latent_cache + q_fp4_output = q_fp4_out + q_sf_output = q_sf_out + write_v_packed = v_packed_base is not None + v_packed_output = v_packed_base if write_v_packed else kv_cache + v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 + v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 + kv_work_blocks = ( + v_head_dim // FP4_BLOCK_SIZE // q1_kv_blocks_per_program + 1 + if max_gen_len == 1 + else num_dim_blocks + ) + launch_grid = (num_gen, kv_work_blocks) + store_k_residual = _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS + + def launch_generation_update( + grid: tuple[int, ...], + *, + page_ids_len: int, + indptr_len: int, + max_gen_tiles_variant: int, + q_prefix_block_dim_variant: int, + q_prefix_blocks_variant: int, + q_prefix_blocks_per_program_variant: int, + q1_kv_blocks_per_program_variant: int, + ) -> None: + q_work_blocks_variant = q_prefix_blocks_variant // q_prefix_blocks_per_program_variant + 1 + _fp4_mla_generation_kv_update_kernel[grid]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + pool, + latent_cache, + global_scale, + q_global_scale, + rotary_table, + q_pe_input, + q_rope_output, + q_full_input, + q_fp4_output, + q_sf_output, + kv_lens_gen, + gen_lens_gen, + helix_position_offsets, + helix_local_slots, + helix_is_inactive_rank, + page_ids, + hp_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode, + page_ids_len, + hp_page_ids.numel(), + indptr_len, + v_sf.shape[1], + num_hp_pages, + v_sf.shape[0], + local_layer, + v_page_offset if write_v_packed else 0, + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + pool.stride(0), + pool.stride(1), + v_sf.stride(0), + v_sf.stride(1), + v_packed_s0, + v_packed_s1, + 0, + 0, + 0, + 0, + 0, + 0, + HEAD_D=hp_head_dim, + V_HEAD_D=v_head_dim, + HP_BLOCK=FP4_BLOCK_SIZE, + HP_POOL_SIZE=hp_pool_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=sf_per_page, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + STORE_K_RESIDUAL=store_k_residual, + FUSE_ROPE_CACHE_STORE=True, + USE_HELIX=use_helix, + USE_HELIX_LOCAL_SLOTS=use_helix_local_slots, + WRITE_V_PACKED=write_v_packed, + MAX_GEN_TILES=max_gen_tiles_variant, + ROPE_DIM=rope_dim, + ROPE_PAIR_BLOCK=triton.next_power_of_2(rope_dim // 2), + NUM_DIM_BLOCKS=num_dim_blocks, + NUM_Q_HEADS=num_q_heads, + Q_HEAD_BLOCKS=max(q_head_blocks, 1), + BLOCK_Q_HEADS=block_q_heads, + Q_PREFIX_D=FP4_MLA_Q_PREFIX_DIM, + Q_PREFIX_BLOCK_D=q_prefix_block_dim_variant, + Q_PREFIX_BLOCKS=q_prefix_blocks_variant, + Q_PREFIX_BLOCKS_PER_PROGRAM=q_prefix_blocks_per_program_variant, + Q_WORK_BLOCKS=q_work_blocks_variant, + Q_SF_COLS=FP4_MLA_Q_SF_GROUPS, + WRITE_Q=True, + Q1_KV_BLOCKS_PER_PROGRAM=q1_kv_blocks_per_program_variant, + Q_ROWS_PER_SEQ=q_rows_per_seq, + WRITE_Q_ROWS=False, + maxnreg=56, + ) + + launch_generation_update( + launch_grid, + page_ids_len=page_ids.shape[0], + indptr_len=metadata.fp4_mla_state.paged_kv_indptr_decode.shape[0], + max_gen_tiles_variant=max(max_gen_tiles, 1), + q_prefix_block_dim_variant=q_prefix_block_dim, + q_prefix_blocks_variant=q_prefix_blocks, + q_prefix_blocks_per_program_variant=q_prefix_blocks_per_program, + q1_kv_blocks_per_program_variant=q1_kv_blocks_per_program, + ) + return kv_lens_gen, gen_lens_gen, page_ids + + +@triton.jit( + do_not_specialize=[ + "page_ids_len", + "hp_page_ids_len", + "indptr_len", + "num_pages", + "num_hp_pages", + "num_layers", + "local_layer", + "v_page_offset", + ], + do_not_specialize_on_alignment=[ + "page_ids_len", + "hp_page_ids_len", + "indptr_len", + "num_pages", + "num_hp_pages", + "num_layers", + "local_layer", + "v_page_offset", + ], +) +def _fp4_mla_generation_kv_update_kernel( + kv_cache_ptr, + sf_cache_ptr, + v_sf_ptr, + v_packed_ptr, + hp_pool_ptr, + latent_cache_ptr, + global_scale_ptr, + q_global_scale_ptr, + rotary_cos_sin_ptr, + q_pe_ptr, + q_rope_out_ptr, + q_full_ptr, + q_fp4_out_ptr, + q_sf_out_ptr, + kv_lens_ptr, + prompt_lens_ptr, + helix_position_offsets_ptr, + helix_local_slots_ptr, + helix_is_inactive_rank_ptr, + page_ids_ptr, + hp_page_ids_ptr, + paged_kv_indptr_ptr, + page_ids_len, + hp_page_ids_len, + indptr_len, + num_pages, + num_hp_pages, + num_layers, + local_layer, + v_page_offset, + page_size: tl.constexpr, + kv_s0, + kv_s2, + kv_s4, + sf_s0, + pool_s0, + pool_s1, + vsf_s0, + vsf_s1, + v_packed_s0, + v_packed_s1, + q_pe_s0, + q_pe_s1, + q_pe_s2, + q_out_s0, + q_out_s1, + q_out_s2, + HEAD_D: tl.constexpr, + V_HEAD_D: tl.constexpr, + HP_BLOCK: tl.constexpr, + HP_POOL_SIZE: tl.constexpr, + FP4_BLOCK: tl.constexpr, + SF_PER_TOKEN: tl.constexpr, + SF_PER_PAGE: tl.constexpr, + K_RESIDUAL_D: tl.constexpr, + STORE_K_RESIDUAL: tl.constexpr, + FUSE_ROPE_CACHE_STORE: tl.constexpr, + USE_HELIX: tl.constexpr, + USE_HELIX_LOCAL_SLOTS: tl.constexpr, + WRITE_V_PACKED: tl.constexpr, + MAX_GEN_TILES: tl.constexpr, + ROPE_DIM: tl.constexpr, + ROPE_PAIR_BLOCK: tl.constexpr, + NUM_DIM_BLOCKS: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + Q_HEAD_BLOCKS: tl.constexpr, + BLOCK_Q_HEADS: tl.constexpr, + Q_PREFIX_D: tl.constexpr, + Q_PREFIX_BLOCK_D: tl.constexpr, + Q_PREFIX_BLOCKS: tl.constexpr, + Q_PREFIX_BLOCKS_PER_PROGRAM: tl.constexpr, + Q_WORK_BLOCKS: tl.constexpr, + Q_SF_COLS: tl.constexpr, + WRITE_Q: tl.constexpr, + Q1_KV_BLOCKS_PER_PROGRAM: tl.constexpr, + Q_ROWS_PER_SEQ: tl.constexpr = 0, + WRITE_Q_ROWS: tl.constexpr = True, +): + q1_shared_main: tl.constexpr = FUSE_ROPE_CACHE_STORE and MAX_GEN_TILES == 1 + q1_grouped_kv: tl.constexpr = q1_shared_main and Q1_KV_BLOCKS_PER_PROGRAM > 1 + q1_main_kv_programs: tl.constexpr = V_HEAD_D // FP4_BLOCK // Q1_KV_BLOCKS_PER_PROGRAM + q_prefix_programs: tl.constexpr = Q_PREFIX_BLOCKS // Q_PREFIX_BLOCKS_PER_PROGRAM + q1_kv_blocks_per_iteration: tl.constexpr = 2 + tl.static_assert(Q1_KV_BLOCKS_PER_PROGRAM == 1 or q1_grouped_kv) + tl.static_assert( + Q1_KV_BLOCKS_PER_PROGRAM == 1 or Q1_KV_BLOCKS_PER_PROGRAM % q1_kv_blocks_per_iteration == 0 + ) + tl.static_assert(V_HEAD_D % (FP4_BLOCK * Q1_KV_BLOCKS_PER_PROGRAM) == 0) + tl.static_assert(Q_PREFIX_BLOCKS % Q_PREFIX_BLOCKS_PER_PROGRAM == 0) + seq_idx = tl.program_id(0) + if (local_layer < 0) | (local_layer >= num_layers): + return + if seq_idx + 1 >= indptr_len: + return + gen_len = tl.load(prompt_lens_ptr + seq_idx) + if gen_len <= 0: + return + kv_len = tl.load(kv_lens_ptr + seq_idx) + # Generation tokens are request-major. Deriving their absolute position + # here avoids materializing per-token metadata in a CUDA graph. + first_new_pos = kv_len - gen_len + rope_first_new_pos = first_new_pos + if USE_HELIX: + rope_first_new_pos = tl.load(helix_position_offsets_ptr + seq_idx * gen_len) + if USE_HELIX_LOCAL_SLOTS: + # Helix owns a contiguous rank-local subset of a linear verify group, + # but that subset can begin at any token in the group. + first_new_pos = kv_len + for token_idx in tl.static_range(0, MAX_GEN_TILES * HP_BLOCK): + local_slot = tl.load( + helix_local_slots_ptr + seq_idx * gen_len + token_idx, + mask=token_idx < gen_len, + other=-1, + ) + first_new_pos = tl.minimum(first_new_pos, tl.where(local_slot >= 0, local_slot, kv_len)) + # HP ring slots are addressed by global position (the context writer and + # the ring transfer use it). Helix cache positions are rank-local; the + # difference only cancels when HP_POOL_SIZE divides the page size. + hp_pos_offset = first_new_pos - first_new_pos + if USE_HELIX_LOCAL_SLOTS: + for token_idx in tl.static_range(0, MAX_GEN_TILES * HP_BLOCK): + local_slot = tl.load( + helix_local_slots_ptr + seq_idx * gen_len + token_idx, + mask=token_idx < gen_len, + other=-1, + ) + global_pos = tl.load( + helix_position_offsets_ptr + seq_idx * gen_len + token_idx, + mask=token_idx < gen_len, + other=0, + ) + hp_pos_offset = tl.where( + local_slot == first_new_pos, global_pos - local_slot, hp_pos_offset + ) + elif USE_HELIX: + hp_pos_offset = rope_first_new_pos - first_new_pos + if FUSE_ROPE_CACHE_STORE: + work_idx = tl.program_id(1) + dim_block = work_idx + if work_idx >= 0: + q_program = work_idx + if ( + q1_grouped_kv + and Q_WORK_BLOCKS == 2 + and Q_HEAD_BLOCKS >= 3 + and q_prefix_programs == 1 + ): + # Move the final prefix program ahead of the last tail + # programs to shorten the grid drain without changing Q work. + penultimate_tail: tl.constexpr = Q_HEAD_BLOCKS * Q_WORK_BLOCKS - 3 + final_prefix: tl.constexpr = penultimate_tail + 1 + q_program = tl.where( + work_idx == penultimate_tail, + final_prefix, + tl.where( + work_idx == final_prefix, + penultimate_tail, + work_idx, + ), + ) + q_token_head_program = q_program // Q_WORK_BLOCKS + q_work_block = q_program - q_token_head_program * Q_WORK_BLOCKS + q_token_idx = q_token_head_program // Q_HEAD_BLOCKS + if not WRITE_Q_ROWS: + # Another kernel writes the query rows: only the KV programs work. + q_token_idx = gen_len + q_head_block = q_token_head_program - q_token_idx * Q_HEAD_BLOCKS + if q_token_idx < gen_len: + q_token = seq_idx * gen_len + q_token_idx + head_lanes = tl.arange(0, BLOCK_Q_HEADS) + head_offsets = q_head_block * BLOCK_Q_HEADS + head_lanes + # The BF16 Q input is token-major; the FP4 output follows the + # decode layout: token-major, or packed per request where each + # request owns Q_ROWS_PER_SEQ rows (a multiple of 128) holding + # its [query token, head] pairs. + q_in_rows = q_token * NUM_Q_HEADS + head_offsets + if Q_ROWS_PER_SEQ > 0: + q_rows = seq_idx * Q_ROWS_PER_SEQ + q_token_idx * NUM_Q_HEADS + head_offsets + else: + q_rows = q_in_rows + q_input_row_stride: tl.constexpr = Q_PREFIX_D + ROPE_DIM + q_packed_row_stride: tl.constexpr = (Q_PREFIX_D + 2 * ROPE_DIM) // 2 + # Every lane is valid only when the launch covers the head + # dimension with an exact number of full head blocks. + if NUM_Q_HEADS == Q_HEAD_BLOCKS * BLOCK_Q_HEADS: + row_mask = tl.full( + [BLOCK_Q_HEADS], + True, + dtype=tl.int1, + ) + else: + row_mask = head_offsets < NUM_Q_HEADS + if WRITE_Q and q_work_block < q_prefix_programs: + if Q_PREFIX_BLOCKS_PER_PROGRAM >= 1: + tl.static_assert(Q_PREFIX_D % FP4_BLOCK == 0) + tl.static_assert(Q_PREFIX_D % Q_PREFIX_BLOCK_D == 0) + tl.static_assert(Q_PREFIX_BLOCK_D % FP4_BLOCK == 0) + q_global_scale = tl.load(q_global_scale_ptr).to(tl.float32) + prefix_dims = tl.arange(0, Q_PREFIX_BLOCK_D) + prefix_base = q_in_rows[:, None].to(tl.int64) * q_input_row_stride + prefix_groups: tl.constexpr = Q_PREFIX_BLOCK_D // FP4_BLOCK + q_prefix_iterations: tl.constexpr = Q_PREFIX_BLOCKS_PER_PROGRAM + for q_prefix_iteration in tl.range( + 0, + q_prefix_iterations, + num_stages=1, + loop_unroll_factor=1, + disable_licm=True, + ): + prefix_work_block = ( + q_work_block * Q_PREFIX_BLOCKS_PER_PROGRAM + q_prefix_iteration + ) + prefix_dim_offset = prefix_work_block * Q_PREFIX_BLOCK_D + prefix_offsets = prefix_base + ( + prefix_dim_offset + prefix_dims[None, :] + ).to(tl.int64) + prefix_values = tl.load( + q_full_ptr + prefix_offsets, + mask=row_mask[:, None], + other=0.0, + ).to(tl.float32) + prefix_pairs = tl.reshape( + prefix_values * q_global_scale, + ( + BLOCK_Q_HEADS, + prefix_groups, + FP4_BLOCK // 2, + 2, + ), + ) + prefix_even, prefix_odd = tl.split(prefix_pairs) + prefix_packed_words, prefix_scales = _fp4_mla_quantize_q_groups_words( + prefix_even, + prefix_odd, + BLOCK_Q_HEADS, + prefix_groups, + ) + prefix_group_offsets = prefix_work_block * prefix_groups + tl.arange( + 0, prefix_groups + ) + prefix_word_offsets = q_rows[:, None, None].to( + tl.int64 + ) * q_packed_row_stride + ( + prefix_group_offsets[None, :, None] * (FP4_BLOCK // 2) + + tl.arange(0, FP4_BLOCK // 8)[None, None, :] * 4 + ).to(tl.int64) + tl.store( + (q_fp4_out_ptr + prefix_word_offsets).to( + tl.pointer_type(tl.uint32) + ), + prefix_packed_words, + mask=row_mask[:, None, None], + ) + if NUM_Q_HEADS == 128 and BLOCK_Q_HEADS == 32: + tl.static_assert(NUM_Q_HEADS == 128) + tl.static_assert(BLOCK_Q_HEADS == 32) + prefix_sf_offsets = _fp4_mla_q_sf_offset_h128( + q_token, + q_head_block, + head_lanes[:, None], + prefix_group_offsets[None, :], + Q_SF_COLS, + ) + else: + prefix_sf_offsets = _fp4_mla_q_sf_offset( + q_rows[:, None], + prefix_group_offsets[None, :], + Q_SF_COLS, + ) + tl.store( + q_sf_out_ptr + prefix_sf_offsets, + prefix_scales, + mask=row_mask[:, None], + ) + else: + position = (rope_first_new_pos + q_token_idx).to(tl.int64) + if USE_HELIX_LOCAL_SLOTS: + position = tl.load(helix_position_offsets_ptr + q_token).to(tl.int64) + valid_position = position >= 0 + pair_offsets = tl.arange(0, ROPE_PAIR_BLOCK) + pair_mask = pair_offsets < ROPE_DIM // 2 + rotary_offsets = position * (ROPE_DIM * 2) + pair_offsets * 2 + cos = tl.load( + rotary_cos_sin_ptr + rotary_offsets, + mask=valid_position & pair_mask, + other=0.0, + ).to(tl.float32) + sin = tl.load( + rotary_cos_sin_ptr + rotary_offsets + 1, + mask=valid_position & pair_mask, + other=0.0, + ).to(tl.float32) + q_base = q_token * q_pe_s0 + head_offsets[:, None].to(tl.int64) * q_pe_s1 + q_dims = tl.arange(0, ROPE_DIM) + q_dim_mask = valid_position & row_mask[:, None] & (q_dims[None, :] < ROPE_DIM) + q_offsets = q_base + q_dims[None, :].to(tl.int64) * q_pe_s2 + q_values = tl.load( + q_pe_ptr + q_offsets, + mask=q_dim_mask, + other=0.0, + ).to(tl.float32) + q_pairs = tl.reshape(q_values, (BLOCK_Q_HEADS, ROPE_DIM // 2, 2)) + q_even, q_odd = tl.split(q_pairs) + q_roped_even, q_roped_odd = _fp4_mla_rope_fp32( + q_even, q_odd, cos[None, :], sin[None, :] + ) + q_roped_even = q_roped_even.to(tl.bfloat16) + q_roped_odd = q_roped_odd.to(tl.bfloat16) + q_out_base = q_token * q_out_s0 + head_offsets[:, None].to(tl.int64) * q_out_s1 + q_out_offsets = q_out_base + q_dims[None, :].to(tl.int64) * q_out_s2 + q_roped = tl.reshape( + tl.join(q_roped_even, q_roped_odd), + (BLOCK_Q_HEADS, ROPE_DIM), + ) + tl.store( + q_rope_out_ptr + q_out_offsets, + q_roped, + mask=q_dim_mask, + ) + if WRITE_Q: + tl.static_assert(ROPE_DIM == 64) + tl.static_assert(ROPE_PAIR_BLOCK == 32) + q_global_scale = tl.load(q_global_scale_ptr).to(tl.float32) + tail_even = tl.reshape( + q_roped_even.to(tl.float32) * q_global_scale, + ( + BLOCK_Q_HEADS, + ROPE_DIM // FP4_BLOCK, + FP4_BLOCK // 2, + ), + ) + tail_odd = tl.reshape( + q_roped_odd.to(tl.float32) * q_global_scale, + ( + BLOCK_Q_HEADS, + ROPE_DIM // FP4_BLOCK, + FP4_BLOCK // 2, + ), + ) + packed, scales, residual_packed, residual_scales = ( + _fp4_mla_quantize_q_residual_groups(tail_even, tail_odd) + ) + prefix_groups: tl.constexpr = Q_PREFIX_D // FP4_BLOCK + tail_group_count: tl.constexpr = ROPE_DIM // FP4_BLOCK + output_groups = prefix_groups + tl.arange(0, tail_group_count * 2) + bytes_in_group = tl.arange(0, FP4_BLOCK // 2) + packed_tail = tl.reshape( + tl.join(packed, residual_packed).permute(0, 1, 3, 2), + ( + BLOCK_Q_HEADS, + tail_group_count * 2, + FP4_BLOCK // 2, + ), + ) + output_offsets = q_rows[:, None, None].to( + tl.int64 + ) * q_packed_row_stride + ( + output_groups[None, :, None] * (FP4_BLOCK // 2) + + bytes_in_group[None, None, :] + ).to(tl.int64) + tail_row_mask = valid_position & row_mask + tl.store( + q_fp4_out_ptr + output_offsets, + packed_tail, + mask=tail_row_mask[:, None, None], + ) + tail_scales = tl.reshape( + tl.join(scales, residual_scales), + (BLOCK_Q_HEADS, tail_group_count * 2), + ) + if NUM_Q_HEADS == 128 and BLOCK_Q_HEADS == 32: + sf_offsets = _fp4_mla_q_sf_offset_h128( + q_token, + q_head_block, + head_lanes[:, None], + output_groups[None, :], + Q_SF_COLS, + ) + else: + sf_offsets = _fp4_mla_q_sf_offset( + q_rows[:, None], + output_groups[None, :], + Q_SF_COLS, + ) + tl.store( + q_sf_out_ptr + sf_offsets, + tail_scales, + mask=tail_row_mask[:, None], + ) + if q1_grouped_kv: + if work_idx > q1_main_kv_programs: + return + dim_block = tl.where( + work_idx == q1_main_kv_programs, + V_HEAD_D // FP4_BLOCK, + work_idx, + ) + else: + if work_idx >= NUM_DIM_BLOCKS: + return + else: + dim_block = tl.program_id(2) + if USE_HELIX_LOCAL_SLOTS: + if first_new_pos >= kv_len: + return + elif USE_HELIX: + if tl.load(helix_is_inactive_rank_ptr + seq_idx): + return + if FUSE_ROPE_CACHE_STORE and MAX_GEN_TILES == 1: + if dim_block * FP4_BLOCK > V_HEAD_D: + return + first_tile_pos = (first_new_pos // HP_BLOCK) * HP_BLOCK + page_start = tl.load(paged_kv_indptr_ptr + seq_idx).to(tl.int64) + page_end = tl.load(paged_kv_indptr_ptr + seq_idx + 1).to(tl.int64) + hp_index = tl.full((), 0, tl.int64) + + if q1_grouped_kv: + # A single V owner can share the K-residual program. Multi-owner + # specializations retain their interleaved dimension ownership. + if q1_main_kv_programs == 1: + kv_owner = work_idx == q1_main_kv_programs + kv_work_idx = work_idx - q1_main_kv_programs + else: + kv_owner = work_idx < q1_main_kv_programs + kv_work_idx = work_idx + if kv_owner: + block_base_pos = first_tile_pos + if block_base_pos >= kv_len: + return + if block_base_pos + HP_BLOCK <= first_new_pos: + return + + page_idx = block_base_pos // page_size + page_pos = block_base_pos - page_idx * page_size + physical_page_offset = page_start + page_idx + if ( + (page_pos < 0) + | (page_pos >= page_size) + | (physical_page_offset < page_start) + | (physical_page_offset >= page_end) + | (physical_page_offset < 0) + | (physical_page_offset >= page_ids_len) + ): + return + physical_page = tl.load(page_ids_ptr + physical_page_offset).to(tl.int64) + if (physical_page < 0) | (physical_page >= num_pages): + return + if physical_page_offset >= hp_page_ids_len: + return + hp_index = tl.load(hp_page_ids_ptr + physical_page_offset).to(tl.int64) + if (hp_index < 0) | (hp_index >= num_hp_pages): + return + + latent_token = seq_idx * gen_len + global_scale = tl.load(global_scale_ptr) + # Each grouped CTA stages only the dimensions it owns. The barrier + # then makes the new BF16 token visible to that CTA's rolled loop. + prestage_linear = tl.arange(0, Q1_KV_BLOCKS_PER_PROGRAM * FP4_BLOCK) + prestage_iteration = prestage_linear // FP4_BLOCK + prestage_lane = prestage_linear % FP4_BLOCK + prestage_dim = ( + kv_work_idx + prestage_iteration * q1_main_kv_programs + ) * FP4_BLOCK + prestage_lane + prestage_values = tl.load( + latent_cache_ptr + + latent_token.to(tl.int64) * (NUM_DIM_BLOCKS * FP4_BLOCK) + + prestage_dim.to(tl.int64) + ) + prestage_slot = (first_new_pos + hp_pos_offset) % HP_POOL_SIZE + tl.store( + hp_pool_ptr + + hp_index * pool_s0 + + local_layer * pool_s1 + + prestage_slot * HEAD_D + + prestage_dim, + prestage_values, + ) + tl.debug_barrier() + tail_count = kv_len - block_base_pos + if tail_count <= HP_BLOCK // 2: + tl.static_assert(Q1_KV_BLOCKS_PER_PROGRAM % 4 == 0) + for dim_group in tl.range( + 0, + Q1_KV_BLOCKS_PER_PROGRAM // 4, + num_stages=1, + loop_unroll_factor=1, + disable_licm=True, + ): + dim_block_base = kv_work_idx + dim_group * 4 * q1_main_kv_programs + _fp4_mla_generation_q1_kv_tail8_update( + kv_cache_ptr, + sf_cache_ptr, + v_sf_ptr, + v_packed_ptr, + hp_pool_ptr, + hp_index, + kv_len, + block_base_pos, + hp_pos_offset, + page_pos, + physical_page, + dim_block_base, + global_scale, + local_layer, + v_page_offset, + kv_s0, + kv_s2, + kv_s4, + sf_s0, + pool_s0, + pool_s1, + vsf_s0, + vsf_s1, + v_packed_s0, + v_packed_s1, + HEAD_D, + V_HEAD_D, + HP_BLOCK, + HP_POOL_SIZE, + FP4_BLOCK, + SF_PER_TOKEN, + SF_PER_PAGE, + WRITE_V_PACKED, + q1_main_kv_programs, + ) + else: + for dim_group in tl.range( + 0, + Q1_KV_BLOCKS_PER_PROGRAM // q1_kv_blocks_per_iteration, + num_stages=1, + loop_unroll_factor=1, + disable_licm=True, + ): + dim_iterations = dim_group * q1_kv_blocks_per_iteration + tl.arange( + 0, q1_kv_blocks_per_iteration + ) + pair_dim_blocks = kv_work_idx + dim_iterations * q1_main_kv_programs + _fp4_mla_generation_q1_kv_tile_update( + kv_cache_ptr, + sf_cache_ptr, + v_sf_ptr, + v_packed_ptr, + hp_pool_ptr, + hp_index, + kv_len, + block_base_pos, + hp_pos_offset, + page_pos, + physical_page, + pair_dim_blocks, + global_scale, + local_layer, + v_page_offset, + kv_s0, + kv_s2, + kv_s4, + sf_s0, + pool_s0, + pool_s1, + vsf_s0, + vsf_s1, + v_packed_s0, + v_packed_s1, + HEAD_D, + V_HEAD_D, + HP_BLOCK, + HP_POOL_SIZE, + FP4_BLOCK, + SF_PER_TOKEN, + SF_PER_PAGE, + WRITE_V_PACKED, + ) + if work_idx < q1_main_kv_programs: + return + + byte_offsets = tl.arange(0, HP_BLOCK // 2) + token_offsets = tl.arange(0, HP_BLOCK) + even_d = dim_block * FP4_BLOCK + byte_offsets * 2 + odd_d = even_d + 1 + all_d = dim_block * FP4_BLOCK + tl.arange(0, FP4_BLOCK) + mask_even_d = even_d < HEAD_D + mask_odd_d = odd_d < HEAD_D + mask_all_d = all_d < HEAD_D + safe_even_d = tl.where(mask_even_d, even_d, 0) + safe_odd_d = tl.where(mask_odd_d, odd_d, 0) + safe_all_d = tl.where(mask_all_d, all_d, 0) + + if FUSE_ROPE_CACHE_STORE and MAX_GEN_TILES == 1: + if dim_block * FP4_BLOCK >= V_HEAD_D: + tl.static_assert(STORE_K_RESIDUAL) + tl.static_assert(HEAD_D - K_RESIDUAL_D == V_HEAD_D) + tl.static_assert(K_RESIDUAL_D % FP4_BLOCK == 0) + position = first_new_pos.to(tl.int64) + if position < 0: + return + # Under helix the cache slot is rank-local while RoPE must use the + # GLOBAL position; they coincide everywhere else. + rope_position = rope_first_new_pos.to(tl.int64) + page_idx = position // page_size + page_pos = position - page_idx * page_size + physical_page_offset = page_start + page_idx + if ( + (page_pos < 0) + | (page_pos >= page_size) + | (physical_page_offset < page_start) + | (physical_page_offset >= page_end) + | (physical_page_offset < 0) + | (physical_page_offset >= page_ids_len) + ): + return + physical_page = tl.load(page_ids_ptr + physical_page_offset).to(tl.int64) + if (physical_page < 0) | (physical_page >= num_pages): + return + if physical_page_offset >= hp_page_ids_len: + return + hp_index = tl.load(hp_page_ids_ptr + physical_page_offset).to(tl.int64) + if (hp_index < 0) | (hp_index >= num_hp_pages): + return + + latent_token = seq_idx * gen_len + tail_byte_offsets = tl.arange(0, K_RESIDUAL_D // 2) + tail_even_d = V_HEAD_D + tail_byte_offsets * 2 + tail_odd_d = tail_even_d + 1 + tail_even = tl.load( + latent_cache_ptr + latent_token * (NUM_DIM_BLOCKS * FP4_BLOCK) + tail_even_d, + ).to(tl.float32) + tail_odd = tl.load( + latent_cache_ptr + latent_token * (NUM_DIM_BLOCKS * FP4_BLOCK) + tail_odd_d, + ).to(tl.float32) + rope_pair_offsets = tail_byte_offsets + rotary_offsets = rope_position * (ROPE_DIM * 2) + rope_pair_offsets * 2 + cos = tl.load(rotary_cos_sin_ptr + rotary_offsets).to(tl.float32) + sin = tl.load(rotary_cos_sin_ptr + rotary_offsets + 1).to(tl.float32) + tail_even, tail_odd = _fp4_mla_rope_fp32(tail_even, tail_odd, cos, sin) + tail_even = tail_even.to(tl.bfloat16).to(tl.float32) + tail_odd = tail_odd.to(tl.bfloat16).to(tl.float32) + + residual_groups: tl.constexpr = K_RESIDUAL_D // FP4_BLOCK + tail_even_groups = tl.reshape( + tail_even, + (residual_groups, FP4_BLOCK // 2), + ) + tail_odd_groups = tl.reshape( + tail_odd, + (residual_groups, FP4_BLOCK // 2), + ) + token_amax = tl.maximum( + tl.max(tl.abs(tail_even_groups), axis=1), + tl.max(tl.abs(tail_odd_groups), axis=1), + ) + token_scale = tl.where(token_amax > 0.0, token_amax / 6.0, 1.0) + global_scale = tl.load(global_scale_ptr) + stored_scale = token_scale * global_scale + quant_scale = stored_scale.to(tl.float8e4nv).to(tl.float32) / global_scale + low = _fp4_e2m1_quantize(tail_even_groups / quant_scale[:, None]) + high = _fp4_e2m1_quantize(tail_odd_groups / quant_scale[:, None]) + packed = low | (high << 4) + + group_offsets = tl.arange(0, residual_groups) + bytes_in_group = tl.arange(0, FP4_BLOCK // 2) + packed_cols = ( + V_HEAD_D // 2 + group_offsets[:, None] * (FP4_BLOCK // 2) + bytes_in_group[None, :] + ) + kv_base = physical_page * kv_s0 + tl.store( + kv_cache_ptr + kv_base + page_pos * kv_s2 + packed_cols * kv_s4, + packed, + ) + k_sf_cols = V_HEAD_D // FP4_BLOCK + group_offsets + k_sf_offset = _fp4_mla_swizzled_sf_offset(page_pos, k_sf_cols, SF_PER_TOKEN) + tl.store( + sf_cache_ptr + physical_page * sf_s0 + k_sf_offset, + stored_scale, + ) + + main_even, main_odd = _fp4_mla_unpack_e2m1x2(packed) + residual_even = tail_even_groups - main_even * quant_scale[:, None] + residual_odd = tail_odd_groups - main_odd * quant_scale[:, None] + residual_amax = tl.maximum( + tl.max(tl.abs(residual_even), axis=1), + tl.max(tl.abs(residual_odd), axis=1), + ) + residual_scale = tl.where(residual_amax > 0.0, residual_amax / 6.0, 1.0) + residual_stored_scale = residual_scale * global_scale + residual_quant_scale = ( + residual_stored_scale.to(tl.float8e4nv).to(tl.float32) / global_scale + ) + residual_low = _fp4_e2m1_quantize(residual_even / residual_quant_scale[:, None]) + residual_high = _fp4_e2m1_quantize(residual_odd / residual_quant_scale[:, None]) + residual_packed = residual_low | (residual_high << 4) + residual_packed_cols = ( + HEAD_D // 2 + group_offsets[:, None] * (FP4_BLOCK // 2) + bytes_in_group[None, :] + ) + tl.store( + kv_cache_ptr + kv_base + page_pos * kv_s2 + residual_packed_cols * kv_s4, + residual_packed, + ) + residual_sf_offset = _fp4_mla_swizzled_sf_offset( + page_pos, + HEAD_D // FP4_BLOCK + group_offsets, + SF_PER_TOKEN, + ) + tl.store( + sf_cache_ptr + physical_page * sf_s0 + residual_sf_offset, + residual_stored_scale, + ) + + hp_slot = (position + hp_pos_offset) % HP_POOL_SIZE + hp_store_base = ( + hp_pool_ptr + hp_index * pool_s0 + local_layer * pool_s1 + hp_slot * HEAD_D + ) + tl.store(hp_store_base + tail_even_d, tail_even) + tl.store(hp_store_base + tail_odd_d, tail_odd) + return + + for tile_iter in tl.static_range(0, MAX_GEN_TILES): + if FUSE_ROPE_CACHE_STORE: + tile_idx = tile_iter + else: + tile_idx = tl.program_id(1) + block_base_pos = first_tile_pos + tile_idx * HP_BLOCK + if block_base_pos >= kv_len: + return + if block_base_pos + HP_BLOCK <= first_new_pos: + return + + page_idx = block_base_pos // page_size + page_pos = block_base_pos - page_idx * page_size + physical_page_offset = page_start + page_idx + if ( + (page_pos < 0) + | (page_pos >= page_size) + | (physical_page_offset < page_start) + | (physical_page_offset >= page_end) + | (physical_page_offset < 0) + | (physical_page_offset >= page_ids_len) + ): + return + physical_page = tl.load(page_ids_ptr + physical_page_offset).to(tl.int64) + if (physical_page < 0) | (physical_page >= num_pages): + return + if physical_page_offset >= hp_page_ids_len: + return + hp_index = tl.load(hp_page_ids_ptr + physical_page_offset).to(tl.int64) + if (hp_index < 0) | (hp_index >= num_hp_pages): + return + + abs_positions = block_base_pos + token_offsets + valid_tokens = abs_positions < kv_len + hp_slots = (abs_positions + hp_pos_offset) % HP_POOL_SIZE + if USE_HELIX_LOCAL_SLOTS: + from_latent = abs_positions < 0 + new_token_offsets = abs_positions - abs_positions + for token_idx in tl.static_range(0, MAX_GEN_TILES * HP_BLOCK): + local_slot = tl.load( + helix_local_slots_ptr + seq_idx * gen_len + token_idx, + mask=token_idx < gen_len, + other=-1, + ) + is_local_token = abs_positions == local_slot + from_latent = from_latent | is_local_token + new_token_offsets = tl.where(is_local_token, token_idx, new_token_offsets) + else: + from_latent = abs_positions >= first_new_pos + new_token_offsets = abs_positions - first_new_pos + # Linear MTP uses a uniform generation length, so each sequence + # occupies one contiguous gen_len slice in latent_cache. + latent_tokens = seq_idx * gen_len + new_token_offsets + safe_latent_tokens = tl.where(valid_tokens & from_latent, latent_tokens, 0).to(tl.int64) + + if q1_shared_main: + # Q1 reads each token from exactly one source. Selecting the pointer + # first removes two masked loads and additions without changing MTP. + hp_even_ptrs = ( + hp_pool_ptr + + hp_index * pool_s0 + + local_layer * pool_s1 + + hp_slots[:, None] * HEAD_D + + safe_even_d[None, :] + ) + hp_odd_ptrs = ( + hp_pool_ptr + + hp_index * pool_s0 + + local_layer * pool_s1 + + hp_slots[:, None] * HEAD_D + + safe_odd_d[None, :] + ) + latent_even_ptrs = ( + latent_cache_ptr + + safe_latent_tokens[:, None] * (NUM_DIM_BLOCKS * FP4_BLOCK) + + safe_even_d[None, :] + ) + latent_odd_ptrs = ( + latent_cache_ptr + + safe_latent_tokens[:, None] * (NUM_DIM_BLOCKS * FP4_BLOCK) + + safe_odd_d[None, :] + ) + even_values = tl.load( + tl.where(from_latent[:, None], latent_even_ptrs, hp_even_ptrs), + mask=valid_tokens[:, None] & mask_even_d[None, :], + other=0.0, + ).to(tl.float32) + odd_values = tl.load( + tl.where(from_latent[:, None], latent_odd_ptrs, hp_odd_ptrs), + mask=valid_tokens[:, None] & mask_odd_d[None, :], + other=0.0, + ).to(tl.float32) + else: + hp_even = tl.load( + hp_pool_ptr + + hp_index * pool_s0 + + local_layer * pool_s1 + + hp_slots[:, None] * HEAD_D + + safe_even_d[None, :], + mask=valid_tokens[:, None] & (~from_latent)[:, None] & mask_even_d[None, :], + other=0.0, + ).to(tl.float32) + hp_odd = tl.load( + hp_pool_ptr + + hp_index * pool_s0 + + local_layer * pool_s1 + + hp_slots[:, None] * HEAD_D + + safe_odd_d[None, :], + mask=valid_tokens[:, None] & (~from_latent)[:, None] & mask_odd_d[None, :], + other=0.0, + ).to(tl.float32) + latent_even = tl.load( + latent_cache_ptr + + safe_latent_tokens[:, None] * (NUM_DIM_BLOCKS * FP4_BLOCK) + + safe_even_d[None, :], + mask=valid_tokens[:, None] & from_latent[:, None] & mask_even_d[None, :], + other=0.0, + ).to(tl.float32) + latent_odd = tl.load( + latent_cache_ptr + + safe_latent_tokens[:, None] * (NUM_DIM_BLOCKS * FP4_BLOCK) + + safe_odd_d[None, :], + mask=valid_tokens[:, None] & from_latent[:, None] & mask_odd_d[None, :], + other=0.0, + ).to(tl.float32) + even_values = hp_even + latent_even + odd_values = hp_odd + latent_odd + + if FUSE_ROPE_CACHE_STORE and not q1_shared_main: + if dim_block * FP4_BLOCK >= V_HEAD_D: + positions = (rope_first_new_pos + new_token_offsets).to(tl.int64) + if USE_HELIX_LOCAL_SLOTS: + positions = tl.load( + helix_position_offsets_ptr + safe_latent_tokens, + mask=valid_tokens & from_latent, + other=-1, + ).to(tl.int64) + valid_position = valid_tokens & from_latent & (positions >= 0) + rope_pair_offsets = (safe_even_d - V_HEAD_D) // 2 + rope_dim_mask = mask_even_d & (even_d >= V_HEAD_D) & (odd_d < V_HEAD_D + ROPE_DIM) + rotary_offsets = ( + positions[:, None] * (ROPE_DIM * 2) + rope_pair_offsets[None, :] * 2 + ) + rotary_mask = valid_position[:, None] & rope_dim_mask[None, :] + cos = tl.load( + rotary_cos_sin_ptr + rotary_offsets, + mask=rotary_mask, + other=0.0, + ).to(tl.float32) + sin = tl.load( + rotary_cos_sin_ptr + rotary_offsets + 1, + mask=rotary_mask, + other=0.0, + ).to(tl.float32) + roped_even, roped_odd = _fp4_mla_rope_fp32(even_values, odd_values, cos, sin) + roped_even = roped_even.to(tl.bfloat16).to(tl.float32) + roped_odd = roped_odd.to(tl.bfloat16).to(tl.float32) + even_values = tl.where(rotary_mask, roped_even, even_values) + odd_values = tl.where(rotary_mask, roped_odd, odd_values) + + amax_per_token = tl.maximum( + tl.max(tl.abs(even_values), axis=1), + tl.max(tl.abs(odd_values), axis=1), + ) + global_scale = tl.load(global_scale_ptr) + token_scale = tl.where(amax_per_token > 0.0, amax_per_token / 6.0, 1.0) + if q1_shared_main: + tile_amax = tl.max(amax_per_token, axis=0) + tile_scale = tl.where(tile_amax > 0.0, tile_amax / 6.0, 1.0) + local_scale = tl.full((HP_BLOCK,), tile_scale, tl.float32) + token_scale_col = page_pos // HP_BLOCK + sf_offsets = _fp4_mla_swizzled_sf_offset(safe_all_d, token_scale_col, SF_PER_PAGE) + v_sf_base = tl.cast(local_layer, tl.int64) * tl.cast( + vsf_s0, tl.int64 + ) + physical_page * tl.cast(vsf_s1, tl.int64) + tl.store( + v_sf_ptr + v_sf_base + sf_offsets.to(tl.int64), + tile_scale * global_scale, + mask=mask_all_d & (all_d < V_HEAD_D), + ) + else: + shared_tile = dim_block * FP4_BLOCK < V_HEAD_D + if shared_tile: + tile_amax = tl.max(amax_per_token, axis=0) + tile_scale = tl.where(tile_amax > 0.0, tile_amax / 6.0, 1.0) + local_scale = tl.full((HP_BLOCK,), tile_scale, tl.float32) + token_scale_col = page_pos // HP_BLOCK + sf_offsets = _fp4_mla_swizzled_sf_offset( + safe_all_d, + token_scale_col, + SF_PER_PAGE, + ) + v_sf_base = tl.cast(local_layer, tl.int64) * tl.cast( + vsf_s0, tl.int64 + ) + physical_page * tl.cast(vsf_s1, tl.int64) + tl.store( + v_sf_ptr + v_sf_base + sf_offsets.to(tl.int64), + tile_scale * global_scale, + mask=mask_all_d & (all_d < V_HEAD_D), + ) + else: + local_scale = token_scale + stored_scale = local_scale * global_scale + + quant_scale = local_scale + if STORE_K_RESIDUAL and not q1_shared_main: + if dim_block * FP4_BLOCK >= HEAD_D - K_RESIDUAL_D: + quant_scale = stored_scale.to(tl.float8e4nv).to(tl.float32) / global_scale + low = _fp4_e2m1_quantize(even_values / quant_scale[:, None]) + high = _fp4_e2m1_quantize(odd_values / quant_scale[:, None]) + packed = low | (high << 4) + v_dim_packed = packed + + packed_cols = dim_block * (FP4_BLOCK // 2) + byte_offsets + page_positions = page_pos + token_offsets + kv_base = physical_page * kv_s0 + tl.store( + kv_cache_ptr + kv_base + page_positions[:, None] * kv_s2 + packed_cols[None, :] * kv_s4, + packed, + mask=valid_tokens[:, None] & mask_even_d[None, :], + ) + + if WRITE_V_PACKED: + if q1_shared_main or dim_block * FP4_BLOCK < V_HEAD_D: + # The canonical cache packs dimensions; PV packs tokens. + # Transpose the selected V codes in registers. + v_pairs = tl.reshape( + v_dim_packed.T, + (FP4_BLOCK // 2, HP_BLOCK // 2, 2), + ) + even_token_packed, odd_token_packed = tl.split(v_pairs) + v_low = (even_token_packed & 0x0F) | ((odd_token_packed & 0x0F) << 4) + v_high = ((even_token_packed >> 4) & 0x0F) | (odd_token_packed & 0xF0) + # Keep the two dimension parities in the 8x8 layout produced + # by the token-pair transpose. Interleaving them into a + # 16x8 tensor forces another cross-warp layout conversion; + # two stores address the exact same coalesced rows without + # changing any packed bytes. + v_even_rows = ( + (v_page_offset + physical_page) * V_HEAD_D + + dim_block * FP4_BLOCK + + tl.arange(0, FP4_BLOCK // 2) * 2 + ).to(tl.int64) + pair_offsets = tl.arange(0, HP_BLOCK // 2) + v_cols = page_pos // 2 + pair_offsets + v_even_offsets = v_even_rows[:, None] * v_packed_s0 + v_cols[None, :] * v_packed_s1 + v_odd_offsets = v_even_offsets + v_packed_s0 + token0_valid = block_base_pos + pair_offsets * 2 < kv_len + token1_valid = block_base_pos + pair_offsets * 2 + 1 < kv_len + partial_pair = token0_valid & (~token1_valid) + old_v_low = tl.load( + v_packed_ptr + v_even_offsets, + mask=partial_pair[None, :], + other=0, + ) + old_v_high = tl.load( + v_packed_ptr + v_odd_offsets, + mask=partial_pair[None, :], + other=0, + ) + v_low = tl.where( + partial_pair[None, :], + (old_v_low & 0xF0) | (v_low & 0x0F), + v_low, + ).to(tl.uint8) + v_high = tl.where( + partial_pair[None, :], + (old_v_high & 0xF0) | (v_high & 0x0F), + v_high, + ).to(tl.uint8) + tl.store( + v_packed_ptr + v_even_offsets, + v_low, + mask=token0_valid[None, :], + ) + tl.store( + v_packed_ptr + v_odd_offsets, + v_high, + mask=token0_valid[None, :], + ) + + k_sf_offsets = _fp4_mla_swizzled_sf_offset(page_positions, dim_block, SF_PER_TOKEN) + tl.store( + sf_cache_ptr + physical_page * sf_s0 + k_sf_offsets, + stored_scale, + mask=valid_tokens, + ) + + if STORE_K_RESIDUAL and not q1_shared_main: + residual_start = HEAD_D - K_RESIDUAL_D + if dim_block * FP4_BLOCK >= residual_start: + if MAX_GEN_TILES == 1: + main_even, main_odd = _fp4_mla_unpack_e2m1x2(packed) + main_even *= quant_scale[:, None] + main_odd *= quant_scale[:, None] + else: + main_even = _fp4_e2m1_to_f32(low) * quant_scale[:, None] + main_odd = _fp4_e2m1_to_f32(high) * quant_scale[:, None] + residual_even = even_values - main_even + residual_odd = odd_values - main_odd + residual_amax = tl.maximum( + tl.max(tl.abs(residual_even), axis=1), + tl.max(tl.abs(residual_odd), axis=1), + ) + residual_scale = tl.where(residual_amax > 0.0, residual_amax / 6.0, 1.0) + residual_stored_scale = residual_scale * global_scale + residual_quant_scale = ( + residual_stored_scale.to(tl.float8e4nv).to(tl.float32) / global_scale + ) + residual_low = _fp4_e2m1_quantize(residual_even / residual_quant_scale[:, None]) + residual_high = _fp4_e2m1_quantize(residual_odd / residual_quant_scale[:, None]) + residual_packed = residual_low | (residual_high << 4) + residual_group = dim_block - residual_start // FP4_BLOCK + residual_packed_cols = ( + HEAD_D // 2 + residual_group * (FP4_BLOCK // 2) + byte_offsets + ) + tl.store( + kv_cache_ptr + + kv_base + + page_positions[:, None] * kv_s2 + + residual_packed_cols[None, :] * kv_s4, + residual_packed, + mask=valid_tokens[:, None] & mask_even_d[None, :], + ) + residual_sf_offsets = _fp4_mla_swizzled_sf_offset( + page_positions, + HEAD_D // FP4_BLOCK + residual_group, + SF_PER_TOKEN, + ) + tl.store( + sf_cache_ptr + physical_page * sf_s0 + residual_sf_offsets, + residual_stored_scale, + mask=valid_tokens, + ) + + if FUSE_ROPE_CACHE_STORE: + hp_store_mask = valid_tokens[:, None] & from_latent[:, None] + hp_store_base = ( + hp_pool_ptr + + hp_index * pool_s0 + + local_layer * pool_s1 + + hp_slots[:, None] * HEAD_D + ) + tl.store( + hp_store_base + safe_even_d[None, :], + even_values, + mask=hp_store_mask & mask_even_d[None, :], + ) + tl.store( + hp_store_base + safe_odd_d[None, :], + odd_values, + mask=hp_store_mask & mask_odd_d[None, :], + ) + if tile_iter + 1 < MAX_GEN_TILES: + # The next tile can wrap onto HP slots read by this tile. + # Keep each dimension block's read-before-overwrite order. + tl.debug_barrier() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_mufu16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_mufu16.py new file mode 100644 index 000000000000..ed5303c5977e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_mufu16.py @@ -0,0 +1,9026 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 variant of the packed FP4 MLA CuTe DSL decode kernel (SM107). + +The same kernel interface as ``attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py``; the K3 MLA +layers run it with its own plan and split reduction (``ops/attention/_fp4_mla_decode.py``, +``fp4_mla_packed_reduce``). Differences from the shared kernel: + +- QK and PV are issued by two warps; QK runs two KV tiles ahead of PV, a score bank is refilled + once every producer warp of both CTAs has read it (``SCORE_CONSUMED``), and the PV issuer waits + on ``P_SOURCE_READY`` only. CTA1 issues the TMA of its own half of each K tile. +- K / V / P scale factors live in disjoint TMEM columns per issuing warp; the producers accumulate + the P row sums in FP16x2 per 8-element word and keep the denominator of their 64-column slice + in a register, combined once after the KV loop. +- The unmasked leading KV tiles of a task run software-pipelined by one tile in the producer + warps (row maxima through a two-entry parity buffer in shared memory). +- K / V load warps poll a busy ring slot with ``try_acquire`` + ``nanosleep``. +- A split stores its partial at ``begins[request] + split`` of the plan, so the plan may order the + grid freely (wave and makespan plans, ``SUPPORTS_ORDERED_PLAN``); the partial may be the + normalized FP16 split output (``SPLIT_PARTIAL_DTYPES``). +- Optional programmatic launch: ``launch_dependents`` (the split reduction or a fused consumer + launches early) and ``pdl_wait`` (the decode waits for the kernel that writes its query). + +The kernel carries the internal CuTe DSL build's scheduling hints (``dsl_hints.cfence()`` / +``dsl_hints.warp_switch()``); under the public DSL they emit nothing. +""" + +import math +import sys +import threading +from collections import OrderedDict +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from functools import partial +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass as ctm +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import torch +from cutlass._mlir.dialects import llvm +from cutlass.base_dsl.dsl import BaseDSL +from cutlass.cute.arch.nvvm_wrappers import inline_ptx as cute_inline_ptx +from cutlass.cute.runtime import make_ptr +from cutlass.experimental import cuda as cuda_tma +from cutlass.experimental import primitives as prims + +from . import dsl_hints + + +def _compile_cutedsl(*args: object, **kwargs: object) -> Callable[..., object]: + """Compile with the PyIR frontend the FP4 MLA kernels require.""" + with BaseDSL.enable_pyir(): + return cute.compile(*args, **kwargs) + + +def _current_cu_stream() -> cuda.CUstream: + """The caller's stream, for launch ordering and CUDA graph capture.""" + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + +nvvm_add_packed_f32x2 = partial(prims.add_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_mul_packed_f32x2 = partial(prims.mul_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) +PREPARED_BUFFER_ALIGNMENT_BYTES = 32 +CUDA_GRID_Z_MAX = 65535 +SUPPORTS_FOLDED_QUERIES = True +# Packed decode rows: one batch entry per request holding [query token, head] +# rows, KV split across grid.z tasks, per-row causal bounds (Helix per-token +# bounds or the linear-MTP window), direct output/stats for one split and +# FP32 partials for several. +SUPPORTS_PACKED_QUERIES = True +# Split partial buffers: FP32 holds the numerator relative to the row maximum; +# FP16 holds the normalized split output (numerator / row sum, in output units), +# which the split reduction weighs by the max-scaled row sums. +SPLIT_PARTIAL_DTYPES = (torch.float32, torch.float16) +# A split stores its partial at ``begins[request] + split`` of the plan rather than at its grid +# slot, so the plan kernel may list the tasks in any grid order (the wave and makespan plans of +# fp4_mla_packed_reduce). +SUPPORTS_ORDERED_PLAN = True +# ``launch_dependents`` modes of the packed decode (``True`` is LAUNCH_DEPENDENTS_ENTRY): release +# the programmatic dependent grid at entry, or once the producer warps have finished the last KV +# tile (LAUNCH_DEPENDENTS_EPILOGUE), so that the dependent's prologue overlaps the end of this +# grid without being resident during all of it. +LAUNCH_DEPENDENTS_ENTRY = 1 +LAUNCH_DEPENDENTS_EPILOGUE = 2 +PACKED_MAX_M = 1024 +# Balanced packed split: no task owns fewer KV tiles than this unless a request is shorter. +PACKED_MIN_TILES_PER_TASK = 2 +INT32_MAX = (1 << 31) - 1 + + +@ctm.dsl_user_op +def _as_shared_cta(address, *, loc=None, ip=None): + address_ir = address.ir_value() if hasattr(address, "ir_value") else address + address_space = llvm.PointerType(address_ir.type).address_space + if address_space == ctm.AddressSpace.dsmem: + return prims.cvta_to(address, prims.CvtaSpace.SHARED, loc=loc, ip=ip) + return address + + +@ctm.dsl_user_op +def _mapa_shared_cluster(address, rank, *, loc=None, ip=None): + return prims.mapa(_as_shared_cta(address, loc=loc, ip=ip), rank, loc=loc, ip=ip) + + +@ctm.dsl_user_op +def _mbarrier_arrive_release_cta_shared_cluster(mbar, count=1, *, loc=None, ip=None) -> None: + prims.mbarrier_arrive( + mbar, + count=count, + scope=prims.MemScope.CTA, + loc=loc, + ip=ip, + ) + + +def _scale_words_for_k(k_dim: int, sf_vec_size: int) -> int: + return (k_dim + sf_vec_size * 4 - 1) // (sf_vec_size * 4) + + +def _sfa_tmem_cols_for_scale_chunks(scale_chunks: int, use_128dp_unique: bool) -> int: + if use_128dp_unique: + return (scale_chunks + 3) // 4 * 4 + return scale_chunks * 4 + + +def _derive_qk_scale_burst_kblocks( + *, + qk_mma_kblocks: int, + qk_scale_burst_target_kblocks: int, + scale_a_col_offset: int, + num_tmem_alloc_cols: int, + qk_mma_m_dim: int, + mma_scale_chunks_per_kblock: int, + use_sfa_128dp_unique: bool, + qk_score_slot_n: int, + qk_compact_sfb_cols_per_scale_chunk: int = 0, +) -> int: + target = max(1, qk_scale_burst_target_kblocks) + if qk_compact_sfb_cols_per_scale_chunk > 0: + sfb_cols_per_kblock = mma_scale_chunks_per_kblock * qk_compact_sfb_cols_per_scale_chunk + else: + sfb_cols_per_kblock = _qk_sfb_cols_per_kblock( + qk_mma_m_dim, qk_score_slot_n, mma_scale_chunks_per_kblock + ) + best = 0 + for burst_kblocks in range(1, target + 1): + if qk_mma_kblocks % burst_kblocks != 0: + continue + sfa_cols = _sfa_tmem_cols_for_scale_chunks( + burst_kblocks * mma_scale_chunks_per_kblock, use_sfa_128dp_unique + ) + sfb_cols = burst_kblocks * sfb_cols_per_kblock + scale_b_col_offset = (scale_a_col_offset + sfa_cols + 15) // 16 * 16 + qk_scale_feed_col_end = max(scale_a_col_offset + sfa_cols, scale_b_col_offset + sfb_cols) + if qk_scale_feed_col_end <= num_tmem_alloc_cols: + best = burst_kblocks + if best == 0: + raise ValueError( + "QK scale feed cannot fit even one K64 block in the TMEM tail; " + f"scale_a_col_offset={scale_a_col_offset}, " + f"num_tmem_alloc_cols={num_tmem_alloc_cols}" + ) + return best + + +def _compact_sfb_mma_cols_per_scale_chunk(m_dim: int, n_dim: int) -> int: + if m_dim == 128: + if n_dim > 256: + raise ValueError(f"Unsupported M=128 compact SFB N dimension {n_dim}") + return 2 * ((n_dim + 127) // 128) + return 2 * ((n_dim + 63) // 64) + + +def _compact_sfb_copy_cols_per_scale_chunk(m_dim: int, n_dim: int) -> int: + if m_dim == 128: + if n_dim > 256: + raise ValueError(f"Unsupported M=128 compact SFB N dimension {n_dim}") + return 4 * ((n_dim + 127) // 128) + return _compact_sfb_mma_cols_per_scale_chunk(m_dim, n_dim) + + +def _qk_compact_sfb_cols_per_scale_chunk(m_dim: int, n_dim: int, warpx2_mode: int) -> int: + if warpx2_mode != 0 and m_dim == 128: + return _compact_sfb_mma_cols_per_scale_chunk(m_dim, n_dim) + return _compact_sfb_copy_cols_per_scale_chunk(m_dim, n_dim) + + +def _qk_sfb_cols_per_kblock(m_dim: int, n_dim: int, mma_scale_chunks_per_kblock: int) -> int: + return mma_scale_chunks_per_kblock * _compact_sfb_copy_cols_per_scale_chunk(m_dim, n_dim) + + +def _qk_sfb_smem_nchunks_per_cta(m_dim: int, n_dim: int) -> int: + if m_dim == 128: + if n_dim > 256: + raise ValueError(f"Unsupported M=128 compact SFB N dimension {n_dim}") + return (n_dim + 127) // 128 + + +def _mma_sf_ids_per_scale_atom(k_dim: int, sf_vec_size: int) -> int: + return max(1, sf_vec_size * 4 // k_dim) + + +def _mma_scale_vec_size_enum(sf_vec_size: int): + if sf_vec_size == 16: + return prims.Tcgen05MMAScaleVecSize.BLOCK16 + raise ValueError(f"Unsupported FP4 scale vector size {sf_vec_size}") + + +def _mma_scale_format(sf_vec_size: int) -> int: + if sf_vec_size == 16: + return 0 + raise ValueError(f"Unsupported FP4 scale vector size {sf_vec_size}") + + +def _mma_block_scale_kind(sf_vec_size: int): + if sf_vec_size == 16: + return prims.MMABlockScaleKind.MXF4NVF4 + raise ValueError(f"Unsupported FP4 scale vector size {sf_vec_size}") + + +def _initial_sf_vec_size_from_argv(options: Sequence[int], default: int) -> int: + args = sys.argv[1:] + for idx, arg in enumerate(args): + value_text: str | None = None + if arg == "--sf-vec-size": + if idx + 1 >= len(args): + raise ValueError("--sf-vec-size requires an integer value") + value_text = args[idx + 1] + elif arg.startswith("--sf-vec-size="): + value_text = arg.split("=", 1)[1] + if value_text is None: + continue + try: + value = int(value_text) + except ValueError as exc: + raise ValueError(f"--sf-vec-size must be an integer, got {value_text!r}") from exc + if value not in options: + raise ValueError(f"--sf-vec-size must be one of {options}, got {value}") + return value + return default + + +def _initial_float_from_argv(option: str, default: float) -> float: + args = sys.argv[1:] + for idx, arg in enumerate(args): + value_text: str | None = None + if arg == option: + if idx + 1 >= len(args): + raise ValueError(f"{option} requires a floating-point value") + value_text = args[idx + 1] + elif arg.startswith(f"{option}="): + value_text = arg.split("=", 1)[1] + if value_text is None: + continue + try: + return float(value_text) + except ValueError as exc: + raise ValueError(f"{option} must be a float, got {value_text!r}") from exc + return default + + +M_TILE = 256 +KV_TILE = 256 +KV_TILES = 4 +KV_TOTAL = KV_TILE * KV_TILES +OUT_DIM = 256 +OUT_DIM_TOTAL = 512 +SMEM_P4_N_OUT_TILES = OUT_DIM_TOTAL // OUT_DIM +QK_LOGICAL_DIM = 640 +QK_PADDED_DIM = QK_LOGICAL_DIM +TRTLLM_K_HEAD_DIM = 576 +TRTLLM_V_HEAD_DIM = 512 +TRTLLM_Q_RESIDUAL_DIM = 64 +TRTLLM_K_RESIDUAL_DIM = TRTLLM_Q_RESIDUAL_DIM +TRTLLM_K_STORAGE_DIM = TRTLLM_K_HEAD_DIM + TRTLLM_K_RESIDUAL_DIM +QK_EFFECTIVE_DIM = QK_LOGICAL_DIM + TRTLLM_K_RESIDUAL_DIM +TRTLLM_PAGE_SIZE = 128 +TRTLLM_K_SF_GROUPS = TRTLLM_K_STORAGE_DIM // 16 +QK_PREFIX_DIM = TRTLLM_V_HEAD_DIM +TRTLLM_QSF_EVEN_PRMT_SELECTOR = 25632 +TRTLLM_QSF_ODD_PRMT_SELECTOR = 30001 +SF_VEC_SIZE_OPTIONS = (16,) +SF_VEC_SIZE = _initial_sf_vec_size_from_argv(SF_VEC_SIZE_OPTIONS, 16) +MMA_KBLOCK_DIM = 64 +MMA_KBLOCK_IDESC_K_DIM = 0 +MMA_SCALE_CHUNK_K = SF_VEC_SIZE * 4 +MMA_SCALE_CHUNKS_PER_KBLOCK = _scale_words_for_k(MMA_KBLOCK_DIM, SF_VEC_SIZE) +MMA_SCALE_VEC_SIZE = _mma_scale_vec_size_enum(SF_VEC_SIZE) +MMA_SCALE_FORMAT = _mma_scale_format(SF_VEC_SIZE) +MMA_BLOCK_SCALE_KIND = _mma_block_scale_kind(SF_VEC_SIZE) +QK_MMA_KBLOCKS = QK_LOGICAL_DIM // MMA_KBLOCK_DIM +QK_EFFECTIVE_MMA_KBLOCKS = QK_EFFECTIVE_DIM // MMA_KBLOCK_DIM +QK_TAIL_EFFECTIVE_MMA_KBLOCKS = QK_EFFECTIVE_MMA_KBLOCKS - QK_PREFIX_DIM // MMA_KBLOCK_DIM +FIRST_ORDER_TAIL_A_LOCAL_KBLOCKS = (0, 1, 0) +FIRST_ORDER_TAIL_B_LOCAL_KBLOCKS = (0, 0, 1) +QK_SF_GROUPS = QK_PADDED_DIM // SF_VEC_SIZE +LOG2_6 = math.log2(6.0) +LOG2_E = 1.4426950408889634 +# Keep the PV-only P scale above E4M3's underflow range. The matching output +# normalization below cancels this factor without changing the attention math. +FP4_MLA_E4M3_MAX_FINITE = 448.0 +FP4_MLA_P_GLOBAL_SCALE = FP4_MLA_E4M3_MAX_FINITE * 6.0 +SMEM_P4_PAGE_PLAN_PROFILE_KV = 160 * 1024 +SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT = 4 +SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT = SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT * TRTLLM_PAGE_SIZE +# Keep the device-side ceil division ``valid_k + KV_TILE - 1`` in Int32 range. +SMEM_P4_RUNTIME_MAX_KV = ( + (INT32_MAX - (KV_TILE - 1)) // SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT +) * SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT +SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_PAGE_PLAN_PROFILE_KV // TRTLLM_PAGE_SIZE +SMEM_P4_PAGE_ID_PLAN_BYTES = SMEM_P4_PAGE_ID_PLAN_INTS * ctm.Int32.bytes +# softmax scale, PV rescale, producer mask bound (Int32 bits), pad. +SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS = 4 +SMEM_P4_RUNTIME_MASK_BOUND_WORD = 2 +SMEM_P4_RUNTIME_SCALE_PAIR_BYTES = SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS * ctm.Float32.bytes +SMEM_P4_LAZY_ANCHOR_REBASE_LOG2 = _initial_float_from_argv("--lazy-anchor-threshold-log2", 1.0) +if not 0.0 <= SMEM_P4_LAZY_ANCHOR_REBASE_LOG2 <= 8.0: + raise ValueError( + f"--lazy-anchor-threshold-log2 must be in [0, 8], got {SMEM_P4_LAZY_ANCHOR_REBASE_LOG2}" + ) +SMEM_P4_PIPELINE_DIRECT_TARGET_PAGE_SIZE = 128 +RUBIN_TMEM_CAPACITY_KB = 288 +TMEM_BYTES_PER_COL = 512 +NUM_TMEM_ALLOC_COLS = RUBIN_TMEM_CAPACITY_KB * 1024 // TMEM_BYTES_PER_COL +SMEM_P4_BMM1_MNK = (64, KV_TILE, QK_EFFECTIVE_DIM) +SMEM_P4_BMM2_MNK = (64, OUT_DIM, KV_TILE) +SMEM_P4_TMEM_WARP_SHAPE_MN = (2, 2) +SMEM_P4_BMM1_M = SMEM_P4_BMM1_MNK[0] +SMEM_P4_BMM1_N = SMEM_P4_BMM1_MNK[1] +SMEM_P4_BMM2_M = SMEM_P4_BMM2_MNK[0] +SMEM_P4_BMM2_N = SMEM_P4_BMM2_MNK[1] +SMEM_P4_BMM2_K = SMEM_P4_BMM2_MNK[2] +SMEM_P4_TMEM_WARP_M = SMEM_P4_TMEM_WARP_SHAPE_MN[0] +SMEM_P4_TMEM_WARP_N = SMEM_P4_TMEM_WARP_SHAPE_MN[1] +SMEM_P4_TMEM_BANKS = 2 +SMEM_P4_SCORE_SLOTS_PER_KV_TILE = 2 +SMEM_P4_SCORE_SLOT_N = SMEM_P4_BMM1_N // SMEM_P4_SCORE_SLOTS_PER_KV_TILE +SMEM_P4_BMM2_STAGE_K = SMEM_P4_BMM2_K // SMEM_P4_SCORE_SLOTS_PER_KV_TILE +SMEM_P4_CTA_GROUP_M = SMEM_P4_BMM1_M * SMEM_P4_TMEM_WARP_M +SMEM_P4_PIPELINE_DIRECT_TARGET_M_OPTIONS = (SMEM_P4_CTA_GROUP_M, M_TILE) +SMEM_P4_SFB_ATOM_ROWS = 128 +SMEM_P4_P_ARCHIVE_SLOTS = SMEM_P4_TMEM_BANKS +SMEM_P4_ROWMETA_SLOTS = SMEM_P4_TMEM_BANKS +SMEM_P4_QK_PIPELINE_SLOTS = SMEM_P4_TMEM_BANKS +SMEM_P4_BMM1_ACC_TMEM_COLS = SMEM_P4_BMM1_N // SMEM_P4_TMEM_WARP_N +SMEM_P4_SCORE_HALF_TMEM_COLS = SMEM_P4_BMM1_ACC_TMEM_COLS // SMEM_P4_SCORE_SLOTS_PER_KV_TILE +SMEM_P4_QK_SCORE_SLOT_N = SMEM_P4_BMM1_N +SMEM_P4_QK_SCORE_SLOTS = SMEM_P4_TMEM_BANKS +SMEM_P4_QK_SCORE_RING_SLOTS = SMEM_P4_TMEM_BANKS +SMEM_P4_QK_PIPELINE_SLOT_STRIDE = SMEM_P4_QK_SCORE_SLOT_N // SMEM_P4_TMEM_WARP_N +SMEM_P4_QK_SCORE_SLOT_STRIDE = SMEM_P4_SCORE_HALF_TMEM_COLS +SMEM_P4_QK_MMA_KBLOCK_DIM = MMA_KBLOCK_DIM +SMEM_P4_QK_MMA_KBLOCK_IDESC_K_DIM = MMA_KBLOCK_IDESC_K_DIM +SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK = MMA_SCALE_CHUNKS_PER_KBLOCK +SMEM_P4_QK_MMA_KBLOCKS = QK_MMA_KBLOCKS +SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS_PER_WARP = 32 * 2 +SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS_PER_SLOT = ( + SMEM_P4_TMEM_WARP_M * SMEM_P4_TMEM_WARP_N * SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS_PER_WARP +) +SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS = ( + SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS_PER_SLOT +) +SMEM_P4_P4_SMEM_ROW_STRIDE_BYTES = SMEM_P4_BMM2_STAGE_K // 2 +SMEM_P4_P4_SMEM_ROWS = SMEM_P4_BMM2_M +SMEM_P4_P4_SMEM_STAGE_BYTES = SMEM_P4_P4_SMEM_ROWS * SMEM_P4_P4_SMEM_ROW_STRIDE_BYTES +SMEM_P4_P4_SMEM_STAGE_WORDS = SMEM_P4_P4_SMEM_STAGE_BYTES // 4 +SMEM_P4_P4_PIPELINE_STAGES = SMEM_P4_TMEM_BANKS +SMEM_P4_P4_CONTROL_MBAR_STAGES = SMEM_P4_TMEM_BANKS * SMEM_P4_SCORE_SLOTS_PER_KV_TILE +SMEM_P4_P4_SMEM_PHYSICAL_STAGES = SMEM_P4_TMEM_BANKS * SMEM_P4_SCORE_SLOTS_PER_KV_TILE +SMEM_P4_P4_SMEM_WORDS = SMEM_P4_P4_SMEM_PHYSICAL_STAGES * SMEM_P4_P4_SMEM_STAGE_WORDS +SMEM_P4_P4_SMEM_MMA_LEADING_BYTE_OFFSET = 0 +SMEM_P4_P4_SMEM_MMA_STRIDE_BYTE_OFFSET = 512 +SMEM_P4_P4_SMEM_MMA_LAYOUT = 4 +SMEM_P4_SCORE_SLOT_ROW_STATE_STRIDE = SMEM_P4_BMM1_M * SMEM_P4_TMEM_WARP_N +SMEM_P4_P_TMEM_KBLOCK_COLS = MMA_KBLOCK_DIM // 2 // 4 +SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT = SMEM_P4_BMM2_STAGE_K // MMA_KBLOCK_DIM +SMEM_P4_PV_KBLOCKS_PER_LI = SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT +SMEM_P4_P_TMEM_STAGE_COLS = SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT * SMEM_P4_P_TMEM_KBLOCK_COLS +SMEM_P4_P_TMEM_SLOT_STRIDE = SMEM_P4_P_TMEM_STAGE_COLS +SMEM_P4_SCORE_TMEM_COLS = ( + 0, + SMEM_P4_QK_SCORE_RING_SLOTS * SMEM_P4_QK_PIPELINE_SLOT_STRIDE, +) +SMEM_P4_SCORE_PONG_TMEM_COLS = (SMEM_P4_SCORE_TMEM_COLS[1], 256) +SMEM_P4_O_ACC_TMEM_COLS = ( + SMEM_P4_SCORE_PONG_TMEM_COLS[1], + SMEM_P4_SCORE_PONG_TMEM_COLS[1] + SMEM_P4_BMM2_N, +) +SMEM_P4_O_ACC_TILE_COLS = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N +SMEM_P4_O_ACC_COL_OFFSETS = tuple( + (SMEM_P4_O_ACC_TMEM_COLS[0] + t * SMEM_P4_O_ACC_TILE_COLS for t in range(SMEM_P4_N_OUT_TILES)) +) +SMEM_P4_SFA_LAYOUT = 1 +SMEM_P4_SFA_COLS_PER_SCALE_CHUNK = 1 +SMEM_P4_P_SFA_LAYOUT = 1 +SMEM_P4_P_SFA_COLS_PER_SCALE_CHUNK = 1 +SMEM_P4_SFA_UNIQUE_CHUNKS_PER_UTCCP = 4 +SMEM_P4_QK_B_MMA_LEADING_BYTES = 16 +SMEM_P4_QK_B_MMA_STRIDE_BYTES = 1024 +SMEM_P4_QK_B_MMA_LAYOUT = 2 +SMEM_P4_QK_B_MMA_SMEM_WORD_DELTA = 0 +SMEM_P4_QK_B_MMA_K_SEGMENT_OFFSET = 0 +SMEM_P4_QK_SFB_MMA_COL_DELTA = 0 +SMEM_P4_QK_SFB_EXTRA_ROW_STRIDE = 0 +SMEM_P4_QK_SFB_EXTRA_COL_STRIDE = -1 +SMEM_P4_QK_SFB_S2T_STRIDE_BYTES = 128 +SMEM_P4_QK_SFB_CP_GROUP_ONE = False +SMEM_P4_QK_SFB_WARPX2_MODE = 1 +SMEM_P4_P4_SCORE_READ_GROUP1_COL = 16 +SMEM_P4_P4_SCORE_READ_GROUP2_COL = 32 +SMEM_P4_P_SCALE_FEED_STAGE_COLS = _sfa_tmem_cols_for_scale_chunks( + SMEM_P4_BMM2_K // MMA_KBLOCK_DIM * MMA_SCALE_CHUNKS_PER_KBLOCK, True +) +SMEM_P4_P_SCALE_A_COL_OFFSET = SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_P_TMEM_STAGE_COLS +SMEM_P4_PV_SFB_SLOT_COL_OFFSET = SMEM_P4_P_SCALE_A_COL_OFFSET + SMEM_P4_P_SCALE_FEED_STAGE_COLS +SMEM_P4_PV_SFB_WARPX2_MODE = 1 +SMEM_P4_PV_SFB_N_CHUNKS = (SMEM_P4_BMM2_N + 127) // 128 +SMEM_P4_PV_COMPACT_SFB_COLS_PER_SCALE_CHUNK = _qk_compact_sfb_cols_per_scale_chunk( + SMEM_P4_CTA_GROUP_M, SMEM_P4_BMM2_N, SMEM_P4_PV_SFB_WARPX2_MODE +) +SMEM_P4_PV_SFB_COLS_PER_KBLOCK = ( + MMA_SCALE_CHUNKS_PER_KBLOCK * SMEM_P4_PV_COMPACT_SFB_COLS_PER_SCALE_CHUNK +) +SMEM_P4_PV_SFB_SCORE_SLOT_COLS = SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT * SMEM_P4_PV_SFB_COLS_PER_KBLOCK +SMEM_P4_P_ARCHIVE_TMEM_COLS = SMEM_P4_SCORE_TMEM_COLS +SMEM_P4_P_ARCHIVE_SLOT_STRIDE = SMEM_P4_QK_PIPELINE_SLOT_STRIDE +SMEM_P4_SCALE_FEED_STAGE_COLS = 28 +SMEM_P4_SCALE_FEED_TMEM_COLS = ( + SMEM_P4_O_ACC_TMEM_COLS[1], + SMEM_P4_O_ACC_TMEM_COLS[1] + SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_SCALE_FEED_STAGE_COLS, +) +SMEM_P4_QK_ACC_COL_OFFSET = SMEM_P4_SCORE_TMEM_COLS[0] +# P scale factors per score bank, outside the score columns so the next QK into +# the bank cannot overwrite a pending PV's scales. +SMEM_P4_PSF_TMEM_COL_OFFSET = SMEM_P4_SCALE_FEED_TMEM_COLS[1] +if SMEM_P4_PSF_TMEM_COL_OFFSET + SMEM_P4_TMEM_BANKS * SMEM_P4_P_SCALE_FEED_STAGE_COLS > ( + NUM_TMEM_ALLOC_COLS +): + raise ValueError("P scale-factor banks do not fit in the TMEM allocation") +SMEM_P4_O_ACC_COL_OFFSET = SMEM_P4_O_ACC_TMEM_COLS[0] +SMEM_P4_SCALE_A_COL_OFFSET = SMEM_P4_SCALE_FEED_TMEM_COLS[0] +SMEM_P4_QK_COMPACT_SFB_S2T_BYTE_OFFSET = 0 +SMEM_P4_QK_COMPACT_SFB_COLS_PER_SCALE_CHUNK = _qk_compact_sfb_cols_per_scale_chunk( + SMEM_P4_CTA_GROUP_M, SMEM_P4_QK_SCORE_SLOT_N, SMEM_P4_QK_SFB_WARPX2_MODE +) +SMEM_P4_QK_SFB_COLS_PER_KBLOCK = ( + SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK * SMEM_P4_QK_COMPACT_SFB_COLS_PER_SCALE_CHUNK +) +SMEM_P4_QK_RESIDENT_SF_CHUNKS = _scale_words_for_k(QK_PADDED_DIM, SF_VEC_SIZE) +SMEM_P4_QK_RESIDENT_SF_COLS = SMEM_P4_QK_RESIDENT_SF_CHUNKS +SMEM_P4_QK_SCALE_BURST_KBLOCKS = _derive_qk_scale_burst_kblocks( + qk_mma_kblocks=QK_MMA_KBLOCKS, + qk_scale_burst_target_kblocks=QK_MMA_KBLOCKS, + scale_a_col_offset=SMEM_P4_SCALE_A_COL_OFFSET, + num_tmem_alloc_cols=NUM_TMEM_ALLOC_COLS, + qk_mma_m_dim=SMEM_P4_CTA_GROUP_M, + mma_scale_chunks_per_kblock=SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK, + use_sfa_128dp_unique=True, + qk_score_slot_n=SMEM_P4_QK_SCORE_SLOT_N, + qk_compact_sfb_cols_per_scale_chunk=SMEM_P4_QK_COMPACT_SFB_COLS_PER_SCALE_CHUNK, +) +SMEM_P4_QK_SFA_FEED_COLS = SMEM_P4_QK_RESIDENT_SF_COLS +SMEM_P4_SCALE_B_COL_OFFSET = (SMEM_P4_SCALE_A_COL_OFFSET + SMEM_P4_QK_SFA_FEED_COLS + 15) // 16 * 16 +SMEM_P4_PV_VSF_TAIL_COL_OFFSET = SMEM_P4_SCALE_B_COL_OFFSET +SMEM_P4_RESIDENT_Q_SF_COL_OFFSET = SMEM_P4_SCALE_FEED_TMEM_COLS[0] +SMEM_P4_QK_SFA_COL_OFFSET = SMEM_P4_RESIDENT_Q_SF_COL_OFFSET +SMEM_P4_OLD_SCALE_TMEM_SLOTS = SMEM_P4_ROWMETA_SLOTS +SMEM_P4_STAGE_ROWSUM_TMEM_COLS = SMEM_P4_TMEM_BANKS * SMEM_P4_SCORE_SLOTS_PER_KV_TILE +SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS = ( + SMEM_P4_TMEM_BANKS * SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_SCORE_SLOT_ROW_STATE_STRIDE +) +SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS = SMEM_P4_BMM1_M +SMEM_P4_TMEM_SCORE_PIPELINE_STAGES = SMEM_P4_QK_PIPELINE_SLOTS +SMEM_P4_MBAR_PARITY_PHASES = 2 +SMEM_P4_PAIR_CORR_MBAR_OFFSET = 0 +SMEM_P4_PAIR_P4_READY_MBAR_OFFSET = ( + SMEM_P4_PAIR_CORR_MBAR_OFFSET + SMEM_P4_N_OUT_TILES * SMEM_P4_QK_PIPELINE_SLOTS +) +SMEM_P4_PAIR_ROWSUM_DONE_MBAR_OFFSET = SMEM_P4_PAIR_P4_READY_MBAR_OFFSET + SMEM_P4_QK_PIPELINE_SLOTS +SMEM_P4_STAGE_ROWSUM_RING_SLOTS = SMEM_P4_STAGE_ROWSUM_TMEM_COLS +SMEM_P4_P_HALF_COUNT = SMEM_P4_PV_KBLOCKS_PER_LI // SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT +SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET = ( + SMEM_P4_PAIR_ROWSUM_DONE_MBAR_OFFSET + SMEM_P4_STAGE_ROWSUM_RING_SLOTS +) +SMEM_P4_PAIR_P_HALF_READY_MBAR_OFFSET = ( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + SMEM_P4_P_HALF_COUNT * SMEM_P4_QK_PIPELINE_SLOTS +) +SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET = SMEM_P4_PAIR_P_HALF_READY_MBAR_OFFSET +SMEM_P4_PAIR_FINAL_PV_ISSUED_MBAR_OFFSET = ( + SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + SMEM_P4_QK_PIPELINE_SLOTS +) +SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET = ( + SMEM_P4_PAIR_P_HALF_READY_MBAR_OFFSET + SMEM_P4_P_HALF_COUNT * SMEM_P4_QK_PIPELINE_SLOTS +) +SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET = SMEM_P4_PAIR_ROWSUM_DONE_MBAR_OFFSET +# Per score bank: all producer warps of both CTAs have read the bank's scores, so +# the MMA warp may already issue the QK two tiles ahead into it. +SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET = SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET +SMEM_P4_QK15_STAGE1_MBAR_OFFSET = ( + SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET + SMEM_P4_QK_PIPELINE_SLOTS +) +SMEM_P4_PAIR_OVERLAP_MBAR_COUNT = ( + SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + SMEM_P4_QK_PIPELINE_SLOTS +) +SMEM_P4_V_PIPELINE_STAGES = 3 +# Poll period of the K/V load warps while their ring slot is still in use. +SMEM_P4_LOAD_WAIT_SLEEP_NS = 128 +SMEM_P4_PV_PAGE_B_MMA_LEADING_BYTES = 0 +SMEM_P4_PV_PAGE_B_MMA_STRIDE_BYTES = 512 +SMEM_P4_PV_PAGE_B_MMA_LAYOUT = 4 +SMEM_P4_PV_B_MMA_K_SEGMENT_OFFSET = 0 +SMEM_P4_PV_SFB_MMA_COL_DELTA = 0 +SMEM_P4_PV_SFB_EXTRA_ROW_STRIDE = 0 +SMEM_P4_PV_SFB_EXTRA_COL_STRIDE = -1 +SMEM_P4_PV_SFB_S2T_STRIDE_BYTES = 128 +SMEM_P4_PV_SFB_CP_GROUP_ONE = False +SMEM_P4_QK_COMPLETION_MBARS = SMEM_P4_TMEM_SCORE_PIPELINE_STAGES + + +@dataclass(frozen=True) +class KvCache3DLayout: + num_pages: int + packed_dim: int + stride_page: int + stride_token: int + stride_packed_dim: int + + +SMEM_P4_PRODUCER_WARPS_PER_SCORE_SLOT = SMEM_P4_TMEM_WARP_M * SMEM_P4_TMEM_WARP_N +SMEM_P4_COMPUTE_WARPS = SMEM_P4_PRODUCER_WARPS_PER_SCORE_SLOT +SMEM_P4_CORRECTION_WARPS = SMEM_P4_PRODUCER_WARPS_PER_SCORE_SLOT +SMEM_P4_COMPUTE_WARP_ID_BEGIN = 0 +SMEM_P4_CORRECTION_WARP_ID_BEGIN = SMEM_P4_COMPUTE_WARP_ID_BEGIN + SMEM_P4_COMPUTE_WARPS +SMEM_P4_TOTAL_PRODUCER_WARPS = SMEM_P4_COMPUTE_WARPS + SMEM_P4_CORRECTION_WARPS +SMEM_P4_STORE_WARPS = SMEM_P4_CORRECTION_WARPS +SMEM_P4_COMPUTE_THREADS = 32 * SMEM_P4_COMPUTE_WARPS +SMEM_P4_CORRECTION_THREADS = 32 * SMEM_P4_CORRECTION_WARPS +SMEM_P4_SCORE_SLOT_PRODUCER_THREADS = SMEM_P4_COMPUTE_THREADS +THREADS_PER_CTA = 32 * 12 +SMEM_P4_TMEM_ALLOC_EXCLUSIVE = True +MMA_TILER_MNK = (SMEM_P4_CTA_GROUP_M, SMEM_P4_BMM2_N, SMEM_P4_BMM2_K) +QK_DATA_STAGE_K_DIM = MMA_TILER_MNK[2] +QK_TAIL_DATA_STAGE_K_DIM = QK_LOGICAL_DIM % QK_DATA_STAGE_K_DIM +SMEM_P4_QK_FULL_DATA_STAGES = QK_LOGICAL_DIM // QK_DATA_STAGE_K_DIM +SMEM_P4_QK_TAIL_STAGE_KBLOCK_START = SMEM_P4_QK_FULL_DATA_STAGES * ( + QK_DATA_STAGE_K_DIM // MMA_KBLOCK_DIM +) +QK_K_TAIL_DATA_STAGE_K_DIM = ( + TRTLLM_K_STORAGE_DIM - SMEM_P4_QK_TAIL_STAGE_KBLOCK_START * MMA_KBLOCK_DIM +) +MMA_TILER_MNK_PER_CTA = (SMEM_P4_BMM1_M, SMEM_P4_BMM2_N, SMEM_P4_BMM2_K) +MMA_KBLOCKS_PER_TILE = MMA_TILER_MNK[2] // MMA_KBLOCK_DIM +SMEM_P4_QK_DATA_STAGE_KBLOCKS = QK_DATA_STAGE_K_DIM // MMA_KBLOCK_DIM +SMEM_P4_QK_DATA_STAGES = ( + QK_MMA_KBLOCKS + SMEM_P4_QK_DATA_STAGE_KBLOCKS - 1 +) // SMEM_P4_QK_DATA_STAGE_KBLOCKS +SMEM_P4_QK_SMEM_PIPELINE_SLOTS = 3 +SMEM_P4_QK_BARRIER_STAGES = SMEM_P4_QK_SMEM_PIPELINE_SLOTS +SMEM_P4_QK_KTILE_RING_STAGES = SMEM_P4_QK_SMEM_PIPELINE_SLOTS +SMEM_P4_QK_BULK_READY_MBARS = SMEM_P4_QK_SMEM_PIPELINE_SLOTS +CLUSTER_SHAPE_MNK = (2, 1, 1) +MMA_SCALE_CHUNKS_PER_TILE = _scale_words_for_k(MMA_TILER_MNK[2], SF_VEC_SIZE) +SMEM_P4_V_N_PER_CTA = SMEM_P4_BMM2_N // CLUSTER_SHAPE_MNK[0] +NUM_QK_Q_STAGE = SMEM_P4_QK_DATA_STAGES +NUM_QK_Q_PIPELINE_STAGE = SMEM_P4_QK_DATA_STAGES +NUM_QK_AB_STAGE = SMEM_P4_QK_DATA_STAGES +MMA_WARP_ID = 8 +TMA_QK_WARP_ID = 9 +TMA_V_WARP_ID = 10 +ROWMETA_WARP_ID = 11 +# QK and PV MMAs are issued by two warps on different SMSPs, so their B scale +# factors use disjoint TMEM columns: a ring holding one QK data stage of K scales +# (4 kblocks) and one PV N tile of V scales (4 kblocks). Each issuer reuses its +# columns in issue order. +PV_ISSUE_WARP_ID = ROWMETA_WARP_ID +SMEM_P4_KSF_RING_COL_OFFSET = SMEM_P4_SCALE_B_COL_OFFSET +SMEM_P4_KSF_RING_COLS = SMEM_P4_QK_DATA_STAGE_KBLOCKS * SMEM_P4_QK_SFB_COLS_PER_KBLOCK +SMEM_P4_VSF_TMEM_COL_OFFSET = SMEM_P4_KSF_RING_COL_OFFSET + SMEM_P4_KSF_RING_COLS +SMEM_P4_VSF_TMEM_COLS = SMEM_P4_PV_KBLOCKS_PER_LI * SMEM_P4_PV_SFB_COLS_PER_KBLOCK +if SMEM_P4_VSF_TMEM_COL_OFFSET + SMEM_P4_VSF_TMEM_COLS > SMEM_P4_PSF_TMEM_COL_OFFSET: + raise ValueError("K/V scale-factor rings overlap the P scale-factor banks") +SMEM_P4_ALU_REGISTER_BUDGET = 224 +SMEM_P4_AUX_REGISTER_BUDGET = 56 +A_STAGE_BYTES = MMA_TILER_MNK_PER_CTA[0] * MMA_TILER_MNK[2] // 2 +B_STAGE_BYTES = SMEM_P4_V_N_PER_CTA * MMA_TILER_MNK[2] // 2 +SFA_STAGE_BYTES = 512 * _scale_words_for_k(MMA_TILER_MNK[2], min(SF_VEC_SIZE_OPTIONS)) +SFB_STAGE_BYTES = SFA_STAGE_BYTES * 2 +SFA_TMA_STAGE_BYTES = 512 * MMA_SCALE_CHUNKS_PER_TILE +SMEM_P4_QK_N_PER_CTA = SMEM_P4_QK_SCORE_SLOT_N // CLUSTER_SHAPE_MNK[0] +QK_A_STAGE_BYTES = MMA_TILER_MNK_PER_CTA[0] * QK_DATA_STAGE_K_DIM // 2 +QK_B_STAGE_BYTES = SMEM_P4_QK_N_PER_CTA * QK_DATA_STAGE_K_DIM // 2 +QK_TAIL_A_STAGE_BYTES = MMA_TILER_MNK_PER_CTA[0] * QK_TAIL_DATA_STAGE_K_DIM // 2 +QK_TAIL_B_STAGE_BYTES = SMEM_P4_QK_N_PER_CTA * QK_K_TAIL_DATA_STAGE_K_DIM // 2 +SMEM_P4_QK_DATA_STAGE_SF_CHUNKS = _scale_words_for_k(QK_DATA_STAGE_K_DIM, min(SF_VEC_SIZE_OPTIONS)) +QK_SFA_STAGE_BYTES = 512 * SMEM_P4_QK_DATA_STAGE_SF_CHUNKS +SMEM_P4_QK_SFB_N_CHUNKS = _qk_sfb_smem_nchunks_per_cta(SMEM_P4_CTA_GROUP_M, SMEM_P4_QK_SCORE_SLOT_N) +QK_SFB_STAGE_BYTES = QK_SFA_STAGE_BYTES * SMEM_P4_QK_SFB_N_CHUNKS +QK_SFA_TMA_STAGE_BYTES = QK_SFA_STAGE_BYTES +QK_SFB_TMA_STAGE_BYTES = QK_SFB_STAGE_BYTES +SMEM_P4_QK_TAIL_DATA_STAGE_SF_CHUNKS = _scale_words_for_k( + QK_TAIL_DATA_STAGE_K_DIM, min(SF_VEC_SIZE_OPTIONS) +) +SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS = _scale_words_for_k( + QK_K_TAIL_DATA_STAGE_K_DIM, min(SF_VEC_SIZE_OPTIONS) +) +QK_TAIL_SFA_STAGE_BYTES = 512 * SMEM_P4_QK_TAIL_DATA_STAGE_SF_CHUNKS +QK_TAIL_SFB_STAGE_BYTES = 512 * SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS * SMEM_P4_QK_SFB_N_CHUNKS +QK_FULL_MMA_LEADING_BYTE_OFFSET = 16 +QK_FULL_MMA_STRIDE_BYTE_OFFSET = 1024 +QK_FULL_MMA_LAYOUT = 2 +QK_TAIL_MMA_LEADING_BYTE_OFFSET = 0 +QK_TAIL_MMA_STRIDE_BYTE_OFFSET = 512 +QK_TAIL_MMA_LAYOUT = 4 +QK_K_TAIL_MMA_LEADING_BYTE_OFFSET = 0 +QK_K_TAIL_MMA_STRIDE_BYTE_OFFSET = 512 +QK_K_TAIL_MMA_LAYOUT = 4 +QK_FULL_TMA_SWIZZLE = cuda_tma.TensorMapSwizzle.s128b +QK_TAIL_TMA_SWIZZLE = cuda_tma.TensorMapSwizzle.s64b +QK_K_TAIL_TMA_SWIZZLE = cuda_tma.TensorMapSwizzle.s64b +SMEM_P4_QK_TMA_CHUNKS_PER_CTA = ( + SMEM_P4_QK_N_PER_CTA + SMEM_P4_PIPELINE_DIRECT_TARGET_PAGE_SIZE - 1 +) // SMEM_P4_PIPELINE_DIRECT_TARGET_PAGE_SIZE +SMEM_P4_QK_TMA_N_PER_CHUNK = ( + SMEM_P4_QK_N_PER_CTA + SMEM_P4_QK_TMA_CHUNKS_PER_CTA - 1 +) // SMEM_P4_QK_TMA_CHUNKS_PER_CTA +SMEM_P4_PAGES_PER_KV_TILE = KV_TILE // TRTLLM_PAGE_SIZE +SMEM_P4_SCALE_CHUNKS_PER_PAGE = TRTLLM_PAGE_SIZE // MMA_SCALE_CHUNK_K +SMEM_P4_V_PAGE_N_TILE_BYTES = TRTLLM_PAGE_SIZE * SMEM_P4_V_N_PER_CTA // 2 +SMEM_P4_VSF_PAGE_BYTES = ( + SMEM_P4_SCALE_CHUNKS_PER_PAGE * SMEM_P4_PV_SFB_N_CHUNKS * SMEM_P4_N_OUT_TILES * 512 +) +SMEM_P4_Q_ONLY_STAGE_BYTES = QK_A_STAGE_BYTES + QK_SFA_TMA_STAGE_BYTES +SMEM_P4_QK_KONLY_STAGE_BYTES = QK_B_STAGE_BYTES + QK_SFB_TMA_STAGE_BYTES +SMEM_P4_Q_ONLY_TAIL_STAGE_BYTES = QK_TAIL_A_STAGE_BYTES + QK_TAIL_SFA_STAGE_BYTES +SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES = QK_TAIL_B_STAGE_BYTES + QK_TAIL_SFB_STAGE_BYTES +SMEM_P4_QK_KONLY_TILE_BYTES = ( + SMEM_P4_QK_FULL_DATA_STAGES * SMEM_P4_QK_KONLY_STAGE_BYTES + SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES +) +SMEM_P4_Q_COMPACT_TILE_BYTES = ( + SMEM_P4_QK_FULL_DATA_STAGES * QK_A_STAGE_BYTES + QK_TAIL_A_STAGE_BYTES +) +SMEM_P4_QSF_COMPACT_TILE_BYTES = ( + SMEM_P4_QK_FULL_DATA_STAGES * QK_SFA_STAGE_BYTES + QK_TAIL_SFA_STAGE_BYTES +) +SMEM_P4_K_COMPACT_SLOT_BYTES = ( + SMEM_P4_QK_FULL_DATA_STAGES * QK_B_STAGE_BYTES + QK_TAIL_B_STAGE_BYTES +) +SMEM_P4_KSF_COMPACT_SLOT_BYTES = ( + SMEM_P4_QK_FULL_DATA_STAGES * QK_SFB_STAGE_BYTES + QK_TAIL_SFB_STAGE_BYTES +) +SMEM_P4_K_COMPACT_RING_BYTES = SMEM_P4_QK_KTILE_RING_STAGES * SMEM_P4_K_COMPACT_SLOT_BYTES +SMEM_P4_KSF_COMPACT_RING_BYTES = SMEM_P4_QK_KTILE_RING_STAGES * SMEM_P4_KSF_COMPACT_SLOT_BYTES +SMEM_P4_V_DATA_STAGE_BYTES = B_STAGE_BYTES * SMEM_P4_N_OUT_TILES +SMEM_P4_V_SFB_TMA_STAGE_BYTES = SFA_TMA_STAGE_BYTES * SMEM_P4_PV_SFB_N_CHUNKS * SMEM_P4_N_OUT_TILES +SMEM_P4_V_TILE_TMA_BYTES = SMEM_P4_V_DATA_STAGE_BYTES + SMEM_P4_V_SFB_TMA_STAGE_BYTES +SMEM_P4_CONTROL_MBAR_COUNT = ( + SMEM_P4_QK_COMPLETION_MBARS + + SMEM_P4_QK_PIPELINE_SLOTS + + 1 + + NUM_QK_Q_PIPELINE_STAGE * 2 + + SMEM_P4_QK_BARRIER_STAGES * 2 + + SMEM_P4_V_PIPELINE_STAGES * 2 + + SMEM_P4_TMEM_SCORE_PIPELINE_STAGES * 2 + + SMEM_P4_P4_CONTROL_MBAR_STAGES * 2 + + SMEM_P4_PAIR_OVERLAP_MBAR_COUNT + + SMEM_P4_QK_SCORE_RING_SLOTS + + 2 + + SMEM_P4_QK_BULK_READY_MBARS +) +SMEM_P4_CONTROL_BYTES_RAW = SMEM_P4_CONTROL_MBAR_COUNT * 8 + 8 +SMEM_P4_CONTROL_AND_ALIGNMENT_BYTES = (SMEM_P4_CONTROL_BYTES_RAW + 127) // 128 * 128 + + +def _refresh_sf_vec_config(sf_vec_size: int) -> None: + if sf_vec_size not in SF_VEC_SIZE_OPTIONS: + raise ValueError(f"sf_vec_size must be one of {SF_VEC_SIZE_OPTIONS}, got {sf_vec_size}") + if QK_LOGICAL_DIM % sf_vec_size != 0: + raise ValueError( + f"QK_LOGICAL_DIM={QK_LOGICAL_DIM} must be divisible by sf_vec_size={sf_vec_size}" + ) + if QK_LOGICAL_DIM % MMA_KBLOCK_DIM != 0: + raise ValueError( + f"QK_LOGICAL_DIM={QK_LOGICAL_DIM} must be divisible by MMA_KBLOCK_DIM={MMA_KBLOCK_DIM}" + ) + if QK_PADDED_DIM < QK_LOGICAL_DIM or QK_PADDED_DIM % MMA_KBLOCK_DIM != 0: + raise ValueError( + "QK_PADDED_DIM must cover QK_LOGICAL_DIM and be a whole Kblock; " + f"got padded={QK_PADDED_DIM}, logical={QK_LOGICAL_DIM}, " + f"kblock={MMA_KBLOCK_DIM}" + ) + if QK_DATA_STAGE_K_DIM != MMA_TILER_MNK[2]: + raise ValueError( + "QK_DATA_STAGE_K_DIM must remain K=256 for the FP4 s128b full-stage " + f"descriptor contract; got {QK_DATA_STAGE_K_DIM}" + ) + if QK_TAIL_DATA_STAGE_K_DIM != 2 * MMA_KBLOCK_DIM: + raise ValueError( + f"QK_TAIL_DATA_STAGE_K_DIM must remain K=128 for the logical K640 tail; got {QK_TAIL_DATA_STAGE_K_DIM}" + ) + if QK_K_TAIL_DATA_STAGE_K_DIM != 2 * MMA_KBLOCK_DIM: + raise ValueError( + "QK_K_TAIL_DATA_STAGE_K_DIM must remain physical K=128 for the " + f"TensorRT-LLM K640 tail; got {QK_K_TAIL_DATA_STAGE_K_DIM}" + ) + qk_data_stage_kblocks = QK_DATA_STAGE_K_DIM // MMA_KBLOCK_DIM + qk_mma_kblocks = QK_LOGICAL_DIM // MMA_KBLOCK_DIM + qk_data_stages = (qk_mma_kblocks + qk_data_stage_kblocks - 1) // qk_data_stage_kblocks + qk_ktile_ring_data_stages = SMEM_P4_QK_KTILE_RING_STAGES * qk_data_stages + qk_tail_stage_end = ( + SMEM_P4_QK_TAIL_STAGE_KBLOCK_START * MMA_KBLOCK_DIM + QK_TAIL_DATA_STAGE_K_DIM + ) + if qk_tail_stage_end != QK_LOGICAL_DIM: + raise ValueError( + "QK full-stage plus tail split must cover QK_LOGICAL_DIM exactly; " + f"got tail_stage_end={qk_tail_stage_end}, logical={QK_LOGICAL_DIM}" + ) + qk_k_tail_stage_end = ( + SMEM_P4_QK_TAIL_STAGE_KBLOCK_START * MMA_KBLOCK_DIM + QK_K_TAIL_DATA_STAGE_K_DIM + ) + if qk_k_tail_stage_end != TRTLLM_K_STORAGE_DIM: + raise ValueError( + "QK K full-stage plus physical tail split must cover K640 exactly; " + f"got tail_stage_end={qk_k_tail_stage_end}, physical={TRTLLM_K_STORAGE_DIM}" + ) + if QK_PADDED_DIM // 2 % 64 != 0: + raise ValueError( + f"QK padded FP4 row stride must be 64B aligned for the K128 tail; got {QK_PADDED_DIM // 2} bytes" + ) + if TRTLLM_K_STORAGE_DIM // 2 % 64 != 0: + raise ValueError( + "TensorRT-LLM K640 packed row stride must be 64B aligned for the " + f"s64b K128 tail; got {TRTLLM_K_STORAGE_DIM // 2} bytes" + ) + if KV_TILE % sf_vec_size != 0: + raise ValueError(f"KV_TILE={KV_TILE} must be divisible by sf_vec_size={sf_vec_size}") + if SMEM_P4_SCORE_SLOT_N % sf_vec_size != 0: + raise ValueError( + "SMEM_P4 score slot must divide the FP4 scale-vector size; " + f"got N={SMEM_P4_SCORE_SLOT_N}, sf_vec_size={sf_vec_size}" + ) + mma_scale_chunk_k = sf_vec_size * 4 + mma_scale_chunks_per_kblock = _scale_words_for_k(MMA_KBLOCK_DIM, sf_vec_size) + mma_scale_chunks_per_tile = _scale_words_for_k(MMA_TILER_MNK[2], sf_vec_size) + qk_resident_sf_chunks = _scale_words_for_k(QK_PADDED_DIM, sf_vec_size) + qk_resident_sf_cols = qk_resident_sf_chunks + qk_data_stage_sf_chunks = _scale_words_for_k(QK_DATA_STAGE_K_DIM, sf_vec_size) + qk_tail_data_stage_sf_chunks = _scale_words_for_k(QK_TAIL_DATA_STAGE_K_DIM, sf_vec_size) + qk_k_tail_data_stage_sf_chunks = _scale_words_for_k(QK_K_TAIL_DATA_STAGE_K_DIM, sf_vec_size) + sfa_tma_stage_bytes = 512 * mma_scale_chunks_per_tile + sfb_tma_stage_bytes = sfa_tma_stage_bytes * 2 + qk_sfa_stage_bytes = 512 * qk_data_stage_sf_chunks + qk_sfa_tma_stage_bytes = qk_sfa_stage_bytes + qk_sfb_n_chunks = _qk_sfb_smem_nchunks_per_cta(SMEM_P4_CTA_GROUP_M, SMEM_P4_QK_SCORE_SLOT_N) + qk_sfb_stage_bytes = qk_sfa_stage_bytes * qk_sfb_n_chunks + qk_sfb_tma_stage_bytes = qk_sfb_stage_bytes + pv_sfb_n_chunks = _qk_sfb_smem_nchunks_per_cta(SMEM_P4_CTA_GROUP_M, SMEM_P4_BMM2_N) + pv_compact_sfb_cols_per_scale_chunk = _qk_compact_sfb_cols_per_scale_chunk( + SMEM_P4_CTA_GROUP_M, SMEM_P4_BMM2_N, SMEM_P4_PV_SFB_WARPX2_MODE + ) + pv_sfb_cols_per_kblock = mma_scale_chunks_per_kblock * pv_compact_sfb_cols_per_scale_chunk + pv_sfb_score_slot_cols = SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT * pv_sfb_cols_per_kblock + p_scale_feed_stage_cols = _sfa_tmem_cols_for_scale_chunks( + SMEM_P4_BMM2_K // MMA_KBLOCK_DIM * mma_scale_chunks_per_kblock, True + ) + pv_sfb_score_col_offset = ( + SMEM_P4_P_SCALE_A_COL_OFFSET + SMEM_P4_SCORE_SLOTS_PER_KV_TILE * p_scale_feed_stage_cols + ) + qk_compact_sfb_cols_per_scale_chunk = _qk_compact_sfb_cols_per_scale_chunk( + SMEM_P4_CTA_GROUP_M, SMEM_P4_QK_SCORE_SLOT_N, SMEM_P4_QK_SFB_WARPX2_MODE + ) + qk_sfb_cols_per_kblock = mma_scale_chunks_per_kblock * qk_compact_sfb_cols_per_scale_chunk + qk_scale_burst_kblocks = _derive_qk_scale_burst_kblocks( + qk_mma_kblocks=qk_mma_kblocks, + qk_scale_burst_target_kblocks=qk_mma_kblocks, + scale_a_col_offset=SMEM_P4_SCALE_A_COL_OFFSET, + num_tmem_alloc_cols=NUM_TMEM_ALLOC_COLS, + qk_mma_m_dim=SMEM_P4_CTA_GROUP_M, + mma_scale_chunks_per_kblock=mma_scale_chunks_per_kblock, + use_sfa_128dp_unique=True, + qk_score_slot_n=SMEM_P4_QK_SCORE_SLOT_N, + qk_compact_sfb_cols_per_scale_chunk=qk_compact_sfb_cols_per_scale_chunk, + ) + qk_sfa_feed_cols = qk_resident_sf_cols + qk_sfb_feed_cols = qk_scale_burst_kblocks * qk_sfb_cols_per_kblock + qk_scale_b_col_offset = (SMEM_P4_SCALE_A_COL_OFFSET + qk_sfa_feed_cols + 15) // 16 * 16 + qk_scale_feed_col_end = max( + SMEM_P4_SCALE_A_COL_OFFSET + qk_sfa_feed_cols, + qk_scale_b_col_offset + qk_sfb_feed_cols, + ) + p_sfa_tail_col_offset = qk_scale_b_col_offset - p_scale_feed_stage_cols + pv_vsf_tail_col_offset = qk_scale_b_col_offset + pv_vsf_tail_col_end = ( + pv_vsf_tail_col_offset + + SMEM_P4_N_OUT_TILES * SMEM_P4_SCORE_SLOTS_PER_KV_TILE * pv_sfb_score_slot_cols + ) + running_rowmeta_tmem_cols = 2 + SMEM_P4_OLD_SCALE_TMEM_SLOTS + stage_rowmax_tmem_cols = 0 + stage_rowsum_tmem_cols = SMEM_P4_TMEM_BANKS * SMEM_P4_SCORE_SLOTS_PER_KV_TILE + rowmeta_tmem_cols = stage_rowmax_tmem_cols + stage_rowsum_tmem_cols + running_rowmeta_tmem_cols + rowmeta_tmem_col_offset = SMEM_P4_SCALE_FEED_TMEM_COLS[1] + q_compact_tile_bytes = SMEM_P4_QK_FULL_DATA_STAGES * QK_A_STAGE_BYTES + QK_TAIL_A_STAGE_BYTES + qsf_compact_tile_bytes = ( + SMEM_P4_QK_FULL_DATA_STAGES * qk_sfa_stage_bytes + 512 * qk_tail_data_stage_sf_chunks + ) + k_compact_slot_bytes = SMEM_P4_QK_FULL_DATA_STAGES * QK_B_STAGE_BYTES + QK_TAIL_B_STAGE_BYTES + ksf_compact_slot_bytes = ( + SMEM_P4_QK_FULL_DATA_STAGES * qk_sfb_stage_bytes + + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks + ) + static_smem_bytes = ( + SMEM_P4_CONTROL_AND_ALIGNMENT_BYTES + + q_compact_tile_bytes + + SMEM_P4_QK_KTILE_RING_STAGES * k_compact_slot_bytes + + qsf_compact_tile_bytes + + SMEM_P4_QK_KTILE_RING_STAGES * ksf_compact_slot_bytes + + SMEM_P4_P4_SMEM_WORDS * 4 + + SMEM_P4_V_PIPELINE_STAGES * B_STAGE_BYTES * SMEM_P4_N_OUT_TILES + + SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS * 4 + + SMEM_P4_V_PIPELINE_STAGES * sfa_tma_stage_bytes * pv_sfb_n_chunks * SMEM_P4_N_OUT_TILES + + SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS * 4 + + SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS * 4 + # Row maxima of the pipelined leading tiles, by tile parity. + + 2 * SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS * 4 + + SMEM_P4_PAGE_ID_PLAN_BYTES + + SMEM_P4_RUNTIME_SCALE_PAIR_BYTES + ) + globals().update( + SF_VEC_SIZE=sf_vec_size, + MMA_SCALE_CHUNK_K=mma_scale_chunk_k, + MMA_SCALE_CHUNKS_PER_KBLOCK=mma_scale_chunks_per_kblock, + MMA_SCALE_VEC_SIZE=_mma_scale_vec_size_enum(sf_vec_size), + MMA_SCALE_FORMAT=_mma_scale_format(sf_vec_size), + MMA_BLOCK_SCALE_KIND=_mma_block_scale_kind(sf_vec_size), + QK_MMA_KBLOCKS=qk_mma_kblocks, + SMEM_P4_QK_DATA_STAGE_KBLOCKS=qk_data_stage_kblocks, + SMEM_P4_QK_DATA_STAGES=qk_data_stages, + SMEM_P4_QK_DATA_STAGE_SF_CHUNKS=qk_data_stage_sf_chunks, + SMEM_P4_QK_TAIL_DATA_STAGE_SF_CHUNKS=qk_tail_data_stage_sf_chunks, + SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS=qk_k_tail_data_stage_sf_chunks, + SMEM_P4_QK_KTILE_RING_DATA_STAGES=qk_ktile_ring_data_stages, + SMEM_P4_QK_BARRIER_STAGES=SMEM_P4_QK_SMEM_PIPELINE_SLOTS, + NUM_QK_Q_STAGE=qk_data_stages, + NUM_QK_Q_PIPELINE_STAGE=qk_data_stages, + NUM_QK_AB_STAGE=qk_data_stages, + SMEM_P4_QK_SFB_N_CHUNKS=qk_sfb_n_chunks, + NUM_QK_K_STAGE=qk_ktile_ring_data_stages, + QK_SF_GROUPS=QK_PADDED_DIM // sf_vec_size, + PV_SF_GROUPS=KV_TILE // sf_vec_size, + PV_TOTAL_SF_GROUPS=KV_TOTAL // sf_vec_size, + SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK=mma_scale_chunks_per_kblock, + SMEM_P4_QK_MMA_KBLOCKS=qk_mma_kblocks, + SMEM_P4_QK_SCALE_BURST_KBLOCKS=qk_scale_burst_kblocks, + SMEM_P4_QK_SFA_FEED_COLS=qk_sfa_feed_cols, + SMEM_P4_QK_SFB_FEED_COLS=qk_sfb_feed_cols, + SMEM_P4_QK_COMPACT_SFB_COLS_PER_SCALE_CHUNK=qk_compact_sfb_cols_per_scale_chunk, + SMEM_P4_QK_SFB_COLS_PER_KBLOCK=qk_sfb_cols_per_kblock, + SMEM_P4_SCALE_B_COL_OFFSET=qk_scale_b_col_offset, + SMEM_P4_QK_SCALE_FEED_COL_END=qk_scale_feed_col_end, + SMEM_P4_P_SFA_TAIL_COL_OFFSET=p_sfa_tail_col_offset, + SMEM_P4_PV_VSF_TAIL_COL_OFFSET=pv_vsf_tail_col_offset, + SMEM_P4_PV_VSF_TAIL_COL_END=pv_vsf_tail_col_end, + SMEM_P4_QK_RESIDENT_SF_CHUNKS=qk_resident_sf_chunks, + SMEM_P4_QK_RESIDENT_SF_COLS=qk_resident_sf_cols, + SMEM_P4_RESIDENT_Q_SF_COL_OFFSET=SMEM_P4_SCALE_FEED_TMEM_COLS[0], + SMEM_P4_SCORE_SLOT_SF_GROUPS=SMEM_P4_SCORE_SLOT_N // sf_vec_size, + SMEM_P4_P_SFA_COLS_PER_SCALE_CHUNK=SMEM_P4_P_SFA_COLS_PER_SCALE_CHUNK, + SMEM_P4_P_SCALE_FEED_STAGE_COLS=p_scale_feed_stage_cols, + SMEM_P4_PV_SFB_SLOT_COL_OFFSET=SMEM_P4_P_SCALE_A_COL_OFFSET + p_scale_feed_stage_cols, + SMEM_P4_PV_SFB_N_CHUNKS=pv_sfb_n_chunks, + SMEM_P4_PV_COMPACT_SFB_COLS_PER_SCALE_CHUNK=pv_compact_sfb_cols_per_scale_chunk, + SMEM_P4_PV_SFB_COLS_PER_KBLOCK=pv_sfb_cols_per_kblock, + SMEM_P4_PV_SFB_SCORE_SLOT_COLS=pv_sfb_score_slot_cols, + SMEM_P4_PV_SFB_SCORE_COL_OFFSET=pv_sfb_score_col_offset, + SMEM_P4_P4_OWNER_TMEM_MAILBOX_SLOT_COL_OFFSET=SMEM_P4_PV_SFB_SLOT_COL_OFFSET + + pv_sfb_score_slot_cols, + SMEM_P4_RUNNING_ROWMETA_TMEM_COLS=running_rowmeta_tmem_cols, + SMEM_P4_ROWMETA_TMEM_COLS=rowmeta_tmem_cols, + SMEM_P4_ROWMETA_TMEM_COL_OFFSET=rowmeta_tmem_col_offset, + SMEM_P4_STAGE_ROWMAX_TMEM_COLS=stage_rowmax_tmem_cols, + SMEM_P4_STAGE_ROWSUM_TMEM_COLS=stage_rowsum_tmem_cols, + SMEM_P4_STAGE_ROWMAX_TMEM_COL_OFFSET=rowmeta_tmem_col_offset, + SMEM_P4_STAGE_ROWSUM_TMEM_COL_OFFSET=rowmeta_tmem_col_offset + stage_rowmax_tmem_cols, + SMEM_P4_RUNNING_ROWMAX_TMEM_COL_OFFSET=rowmeta_tmem_col_offset + + stage_rowmax_tmem_cols + + stage_rowsum_tmem_cols, + SMEM_P4_RUNNING_ROWSUM_TMEM_COL_OFFSET=rowmeta_tmem_col_offset + + stage_rowmax_tmem_cols + + stage_rowsum_tmem_cols + + 1, + SMEM_P4_OLD_SCALE_TMEM_COL_OFFSET=rowmeta_tmem_col_offset + + stage_rowmax_tmem_cols + + stage_rowsum_tmem_cols + + 2, + SMEM_P4_PAIR_FIRST_OLD_SCALE_TMEM_COL_OFFSET=rowmeta_tmem_col_offset + + stage_rowmax_tmem_cols + + stage_rowsum_tmem_cols + + 2 + + SMEM_P4_OLD_SCALE_TMEM_SLOTS + - 1, + SMEM_P4_ROWMETA_TMEM_COL_END=rowmeta_tmem_col_offset + + stage_rowmax_tmem_cols + + stage_rowsum_tmem_cols + + running_rowmeta_tmem_cols, + MMA_SCALE_CHUNKS_PER_TILE=mma_scale_chunks_per_tile, + SFA_TMA_STAGE_BYTES=sfa_tma_stage_bytes, + SFB_TMA_STAGE_BYTES=sfb_tma_stage_bytes, + QK_SFA_TMA_STAGE_BYTES=qk_sfa_tma_stage_bytes, + QK_SFB_TMA_STAGE_BYTES=qk_sfb_tma_stage_bytes, + QK_SFA_STAGE_BYTES=qk_sfa_stage_bytes, + QK_SFB_STAGE_BYTES=qk_sfb_stage_bytes, + QK_TAIL_SFA_STAGE_BYTES=512 * qk_tail_data_stage_sf_chunks, + QK_TAIL_SFB_STAGE_BYTES=512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks, + SMEM_P4_QK_STAGE_BYTES=QK_A_STAGE_BYTES + + QK_B_STAGE_BYTES + + qk_sfa_tma_stage_bytes + + qk_sfb_tma_stage_bytes, + SMEM_P4_Q_ONLY_STAGE_BYTES=QK_A_STAGE_BYTES + qk_sfa_tma_stage_bytes, + SMEM_P4_QK_KONLY_STAGE_BYTES=QK_B_STAGE_BYTES + qk_sfb_tma_stage_bytes, + SMEM_P4_QK_TAIL_STAGE_BYTES=QK_TAIL_A_STAGE_BYTES + + QK_TAIL_B_STAGE_BYTES + + 512 * qk_tail_data_stage_sf_chunks + + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks, + SMEM_P4_Q_ONLY_TAIL_STAGE_BYTES=QK_TAIL_A_STAGE_BYTES + 512 * qk_tail_data_stage_sf_chunks, + SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES=QK_TAIL_B_STAGE_BYTES + + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks, + SMEM_P4_QK_TILE_BYTES=SMEM_P4_QK_FULL_DATA_STAGES + * (QK_A_STAGE_BYTES + QK_B_STAGE_BYTES + qk_sfa_tma_stage_bytes + qk_sfb_tma_stage_bytes) + + QK_TAIL_A_STAGE_BYTES + + QK_TAIL_B_STAGE_BYTES + + 512 * qk_tail_data_stage_sf_chunks + + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks, + SMEM_P4_QK_KONLY_TILE_BYTES=SMEM_P4_QK_FULL_DATA_STAGES + * (QK_B_STAGE_BYTES + qk_sfb_tma_stage_bytes) + + QK_TAIL_B_STAGE_BYTES + + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks, + SMEM_P4_Q_COMPACT_TILE_BYTES=SMEM_P4_QK_FULL_DATA_STAGES * QK_A_STAGE_BYTES + + QK_TAIL_A_STAGE_BYTES, + SMEM_P4_QSF_COMPACT_TILE_BYTES=SMEM_P4_QK_FULL_DATA_STAGES * qk_sfa_stage_bytes + + 512 * qk_tail_data_stage_sf_chunks, + SMEM_P4_K_COMPACT_SLOT_BYTES=SMEM_P4_QK_FULL_DATA_STAGES * QK_B_STAGE_BYTES + + QK_TAIL_B_STAGE_BYTES, + SMEM_P4_KSF_COMPACT_SLOT_BYTES=SMEM_P4_QK_FULL_DATA_STAGES * qk_sfb_stage_bytes + + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks, + SMEM_P4_K_COMPACT_RING_BYTES=SMEM_P4_QK_KTILE_RING_STAGES + * (SMEM_P4_QK_FULL_DATA_STAGES * QK_B_STAGE_BYTES + QK_TAIL_B_STAGE_BYTES), + SMEM_P4_KSF_COMPACT_RING_BYTES=SMEM_P4_QK_KTILE_RING_STAGES + * ( + SMEM_P4_QK_FULL_DATA_STAGES * qk_sfb_stage_bytes + + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks + ), + SMEM_P4_STATIC_SMEM_BYTES=static_smem_bytes, + SMEM_P4_V_DATA_STAGE_BYTES=B_STAGE_BYTES * SMEM_P4_N_OUT_TILES, + SMEM_P4_V_SFB_TMA_STAGE_BYTES=sfa_tma_stage_bytes * pv_sfb_n_chunks * SMEM_P4_N_OUT_TILES, + SMEM_P4_V_TILE_TMA_BYTES=B_STAGE_BYTES * SMEM_P4_N_OUT_TILES + + sfa_tma_stage_bytes * pv_sfb_n_chunks * SMEM_P4_N_OUT_TILES, + ) + + +_refresh_sf_vec_config(SF_VEC_SIZE) + + +def _qk_compact_ring_stage_byte_offset( + stage: int, full_stage_bytes: int, compact_slot_bytes: int +) -> int: + slot = stage // SMEM_P4_QK_DATA_STAGES + local_stage = stage % SMEM_P4_QK_DATA_STAGES + return slot * compact_slot_bytes + local_stage * full_stage_bytes + + +TMEM_BAR_ID = 1 +# TMEM allocation is published to the producers and the PV issue warp. +TMEM_BAR_THREADS = 32 * (2 + SMEM_P4_TOTAL_PRODUCER_WARPS) +O_STORE_BAR_ID = 2 +SMEM_P4_P4_SLOT0_BAR_ID = 6 +SMEM_P4_P4_SLOT1_BAR_ID = 7 +SMEM_P4_QK1_SPLIT_ARM_BAR_ID = 8 +SMEM_P4_QK1_SPLIT_ARM_BAR_THREADS = 32 * 2 +SMEM_P4_QSF_OWNER_BAR_ID = 9 +SMEM_P4_P_SFA_OWNER_BAR_ID = 10 +SMEM_P4_DUAL_HALF_BAR_ID = SMEM_P4_P_SFA_OWNER_BAR_ID +SMEM_P4_DUAL_HALF_BAR_THREADS = 32 * SMEM_P4_TOTAL_PRODUCER_WARPS +O_STORE_BAR_THREADS = 32 * SMEM_P4_STORE_WARPS +SMEM_P4_QSF_OWNER_BAR_THREADS = 32 * (SMEM_P4_COMPUTE_WARPS + 1) + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def _select_runtime_kv_profile(max_kv_len: int) -> int: + if type(max_kv_len) is not int: + raise TypeError(f"max_kv_len must be an int, got {type(max_kv_len).__name__}") + if max_kv_len <= 0: + raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") + if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: + raise ValueError( + "max_kv_len exceeds the Int32-safe runtime-KV limit: " + f"{max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" + ) + if max_kv_len <= SMEM_P4_PAGE_PLAN_PROFILE_KV: + return SMEM_P4_PAGE_PLAN_PROFILE_KV + # Eager execution can report a different batch maximum on every step. + # Shifted power-of-two buckets bound the number of compiled variants and + # retain the established 1 Mi-token plus four-page-reserve profile. + profile_payload = max_kv_len - SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT + profile_kv = (1 << (profile_payload - 1).bit_length()) + (SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT) + return min(profile_kv, SMEM_P4_RUNTIME_MAX_KV) + + +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if output_dtype not in {torch.float16, torch.bfloat16}: + raise TypeError(f"output dtype must be torch.float16 or torch.bfloat16, got {output_dtype}") + + +def _validate_tensor_pointer_alignment( + name: str, + tensor: torch.Tensor, + alignment_bytes: int = PREPARED_BUFFER_ALIGNMENT_BYTES, +) -> None: + pointer = tensor.data_ptr() + if pointer % alignment_bytes != 0: + raise ValueError( + f"{name} data pointer must be {alignment_bytes}B aligned, got address remainder {pointer % alignment_bytes}" + ) + + +@cute.jit +def _pack_e2m1x8( + p0: ctm.Float32, + p1: ctm.Float32, + p2: ctm.Float32, + p3: ctm.Float32, + p4: ctm.Float32, + p5: ctm.Float32, + p6: ctm.Float32, + p7: ctm.Float32, +) -> ctm.Int32: + return cute_inline_ptx( + "{\n" + "\t.reg .b8 b0, b1, b2, b3;\n" + "\tcvt.rn.satfinite.e2m1x2.f32 b0, {$r1}, {$r0};\n" + "\tcvt.rn.satfinite.e2m1x2.f32 b1, {$r3}, {$r2};\n" + "\tcvt.rn.satfinite.e2m1x2.f32 b2, {$r5}, {$r4};\n" + "\tcvt.rn.satfinite.e2m1x2.f32 b3, {$r7}, {$r6};\n" + "\tmov.b32 {$w0}, {b0, b1, b2, b3};\n" + "\t}\n", + write_only_types=[ctm.Int32], + read_only_args=[p0, p1, p2, p3, p4, p5, p6, p7], + ) + + +@cute.jit +def fused_fp4_mla_decode_ctm( + q_ptr: cute.Pointer, + k_ptr: cute.Pointer, + q_sf_ptr: cute.Pointer, + k_sf_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + page_table_ptr: cute.Pointer, + valid_k_ptr: cute.Pointer, + helix_kv_bounds_ptr: cute.Pointer, + plan_ptr: cute.Pointer, + c_ptr: cute.Pointer, + accum_ptr: cute.Pointer, + row_max_ptr: cute.Pointer, + row_sum_ptr: cute.Pointer, + page_indptr_ptr: cute.Pointer, + q_global_scale_ptr: cute.Pointer, + kv_global_scale_ptr: cute.Pointer, + problem_size: tuple, + runtime_m: ctm.Int32, + runtime_l: ctm.Int32, + q_tma_l: ctm.Int32, + num_requests: ctm.Int32, + num_splits: ctm.Int32, + valid_m: ctm.Int32, + row_state_row_stride: ctm.Int32, + row_state_batch_stride: ctm.Int32, + page_table_len: ctm.Int32, + num_cache_pages: ctm.Int32, + num_v_cache_pages: ctm.Int32, + v_page_offset: ctm.Int32, + kv_page_stride_bytes: ctm.Int64, + ksf_page_stride_bytes: ctm.Int64, + vsf_page_stride_bytes: ctm.Int64, + softmax_scale_log2: ctm.Float32, + pv_output_scale: ctm.Float32, + stream: cuda.CUstream, + page_size: ctm.Constexpr = KV_TILE, + use_mixed_imlp: ctm.Constexpr = False, + query_len_per_seq: ctm.Constexpr = 1, + use_helix_kv_bounds: ctm.Constexpr = False, + use_smem_page_plan: ctm.Constexpr = True, + use_consecutive_page_pair: ctm.Constexpr = False, + use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, + store_split: ctm.Constexpr = False, + folded_query_heads: ctm.Constexpr = 0, + packed_query_heads: ctm.Constexpr = 0, + launch_dependents: ctm.Constexpr = False, + pdl_wait: ctm.Constexpr = False, +) -> None: + n, k = problem_size + m = runtime_m + batch_size = runtime_l + num_sequences = batch_size // query_len_per_seq + bounds_len = batch_size + if cutlass.const_expr(packed_query_heads > 0): + num_sequences = num_requests + bounds_len = num_requests * ctm.Int32(query_len_per_seq) + q_f4_tensor = cute.make_tensor( + q_ptr, + cute.make_layout( + (QK_PADDED_DIM, m, q_tma_l), + stride=(1, QK_PADDED_DIM, cute.assume(m * QK_PADDED_DIM, 32)), + ), + ) + k_f4_tensor = cute.make_tensor( + k_ptr, + cute.make_layout( + (TRTLLM_K_STORAGE_DIM, page_size, num_cache_pages), + stride=(1, TRTLLM_K_STORAGE_DIM, kv_page_stride_bytes * 2), + ), + ) + page_table_tensor = cute.make_tensor( + cute.recast_ptr(page_table_ptr, dtype=cutlass.Int32), + cute.make_layout((page_table_len,), stride=(1,)), + ) + page_indptr_tensor = cute.make_tensor( + cute.recast_ptr(page_indptr_ptr, dtype=cutlass.Int32), + cute.make_layout((num_sequences + 1,), stride=(1,)), + ) + q_global_scale_tensor = cute.make_tensor( + cute.recast_ptr(q_global_scale_ptr, dtype=cutlass.Float32), + cute.make_layout((1,), stride=(1,)), + ) + kv_global_scale_tensor = cute.make_tensor( + cute.recast_ptr(kv_global_scale_ptr, dtype=cutlass.Float32), + cute.make_layout((1,), stride=(1,)), + ) + valid_k_tensor = cute.make_tensor( + cute.recast_ptr(valid_k_ptr, dtype=cutlass.Int32), + cute.make_layout((num_sequences,), stride=(1,)), + ) + helix_kv_bounds_tensor = cute.make_tensor( + cute.recast_ptr(helix_kv_bounds_ptr, dtype=cutlass.Int32), + cute.make_layout((bounds_len,), stride=(1,)), + ) + # Packed task plan written by the plan kernel (see fp4_mla_packed_reduce): + # [tiles_per_task, used_tasks, begins[R + 1], request[C], split[C], count[C]]. + plan_len = ctm.Int32(3) + num_requests + ctm.Int32(3) * num_requests * num_splits + plan_tensor = cute.make_tensor( + cute.recast_ptr(plan_ptr, dtype=cutlass.Int32), + cute.make_layout((plan_len,), stride=(1,)), + ) + v_tma_tensor = cute.make_tensor( + b_ptr, + cute.make_layout( + ( + page_size, + SMEM_P4_V_N_PER_CTA, + n // SMEM_P4_V_N_PER_CTA, + num_v_cache_pages, + ), + stride=(1, page_size, page_size * SMEM_P4_V_N_PER_CTA, page_size * n), + ), + ) + v_pair_tma_tensor = cute.make_tensor( + b_ptr, + cute.make_layout( + ( + page_size, + SMEM_P4_V_N_PER_CTA, + SMEM_P4_N_OUT_TILES, + num_v_cache_pages, + CLUSTER_SHAPE_MNK[0], + ), + stride=( + 1, + page_size, + page_size * SMEM_P4_BMM2_N, + page_size * n, + page_size * SMEM_P4_V_N_PER_CTA, + ), + ), + ) + # The direct output holds valid_m rows per batch entry (packed rows drop + # the padded tail); split partials keep the padded m rows per task. + c_batch_stride = cute.assume(valid_m * n, 512) + c_layout = cute.make_layout( + (cute.assume(m, 32), cute.assume(n, 16), batch_size), + stride=(cute.assume(n, 16), 1, c_batch_stride), + ) + accum_layout = cute.make_layout( + (cute.assume(m, 32), cute.assume(n, 16), batch_size), + stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)), + ) + c_tensor = cute.make_tensor(c_ptr, c_layout) + accum_tensor = cute.make_tensor(accum_ptr, accum_layout) + row_state_layout = cute.make_layout( + (cute.assume(m, 32), batch_size), stride=(row_state_row_stride, row_state_batch_stride) + ) + row_max_tensor = cute.make_tensor(row_max_ptr, row_state_layout) + row_sum_tensor = cute.make_tensor(row_sum_ptr, row_state_layout) + q_sf_blocks = QK_SF_GROUPS // 4 + k_sf_blocks = TRTLLM_K_SF_GROUPS // 4 + vsf_mma_tensor = cute.make_tensor( + cute.recast_ptr(sfb_ptr, dtype=cutlass.Uint16), + cute.make_layout( + ( + 256, + SMEM_P4_SCALE_CHUNKS_PER_PAGE, + n // SMEM_P4_SFB_ATOM_ROWS, + num_v_cache_pages, + ), + stride=( + 1, + 256, + 256 * SMEM_P4_SCALE_CHUNKS_PER_PAGE, + vsf_page_stride_bytes // 2, + ), + ), + ) + qsf_mma_tile_m = SMEM_P4_CTA_GROUP_M + qsf_mma_m_tiles = cute.assume(m // qsf_mma_tile_m, 1) + qsf_mma_tensor = cute.make_tensor( + cute.recast_ptr(q_sf_ptr, dtype=cutlass.Uint16), + cute.make_layout( + (256, q_sf_blocks, qsf_mma_m_tiles, q_tma_l), + stride=( + 1, + 256, + 256 * q_sf_blocks, + cute.assume(256 * q_sf_blocks * qsf_mma_m_tiles, 1), + ), + ), + ) + ksf_mma_tensor = cute.make_tensor( + cute.recast_ptr(k_sf_ptr, dtype=cutlass.Uint16), + cute.make_layout( + (256, k_sf_blocks, num_cache_pages, 1), + stride=( + 1, + 256, + ksf_page_stride_bytes // 2, + ksf_page_stride_bytes // 2 * num_cache_pages, + ), + ), + ) + ksf_gather_tensor = cute.make_tensor( + cute.recast_ptr(k_sf_ptr, dtype=cutlass.Uint16), + cute.make_layout( + (256, k_sf_blocks * num_cache_pages), + stride=(1, 256), + ), + ) + tma_q_desc = cuda_tma.create_tensor_map_tiled_from_view( + q_f4_tensor, + box_dims=(QK_DATA_STAGE_K_DIM, MMA_TILER_MNK_PER_CTA[0], 1), + stride_order=(0, 1, 2), + swizzle=QK_FULL_TMA_SWIZZLE, + tma_format=cuda_tma.TensorMapDataFormat.B4X16, + ) + tma_k_page_desc = cuda_tma.create_tensor_map_tiled_from_view( + k_f4_tensor, + box_dims=(QK_DATA_STAGE_K_DIM, SMEM_P4_QK_TMA_N_PER_CHUNK, 1), + stride_order=(0, 1, 2), + swizzle=QK_FULL_TMA_SWIZZLE, + tma_format=cuda_tma.TensorMapDataFormat.B4X16, + ) + tma_q_tail_desc = cuda_tma.create_tensor_map_tiled_from_view( + q_f4_tensor, + box_dims=(QK_TAIL_DATA_STAGE_K_DIM, MMA_TILER_MNK_PER_CTA[0], 1), + stride_order=(0, 1, 2), + swizzle=QK_TAIL_TMA_SWIZZLE, + tma_format=cuda_tma.TensorMapDataFormat.B4X16, + ) + tma_k_tail_page_desc = cuda_tma.create_tensor_map_tiled_from_view( + k_f4_tensor, + box_dims=(QK_K_TAIL_DATA_STAGE_K_DIM, SMEM_P4_QK_TMA_N_PER_CHUNK, 1), + stride_order=(0, 1, 2), + swizzle=QK_K_TAIL_TMA_SWIZZLE, + tma_format=cuda_tma.TensorMapDataFormat.B4X16, + ) + tma_qsf_desc = cuda_tma.create_tensor_map_tiled_from_view( + qsf_mma_tensor, + box_dims=(256, SMEM_P4_QK_DATA_STAGE_SF_CHUNKS, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_qsf_tail_desc = cuda_tma.create_tensor_map_tiled_from_view( + qsf_mma_tensor, + box_dims=(256, SMEM_P4_QK_TAIL_DATA_STAGE_SF_CHUNKS, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_ksf_desc = cuda_tma.create_tensor_map_tiled_from_view( + ksf_mma_tensor, + box_dims=(256, 1, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + if cutlass.const_expr(use_consecutive_page_pair): + # A caller promise that covers every active tile lets one TMA span the + # adjacent physical-page pair. Mode order (atom, page, scale, batch) + # preserves the consumer's scale-major SMEM image. + tma_ksf_stage_desc = cuda_tma.create_tensor_map_tiled_from_view( + ksf_mma_tensor, + box_dims=( + 256, + SMEM_P4_QK_DATA_STAGE_SF_CHUNKS, + SMEM_P4_QK_SFB_N_CHUNKS, + 1, + ), + stride_order=(0, 2, 1, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_ksf_tail_stage_desc = cuda_tma.create_tensor_map_tiled_from_view( + ksf_mma_tensor, + box_dims=( + 256, + SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS, + SMEM_P4_QK_SFB_N_CHUNKS, + 1, + ), + stride_order=(0, 2, 1, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + elif cutlass.const_expr(use_ksf_gather4): + # Each 512B KSF atom is one gather row. The generic fast path is + # selected only when physical pages have no padding, so flattening + # (page, scale chunk) into a uniform row index is exact. + tma_ksf_stage_desc = cuda_tma.create_tensor_map_tiled_from_view( + ksf_gather_tensor, + box_dims=(256, 1), + stride_order=(0, 1), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_ksf_tail_stage_desc = cuda_tma.create_tensor_map_tiled_from_view( + ksf_gather_tensor, + box_dims=(256, 1), + stride_order=(0, 1), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + else: + tma_ksf_stage_desc = cuda_tma.create_tensor_map_tiled_from_view( + ksf_mma_tensor, + box_dims=(256, SMEM_P4_QK_DATA_STAGE_SF_CHUNKS, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_ksf_tail_stage_desc = cuda_tma.create_tensor_map_tiled_from_view( + ksf_mma_tensor, + box_dims=(256, SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_vsf_desc = cuda_tma.create_tensor_map_tiled_from_view( + vsf_mma_tensor, + box_dims=( + 256, + SMEM_P4_SCALE_CHUNKS_PER_PAGE, + SMEM_P4_PV_SFB_N_CHUNKS * SMEM_P4_N_OUT_TILES, + SMEM_P4_PAGES_PER_KV_TILE, + ), + stride_order=(0, 2, 1, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_vsf_stage_desc = cuda_tma.create_tensor_map_tiled_from_view( + vsf_mma_tensor, + box_dims=( + 256, + SMEM_P4_SCALE_CHUNKS_PER_PAGE, + SMEM_P4_PV_SFB_N_CHUNKS * SMEM_P4_N_OUT_TILES, + 1, + ), + stride_order=(0, 2, 1, 3), + swizzle=cuda_tma.TensorMapSwizzle.none, + tma_format=cuda_tma.TensorMapDataFormat.DEFAULT, + ) + tma_v_tile_desc = cuda_tma.create_tensor_map_tiled_from_view( + v_tma_tensor, + box_dims=(TRTLLM_PAGE_SIZE, SMEM_P4_V_N_PER_CTA, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda_tma.TensorMapSwizzle.s64b, + tma_format=cuda_tma.TensorMapDataFormat.B4X16, + ) + v_pair_page_box: ctm.Constexpr = 1 + if cutlass.const_expr(use_consecutive_page_pair): + v_pair_page_box = SMEM_P4_PAGES_PER_KV_TILE + tma_v_pair_desc = cuda_tma.create_tensor_map_tiled_from_view( + v_pair_tma_tensor, + box_dims=( + TRTLLM_PAGE_SIZE, + SMEM_P4_V_N_PER_CTA, + SMEM_P4_N_OUT_TILES, + v_pair_page_box, + 1, + ), + stride_order=(0, 1, 2, 3, 4), + swizzle=cuda_tma.TensorMapSwizzle.s64b, + tma_format=cuda_tma.TensorMapDataFormat.B4X16, + ) + tma_sfb_desc = tma_vsf_stage_desc + # features.ull2_mla_decode_pdl: launched as a programmatic dependent (the launch attribute is + # only passed then). + pdl_launch = {} + if cutlass.const_expr(pdl_wait): + pdl_launch["use_pdl"] = True + kernel( + page_table_tensor, + page_indptr_tensor, + valid_k_tensor, + helix_kv_bounds_tensor, + plan_tensor, + q_global_scale_tensor, + kv_global_scale_tensor, + tma_q_desc, + tma_k_page_desc, + tma_q_tail_desc, + tma_k_tail_page_desc, + tma_v_tile_desc, + tma_qsf_desc, + tma_qsf_tail_desc, + tma_ksf_desc, + tma_ksf_stage_desc, + tma_ksf_tail_stage_desc, + tma_vsf_desc, + tma_vsf_stage_desc, + tma_v_pair_desc, + tma_sfb_desc, + c_tensor, + accum_tensor, + row_max_tensor, + row_sum_tensor, + problem_size, + runtime_m, + num_requests, + num_splits, + valid_m, + c_batch_stride, + v_page_offset, + softmax_scale_log2, + pv_output_scale, + page_size, + query_len_per_seq, + use_helix_kv_bounds, + use_smem_page_plan, + use_mixed_imlp, + use_consecutive_page_pair, + use_ksf_gather4, + write_softmax_stats, + store_split, + folded_query_heads, + packed_query_heads, + launch_dependents, + pdl_wait, + ).launch( + grid=( + cute.ceil_div(c_tensor.shape[0], SMEM_P4_CTA_GROUP_M) * CLUSTER_SHAPE_MNK[0], + cute.ceil_div(OUT_DIM, MMA_TILER_MNK[1]), + c_tensor.shape[2], + ), + block=[THREADS_PER_CTA, 1, 1], + cluster=CLUSTER_SHAPE_MNK, + stream=stream, + min_blocks_per_mp=1, + **pdl_launch, + ) + + +@cute.jit +def _load_qk_qonly_kblock_stage( + tma_q_ptr, + tma_q_tail_ptr, + tma_qsf_ptr, + tma_qsf_tail_ptr, + q_tma_mbar, + sA, + sSFA, + tidx: ctm.Int32, + bidx: ctm.Int32, + bidz: ctm.Int32, + cta_rank: ctm.Int32, + qk_kblock_idx: ctm.Constexpr, + stage: ctm.Constexpr, + q_tma_phase: ctm.Constexpr, + manage_mbarrier: ctm.Constexpr = True, +) -> None: + del tidx + sA_stage = sA.subview(QK_A_STAGE_BYTES // 4 * stage) + sSFA_stage = sSFA.subview(QK_SFA_STAGE_BYTES * stage) + if cutlass.const_expr(qk_kblock_idx >= SMEM_P4_QK_TAIL_STAGE_KBLOCK_START): + tma_q_stage_ptr = tma_q_tail_ptr + tma_qsf_stage_ptr = tma_qsf_tail_ptr + q_tma_bytes: ctm.Constexpr = SMEM_P4_Q_ONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + else: + tma_q_stage_ptr = tma_q_ptr + tma_qsf_stage_ptr = tma_qsf_ptr + q_tma_bytes: ctm.Constexpr = SMEM_P4_Q_ONLY_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + mcast_mask = cutlass.Int16(1) << cutlass.Int16(cta_rank) + if cutlass.const_expr(manage_mbarrier): + if cta_rank == ctm.Int32(0): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(q_tma_mbar, q_tma_bytes) + if prims.elect_sync(): + m_tile_idx = bidx // ctm.Int32(CLUSTER_SHAPE_MNK[0]) + prims.cp_async_bulk_tensor_shared_cluster_global( + sA_stage, + tma_q_stage_ptr, + ( + ctm.Int32(qk_kblock_idx * MMA_KBLOCK_DIM), + m_tile_idx * ctm.Int32(SMEM_P4_CTA_GROUP_M) + cta_rank * ctm.Int32(SMEM_P4_BMM1_M), + bidz, + ), + _as_shared_cta(q_tma_mbar), + [], + multicast_mask=mcast_mask, + group=prims.CTAGroup.CTA_2, + ) + qsf_m_tile = m_tile_idx + prims.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_qsf_stage_ptr, + ( + ctm.Int32(0), + ctm.Int32(qk_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK), + qsf_m_tile, + bidz, + ), + _as_shared_cta(q_tma_mbar), + [], + multicast_mask=mcast_mask, + group=prims.CTAGroup.CTA_2, + ) + if cutlass.const_expr(manage_mbarrier): + if cta_rank == ctm.Int32(0): + while not prims.mbarrier_try_wait_parity(q_tma_mbar, q_tma_phase, time_limit=10000000): + pass + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + + +@cute.jit +def _lookup_physical_page( + mPageTable_pl: cute.Tensor, + logical_page: ctm.Int32, + batch: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, +) -> ctm.Int32: + has_pages = ctm.Int32(page_count > ctm.Int32(0)) + last_logical_page = ctm.max(page_count - ctm.Int32(1), ctm.Int32(0)) + safe_logical_page = ctm.min(logical_page, last_logical_page) + page_table_idx = (page_begin + safe_logical_page) * has_pages + return ctm.Int32(mPageTable_pl[page_table_idx]) + + +@cute.jit +def _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx: ctm.Int32) -> tuple: + physical_page0 = ctm.Int32(0) + physical_page1 = ctm.Int32(0) + tile_page_idx = ctm.Int32(kv_tile_idx * SMEM_P4_PAGES_PER_KV_TILE) + staged_page_pair = (sPageIdPlan.data_ptr() + tile_page_idx).load(count=2, alignment=8) + physical_page0 = staged_page_pair[0] + physical_page1 = staged_page_pair[1] + physical_page0 = cute.arch.make_warp_uniform(physical_page0) + physical_page1 = cute.arch.make_warp_uniform(physical_page1) + return (physical_page0, physical_page1) + + +@cute.jit +def _load_global_page_native_tile_pair( + mPageTable_pl: cute.Tensor, + kv_tile_idx: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, +) -> tuple: + tile_page_idx = ctm.Int32(kv_tile_idx * SMEM_P4_PAGES_PER_KV_TILE) + # Keep these as two scalar loads: a CSR row may start at an odd Int32, and + # the second logical page must clamp to the first for a one-page tail. + physical_page0 = _lookup_physical_page( + mPageTable_pl, tile_page_idx, ctm.Int32(0), page_begin, page_count + ) + physical_page1 = _lookup_physical_page( + mPageTable_pl, + tile_page_idx + ctm.Int32(1), + ctm.Int32(0), + page_begin, + page_count, + ) + physical_page0 = cute.arch.make_warp_uniform(physical_page0) + physical_page1 = cute.arch.make_warp_uniform(physical_page1) + return (physical_page0, physical_page1) + + +@cute.jit +def _tma_gather4_cluster( + smem_dst, + tma_desc, + col_coord: ctm.Int32, + row0: ctm.Int32, + row1: ctm.Int32, + row2: ctm.Int32, + row3: ctm.Int32, + barrier, + multicast_mask, +) -> None: + """Gather four KSF atoms and multicast the compact image to both CTAs.""" + smem_ptr = smem_dst.data_ptr() + tma_ptr = tma_desc.data_ptr() if hasattr(tma_desc, "data_ptr") else tma_desc + leader_barrier = _as_shared_cta(_mapa_shared_cluster(barrier, ctm.Int32(0))) + barrier_ptr = leader_barrier.data_ptr() + multicast_mask_u16 = ctm.Uint16(multicast_mask) + cute_inline_ptx( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], " + "[{$r7}], {$r8};", + read_only_args=[ + smem_ptr, + tma_ptr, + col_coord, + row0, + row1, + row2, + row3, + barrier_ptr, + multicast_mask_u16, + ], + ) + + +@cute.jit +def _load_qk_konly_kblock_stage( + mPageTable_pl: cute.Tensor, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_tma_mbar, + sB_slot, + sSFB_slot, + tidx: ctm.Int32, + bidz: ctm.Int32, + cta_rank: ctm.Int32, + physical_sfb_blocks_per_l: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, + tile_physical_page0: ctm.Int32, + tile_physical_page1: ctm.Int32, + kv_tile_idx: ctm.Int32, + qk_kblock_idx: ctm.Constexpr, + score_slot_idx: ctm.Constexpr, + stage: ctm.Constexpr, + page_size: ctm.Constexpr, + q_tma_phase: ctm.Constexpr, + issue_ksf: ctm.Constexpr = True, + manage_mbarrier: ctm.Constexpr = True, + use_consecutive_page_pair: ctm.Constexpr = False, + use_ksf_gather4: ctm.Constexpr = False, +) -> None: + del ( + tidx, + physical_sfb_blocks_per_l, + page_begin, + page_count, + score_slot_idx, + q_tma_phase, + manage_mbarrier, + ) + sB_stage = sB_slot.subview(ctm.Int32(stage * QK_B_STAGE_BYTES // 4)) + sSFB_stage = sSFB_slot.subview(ctm.Int32(stage * QK_SFB_STAGE_BYTES)) + if cutlass.const_expr(qk_kblock_idx >= SMEM_P4_QK_TAIL_STAGE_KBLOCK_START): + tma_k_stage_ptr = tma_k_tail_page_ptr + tma_ksf_merged_stage_ptr = tma_ksf_tail_stage_ptr + qk_k_stage_dim: ctm.Constexpr = QK_K_TAIL_DATA_STAGE_K_DIM + qk_k_stage_sf_chunks: ctm.Constexpr = SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS + else: + tma_k_stage_ptr = tma_k_page_ptr + tma_ksf_merged_stage_ptr = tma_ksf_stage_ptr + qk_k_stage_dim: ctm.Constexpr = QK_DATA_STAGE_K_DIM + qk_k_stage_sf_chunks: ctm.Constexpr = SMEM_P4_QK_DATA_STAGE_SF_CHUNKS + mcast_mask = cutlass.Int16(1) << cutlass.Int16(cta_rank) + physical_page = tile_physical_page0 + if cta_rank == ctm.Int32(1): + physical_page = tile_physical_page1 + for page_chunk in ctm.range_constexpr(SMEM_P4_QK_TMA_CHUNKS_PER_CTA): + sB_page_stage = sB_stage.subview( + ctm.Int32(page_chunk * SMEM_P4_QK_TMA_N_PER_CHUNK * qk_k_stage_dim // 8) + ) + k_tma_coord = ( + ctm.Int32(qk_kblock_idx * MMA_KBLOCK_DIM), + ctm.Int32(page_chunk * SMEM_P4_QK_TMA_N_PER_CHUNK), + physical_page, + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + sB_page_stage, + tma_k_stage_ptr, + k_tma_coord, + _as_shared_cta(qk_tma_mbar), + [], + multicast_mask=mcast_mask, + group=prims.CTAGroup.CTA_2, + ) + if cutlass.const_expr(use_consecutive_page_pair and issue_ksf): + # ``stride_order=(0, 2, 1, 3)`` makes the TMA coordinate order + # (atom, physical page, scale chunk, batch). The promise guarantees + # that page1 is exactly page0 + 1 for every active tile. + prims.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_ksf_merged_stage_ptr, + ( + ctm.Int32(0), + tile_physical_page0, + ctm.Int32(qk_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK), + ctm.Int32(0), + ), + _as_shared_cta(qk_tma_mbar), + [], + multicast_mask=cutlass.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(use_ksf_gather4 and issue_ksf): + # A gather4 transaction shares one column coordinate, so flatten each + # (physical page, scale chunk) atom into a row. Pairing two adjacent + # scale chunks produces the exact compact image expected by WARPX2: + # [scale0/page0, scale0/page1, scale1/page0, scale1/page1]. + scale_chunk_base = ctm.Int32(qk_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK) + for scale_pair_idx in ctm.range_constexpr(qk_k_stage_sf_chunks // 2): + scale_chunk0 = scale_chunk_base + ctm.Int32(scale_pair_idx * 2) + scale_chunk1 = scale_chunk0 + ctm.Int32(1) + row_page0_base = tile_physical_page0 * ctm.Int32(TRTLLM_K_SF_GROUPS // 4) + row_page1_base = tile_physical_page1 * ctm.Int32(TRTLLM_K_SF_GROUPS // 4) + _tma_gather4_cluster( + sSFB_stage.subview(ctm.Int32(scale_pair_idx * 4 * 512)), + tma_ksf_merged_stage_ptr, + ctm.Int32(0), + row_page0_base + scale_chunk0, + row_page1_base + scale_chunk0, + row_page0_base + scale_chunk1, + row_page1_base + scale_chunk1, + qk_tma_mbar, + cutlass.Int16(3), + ) + elif cutlass.const_expr(issue_ksf): + # Match the scale-major KSF image consumed by the WARPX2 S2T path: + # [scale chunk][logical page / N128 band][512-byte atom]. The prior + # page-major merged transaction left only one compact-ring residue + # numerically valid for nonzero K payloads. + for scale_chunk_idx in ctm.range_constexpr(qk_k_stage_sf_chunks): + for logical_page_idx in ctm.range_constexpr(SMEM_P4_PAGES_PER_KV_TILE): + physical_sf_page = tile_physical_page0 + if cutlass.const_expr(logical_page_idx == 1): + physical_sf_page = tile_physical_page1 + sSFB_atom = sSFB_stage.subview( + ctm.Int32( + (scale_chunk_idx * SMEM_P4_PAGES_PER_KV_TILE + logical_page_idx) * 512 + ) + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + sSFB_atom, + tma_ksf_ptr, + ( + ctm.Int32(0), + ctm.Int32( + qk_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK + scale_chunk_idx + ), + physical_sf_page, + ctm.Int32(0), + ), + _as_shared_cta(qk_tma_mbar), + [], + multicast_mask=cutlass.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + + +@cute.jit +def _issue_runtime_t336_qk_stage( + mPageTable_pl: cute.Tensor, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_stage_mbar, + sB_slot, + sSFB_slot, + tidx: ctm.Int32, + bidz: ctm.Int32, + physical_sfb_blocks_per_l: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, + physical_page0: ctm.Int32, + physical_page1: ctm.Int32, + kv_tile_idx: ctm.Int32, + page_size: ctm.Constexpr, + qk_stage_idx: ctm.Constexpr, + issue_rank0: ctm.Constexpr = True, + issue_rank1: ctm.Constexpr = True, + rank1_first: ctm.Constexpr = False, + use_consecutive_page_pair: ctm.Constexpr = False, + use_ksf_gather4: ctm.Constexpr = False, +) -> None: + qk_stage_kblock_idx: ctm.Constexpr = qk_stage_idx * SMEM_P4_QK_DATA_STAGE_KBLOCKS + if cutlass.const_expr(rank1_first and issue_rank1): + _load_qk_konly_kblock_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_stage_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + ctm.Int32(1), + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + qk_stage_kblock_idx, + 0, + qk_stage_idx, + page_size, + 0, + issue_ksf=False, + manage_mbarrier=False, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + if cutlass.const_expr(issue_rank0): + _load_qk_konly_kblock_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_stage_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + ctm.Int32(0), + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + qk_stage_kblock_idx, + 0, + qk_stage_idx, + page_size, + 0, + issue_ksf=True, + manage_mbarrier=False, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + if cutlass.const_expr(not rank1_first and issue_rank1): + _load_qk_konly_kblock_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_stage_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + ctm.Int32(1), + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + qk_stage_kblock_idx, + 0, + qk_stage_idx, + page_size, + 0, + issue_ksf=False, + manage_mbarrier=False, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + + +@cute.jit +def _load_runtime_t336_qk_tile( + mPageTable_pl: cute.Tensor, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_tma_mbar, + qk_bulk_ready_mbar_ptr, + pair_overlap_sync_mbars, + sB, + sSFB, + tidx: ctm.Int32, + bidz: ctm.Int32, + physical_sfb_blocks_per_l: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, + physical_page0: ctm.Int32, + physical_page1: ctm.Int32, + kv_tile_idx: ctm.Int32, + qk_slot_idx: ctm.Int32, + page_size: ctm.Constexpr, + initial_tile: ctm.Constexpr = False, + split_prefix_owner: ctm.Constexpr = False, + qk15_final: ctm.Constexpr = False, + use_consecutive_page_pair: ctm.Constexpr = False, + use_ksf_gather4: ctm.Constexpr = False, +) -> None: + # Steady tiles: CTA1 loads its own K half (_load_runtime_t336_qk_tile_rank1). + rank1_self_issued: ctm.Constexpr = ( + not initial_tile and not split_prefix_owner and not qk15_final + ) + qk_bulk_tma_mbar = qk_tma_mbar + if cutlass.const_expr(not initial_tile): + qk_bulk_tma_mbar = ctm.Array(qk_bulk_ready_mbar_ptr + qk_slot_idx, shape=1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + qk_bulk_tma_mbar, + (1 if qk15_final else SMEM_P4_QK_FULL_DATA_STAGES) + * SMEM_P4_QK_KONLY_STAGE_BYTES + * CLUSTER_SHAPE_MNK[0], + ) + if cutlass.const_expr(split_prefix_owner): + prims.barrier( + barrier_id=SMEM_P4_QK1_SPLIT_ARM_BAR_ID, + number_of_threads=SMEM_P4_QK1_SPLIT_ARM_BAR_THREADS, + ) + if prims.elect_sync(): + sB_slot = sB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_K_COMPACT_SLOT_BYTES // 4)) + sSFB_slot = sSFB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_KSF_COMPACT_SLOT_BYTES)) + if cutlass.const_expr(initial_tile): + _issue_runtime_t336_qk_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_tma_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + page_size, + SMEM_P4_QK_FULL_DATA_STAGES, + issue_rank0=False, + issue_rank1=True, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + _issue_runtime_t336_qk_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_tma_mbar if initial_tile else qk_bulk_tma_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + page_size, + 0, + issue_rank1=not split_prefix_owner and not rank1_self_issued, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + qk_stage1_mbar = qk_tma_mbar if initial_tile else qk_bulk_tma_mbar + if cutlass.const_expr(qk15_final): + qk_stage1_mbar = pair_overlap_sync_mbars.subview(SMEM_P4_QK15_STAGE1_MBAR_OFFSET) + prims.mbarrier_arrive_expect_tx( + qk_stage1_mbar, SMEM_P4_QK_KONLY_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + ) + _issue_runtime_t336_qk_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_stage1_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + page_size, + 1, + issue_rank1=not split_prefix_owner and not rank1_self_issued, + rank1_first=qk15_final, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + _issue_runtime_t336_qk_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_tma_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + page_size, + SMEM_P4_QK_FULL_DATA_STAGES, + issue_rank0=True, + issue_rank1=not initial_tile and not rank1_self_issued, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + + +@cute.jit +def _load_runtime_t336_qk_tile_rank1( + mPageTable_pl: cute.Tensor, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_full_mbar_ptr, + qk_bulk_ready_mbar_ptr, + sB, + sSFB, + physical_page0: ctm.Int32, + physical_page1: ctm.Int32, + kv_tile_idx: ctm.Int32, + qk_slot_idx: ctm.Int32, + page_size: ctm.Constexpr, +) -> None: + """CTA1's half of a steady K tile: its own K rows, completing on the leader's barriers.""" + leader_bulk_mbar = _as_shared_cta( + _mapa_shared_cluster(ctm.Array(qk_bulk_ready_mbar_ptr + qk_slot_idx, shape=1), ctm.Int32(0)) + ) + leader_full_mbar = _as_shared_cta( + _mapa_shared_cluster(ctm.Array(qk_full_mbar_ptr, shape=1), ctm.Int32(0)) + ) + if prims.elect_sync(): + sB_slot = sB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_K_COMPACT_SLOT_BYTES // 4)) + sSFB_slot = sSFB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_KSF_COMPACT_SLOT_BYTES)) + for qk_stage_idx in ctm.range_constexpr(SMEM_P4_QK_FULL_DATA_STAGES + 1): + _load_qk_konly_kblock_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + leader_full_mbar + if qk_stage_idx == SMEM_P4_QK_FULL_DATA_STAGES + else leader_bulk_mbar, + sB_slot, + sSFB_slot, + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(1), + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(0), + physical_page0, + physical_page1, + kv_tile_idx, + qk_stage_idx * SMEM_P4_QK_DATA_STAGE_KBLOCKS, + 0, + qk_stage_idx, + page_size, + 0, + issue_ksf=False, + manage_mbarrier=False, + ) + + +@cute.jit +def _issue_runtime_t336_qk_prefix_rank1( + mPageTable_pl: cute.Tensor, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_bulk_ready_mbar_ptr, + sB, + sSFB, + tidx: ctm.Int32, + bidz: ctm.Int32, + physical_sfb_blocks_per_l: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, + physical_page0: ctm.Int32, + physical_page1: ctm.Int32, + kv_tile_idx: ctm.Int32, + qk_slot_idx: ctm.Int32, + page_size: ctm.Constexpr, +) -> None: + prims.barrier( + barrier_id=SMEM_P4_QK1_SPLIT_ARM_BAR_ID, + number_of_threads=SMEM_P4_QK1_SPLIT_ARM_BAR_THREADS, + ) + qk_bulk_tma_mbar = ctm.Array(qk_bulk_ready_mbar_ptr + qk_slot_idx, shape=1) + if prims.elect_sync(): + sB_slot = sB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_K_COMPACT_SLOT_BYTES // 4)) + sSFB_slot = sSFB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_KSF_COMPACT_SLOT_BYTES)) + for qk_stage_idx in ctm.range_constexpr(SMEM_P4_QK_FULL_DATA_STAGES): + _issue_runtime_t336_qk_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_bulk_tma_mbar, + sB_slot, + sSFB_slot, + tidx, + bidz, + physical_sfb_blocks_per_l, + page_begin, + page_count, + physical_page0, + physical_page1, + kv_tile_idx, + page_size, + qk_stage_idx, + issue_rank0=False, + issue_rank1=True, + ) + + +@cute.jit +def _load_v_tile_stage( + mPageTable_pl: cute.Tensor, + sPageIdPlan, + tma_v_tile_ptr, + tma_v_pair_ptr, + tma_vsf_ptr, + tma_vsf_stage_ptr, + v_tma_mbar, + sB, + sSFB, + tidx: ctm.Int32, + bidy: ctm.Int32, + bidz: ctm.Int32, + cta_rank: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, + tile_physical_page0: ctm.Int32, + tile_physical_page1: ctm.Int32, + kv_tile_idx: ctm.Int32, + v_page_offset: ctm.Int32, + page_size: ctm.Constexpr, + v_tma_phase: ctm.Int32, + stage: ctm.Int32 = 0, + manage_mbarrier: ctm.Constexpr = True, + use_smem_page_plan: ctm.Constexpr = True, + use_consecutive_page_pair: ctm.Constexpr = False, +) -> None: + del tidx + v_smem_stage = stage + sB_stage = sB.subview(SMEM_P4_V_DATA_STAGE_BYTES // 4 * v_smem_stage) + sSFB_stage = sSFB.subview(SMEM_P4_V_SFB_TMA_STAGE_BYTES * v_smem_stage) + v_tma_bytes: ctm.Constexpr = SMEM_P4_V_TILE_TMA_BYTES * CLUSTER_SHAPE_MNK[0] + mcast_mask = cutlass.Int16(1) << cutlass.Int16(cta_rank) + if cutlass.const_expr(manage_mbarrier): + should_wait_v_tma = cta_rank == ctm.Int32(0) + if should_wait_v_tma: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(v_tma_mbar, v_tma_bytes) + del bidz, tile_physical_page0, tile_physical_page1 + if cutlass.const_expr(use_smem_page_plan): + resolved_physical_page0, resolved_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) + else: + resolved_physical_page0, resolved_physical_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, page_begin, page_count + ) + if cutlass.const_expr(use_consecutive_page_pair): + resolved_physical_page1 = resolved_physical_page0 + ctm.Int32(1) + if prims.elect_sync(): + physical_page0 = resolved_physical_page0 + physical_page1 = resolved_physical_page1 + if cutlass.const_expr(use_consecutive_page_pair): + prims.cp_async_bulk_tensor_shared_cluster_global( + sB_stage, + tma_v_pair_ptr, + ( + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(0), + physical_page0 + v_page_offset, + cta_rank, + ), + _as_shared_cta(v_tma_mbar), + [], + multicast_mask=mcast_mask, + group=prims.CTAGroup.CTA_2, + ) + else: + for logical_page_idx in ctm.range_constexpr(SMEM_P4_PAGES_PER_KV_TILE): + physical_page = physical_page0 + if cutlass.const_expr(logical_page_idx == 1): + physical_page = physical_page1 + sB_page = sB_stage.subview(ctm.Int32(logical_page_idx * B_STAGE_BYTES // 4)) + prims.cp_async_bulk_tensor_shared_cluster_global( + sB_page, + tma_v_pair_ptr, + ( + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(0), + physical_page + v_page_offset, + cta_rank, + ), + _as_shared_cta(v_tma_mbar), + [], + multicast_mask=mcast_mask, + group=prims.CTAGroup.CTA_2, + ) + if cta_rank == ctm.Int32(0): + if prims.elect_sync(): + if cutlass.const_expr(use_consecutive_page_pair): + prims.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_vsf_ptr, + ( + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(0), + resolved_physical_page0 + v_page_offset, + ), + _as_shared_cta(v_tma_mbar), + [], + multicast_mask=cutlass.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + else: + for logical_page_idx in ctm.range_constexpr(SMEM_P4_PAGES_PER_KV_TILE): + physical_page = resolved_physical_page0 + if cutlass.const_expr(logical_page_idx == 1): + physical_page = resolved_physical_page1 + sSFB_page = sSFB_stage.subview( + ctm.Int32(logical_page_idx * SMEM_P4_VSF_PAGE_BYTES) + ) + vsf_tma_coord = ( + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(0), + physical_page + v_page_offset, + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + sSFB_page, + tma_vsf_stage_ptr, + vsf_tma_coord, + _as_shared_cta(v_tma_mbar), + [], + multicast_mask=cutlass.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + if cutlass.const_expr(manage_mbarrier): + should_wait_v_tma = cta_rank == ctm.Int32(0) + if should_wait_v_tma: + while not prims.mbarrier_try_wait_parity(v_tma_mbar, v_tma_phase, time_limit=10000000): + pass + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + + +@cute.jit +def _fence_async_shared_cta() -> None: + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + + +@cute.jit +def _sleep_until_producer_slot_free(producer) -> ctm.Boolean: + """Poll a load pipeline slot without the SYNCS-woken spin of a blocking wait. + + A blocking mbarrier wait lowers to a NANOSLEEP.SYNCS loop that wakes on every + barrier update in the CTA; the load warps wait most of each KV tile, so that + loop burned issue slots the producer warps on the same SMSP need. + """ + acquired = producer.try_acquire() + while not acquired: + cute_inline_ptx( + "nanosleep.u32 {$r0};", + read_only_args=[ctm.Int32(SMEM_P4_LOAD_WAIT_SLEEP_NS)], + ) + acquired = producer.try_acquire() + return acquired + + +@cute.jit +def _permute_raw_trtllm_q_tail_in_smem(sQ, local_warp_idx: ctm.Int32, lane_idx: ctm.Int32) -> None: + tail_stage_word_offset: ctm.Constexpr = ( + SMEM_P4_QK_FULL_DATA_STAGES * QK_A_STAGE_BYTES // ctm.Int32.bytes + ) + tail_words_per_row: ctm.Constexpr = QK_TAIL_DATA_STAGE_K_DIM // 8 + words_per_group: ctm.Constexpr = SF_VEC_SIZE // 8 + local_row = local_warp_idx * ctm.Int32(32) + lane_idx + row_word_offset = local_row * ctm.Int32(tail_words_per_row) + tail_row_ptr = sQ.data_ptr() + ctm.Int32(tail_stage_word_offset) + row_word_offset + tail_swizzle = ctm.Swizzle.from_name("s64b") + raw_group1 = (tail_row_ptr + ctm.Int32(1 * words_per_group)).load_swizzled( + tail_swizzle, alignment=8, count=words_per_group + ) + raw_group2 = (tail_row_ptr + ctm.Int32(2 * words_per_group)).load_swizzled( + tail_swizzle, alignment=8, count=words_per_group + ) + raw_group4 = (tail_row_ptr + ctm.Int32(4 * words_per_group)).load_swizzled( + tail_swizzle, alignment=8, count=words_per_group + ) + (tail_row_ptr + ctm.Int32(4 * words_per_group)).store_swizzled( + raw_group1, swizzle=tail_swizzle, alignment=8 + ) + (tail_row_ptr + ctm.Int32(1 * words_per_group)).store_swizzled( + raw_group2, swizzle=tail_swizzle, alignment=8 + ) + (tail_row_ptr + ctm.Int32(2 * words_per_group)).store_swizzled( + raw_group4, swizzle=tail_swizzle, alignment=8 + ) + raw_group3 = (tail_row_ptr + ctm.Int32(3 * words_per_group)).load_swizzled( + tail_swizzle, alignment=8, count=words_per_group + ) + raw_group5 = (tail_row_ptr + ctm.Int32(5 * words_per_group)).load_swizzled( + tail_swizzle, alignment=8, count=words_per_group + ) + raw_group6 = (tail_row_ptr + ctm.Int32(6 * words_per_group)).load_swizzled( + tail_swizzle, alignment=8, count=words_per_group + ) + (tail_row_ptr + ctm.Int32(5 * words_per_group)).store_swizzled( + raw_group3, swizzle=tail_swizzle, alignment=8 + ) + (tail_row_ptr + ctm.Int32(6 * words_per_group)).store_swizzled( + raw_group5, swizzle=tail_swizzle, alignment=8 + ) + (tail_row_ptr + ctm.Int32(3 * words_per_group)).store_swizzled( + raw_group6, swizzle=tail_swizzle, alignment=8 + ) + + +@cute.jit +def _mask_invalid_score( + score: ctm.Float32, global_k_idx: ctm.Int32, valid_k: ctm.Int32 +) -> ctm.Float32: + masked = score + if global_k_idx >= valid_k: + masked = ctm.Float32(-ctm.Float32.inf) + return masked + + +@cute.jit +def _pack_e4m3x4_natural( + sf0: ctm.Float32, sf1: ctm.Float32, sf2: ctm.Float32, sf3: ctm.Float32 +) -> ctm.Int32: + return cute_inline_ptx( + "{\n" + "\t.reg .b16 lo, hi;\n" + "\tcvt.rn.satfinite.e4m3x2.f32 lo, {$r1}, {$r0};\n" + "\tcvt.rn.satfinite.e4m3x2.f32 hi, {$r3}, {$r2};\n" + "\tmov.b32 {$w0}, {lo, hi};\n" + "\t}\n", + write_only_types=[ctm.Int32], + read_only_args=[sf0, sf1, sf2, sf3], + ) + + +@cute.jit +def _pack_e4m3x4_natural_from_f16x2(pair01: ctm.Int32, pair23: ctm.Int32) -> ctm.Int32: + return cute_inline_ptx( + "{\n" + "\t.reg .b16 lo, hi;\n" + "\tcvt.rn.satfinite.e4m3x2.f16x2 lo, {$r0};\n" + "\tcvt.rn.satfinite.e4m3x2.f16x2 hi, {$r1};\n" + "\tmov.b32 {$w0}, {lo, hi};\n" + "\t}\n", + write_only_types=[ctm.Int32], + read_only_args=[pair01, pair23], + ) + + +@cute.jit +def _scale_pack_e4m3x4_natural_from_f16x2( + pair01: ctm.Int32, pair23: ctm.Int32, scale: ctm.Float32 +) -> ctm.Int32: + """Scale two FP16x2 pairs and pack the PV-only E4M3 scale word.""" + return cute_inline_ptx( + "{\n" + "\t.reg .f16 scale_h;\n" + "\t.reg .b32 scale_h2, scaled01, scaled23;\n" + "\t.reg .b16 lo, hi;\n" + "\tcvt.rn.f16.f32 scale_h, {$r2};\n" + "\tmov.b32 scale_h2, {scale_h, scale_h};\n" + "\tmul.f16x2 scaled01, {$r0}, scale_h2;\n" + "\tmul.f16x2 scaled23, {$r1}, scale_h2;\n" + "\tcvt.rn.satfinite.e4m3x2.f16x2 lo, scaled01;\n" + "\tcvt.rn.satfinite.e4m3x2.f16x2 hi, scaled23;\n" + "\tmov.b32 {$w0}, {lo, hi};\n" + "\t}\n", + write_only_types=[ctm.Int32], + read_only_args=[pair01, pair23, scale], + ) + + +@cute.jit +def _softmax_exp2_pair( + score0: ctm.Float32, + score1: ctm.Float32, + softmax_scale_log2: ctm.Float32, + p_bias: ctm.Float32, +) -> tuple: + shifted = nvvm_fma_packed_f32x2( + (score0, score1), (softmax_scale_log2, softmax_scale_log2), (p_bias, p_bias) + ) + return (cute.exp2(shifted[0], fastmath=True), cute.exp2(shifted[1], fastmath=True)) + + +@cute.jit +def _bitcast_i32_to_f32(value: ctm.Int32) -> ctm.Float32: + return ctm.Vector.from_elements((value,), dtype=ctm.Int32).bitcast(ctm.Float32)[0] + + +@cute.jit +def _load_smem_i32(ptr) -> ctm.Int32: + """Plain ld.shared through inline PTX so the load stays inside the loop.""" + return cute_inline_ptx( + "ld.shared.b32 {$w0}, [{$r0}];", + write_only_types=[ctm.Int32], + read_only_args=[ctm.Int32(ptr.toint())], + ) + + +@cute.jit +def _warp_uniform_f32(value: ctm.Float32) -> ctm.Float32: + bits = ctm.Vector.from_elements((value,), dtype=ctm.Float32).bitcast(ctm.Int32)[0] + uniform_bits = cute.arch.make_warp_uniform(bits) + return ctm.Vector.from_elements((uniform_bits,), dtype=ctm.Int32).bitcast(ctm.Float32)[0] + + +@cute.jit +def _pack_f32_pair_to_f16x2(value0: ctm.Float32, value1: ctm.Float32) -> ctm.Int32: + return ( + ctm.Vector.from_elements((value0, value1), dtype=ctm.Float32) + .to(ctm.Float16) + .bitcast(ctm.Int32)[0] + ) + + +@cute.jit +def _softmax_exp2_pair_packed_f16x2( + score0: ctm.Float32, + score1: ctm.Float32, + softmax_scale_log2: ctm.Float32, + p_bias: ctm.Float32, +) -> ctm.Int32: + shifted = nvvm_fma_packed_f32x2( + (score0, score1), (softmax_scale_log2, softmax_scale_log2), (p_bias, p_bias) + ) + shifted_h2 = _pack_f32_pair_to_f16x2(shifted[0], shifted[1]) + return cute_inline_ptx( + "ex2.approx.f16x2 {$w0}, {$r0};", + write_only_types=[ctm.Int32], + read_only_args=[shifted_h2], + ) + + +@cute.jit +def _add_f16x2(a: ctm.Int32, b: ctm.Int32) -> ctm.Int32: + return cute_inline_ptx( + "add.rn.f16x2 {$w0}, {$r0}, {$r1};", + write_only_types=[ctm.Int32], + read_only_args=[a, b], + ) + + +@cute.jit +def _unpack_f16x2_to_f32_pair(packed: ctm.Int32) -> tuple: + """Widen one packed probability pair for the denominator only.""" + values_f16 = ctm.Vector.from_elements((packed,), dtype=ctm.Int32).bitcast(ctm.Float16) + values_f32 = values_f16.to(ctm.Float32) + return (values_f32[0], values_f32[1]) + + +@cute.jit +def _p4_n256_prequant_rowsum_from_f16x2_psf( + sf_h2_01: ctm.Int32, + sf_h2_23: ctm.Int32, + group_p_sums, +) -> ctm.Float32: + """Weight FP32 P sums before the FP16 PSF values narrow to E4M3.""" + sf_exact_0, sf_exact_1 = _unpack_f16x2_to_f32_pair(sf_h2_01) + sf_exact_2, sf_exact_3 = _unpack_f16x2_to_f32_pair(sf_h2_23) + row_sum_parts01 = nvvm_mul_packed_f32x2( + (sf_exact_0, sf_exact_1), (group_p_sums[0], group_p_sums[1]) + ) + row_sum_parts23 = nvvm_mul_packed_f32x2( + (sf_exact_2, sf_exact_3), (group_p_sums[2], group_p_sums[3]) + ) + return _sum_p4_n256_psf_row_parts( + row_sum_parts01[0], + row_sum_parts01[1], + row_sum_parts23[0], + row_sum_parts23[1], + ) + + +@cute.jit +def _pack_e2m1x8_from_f16x2( + pair0: ctm.Int32, pair1: ctm.Int32, pair2: ctm.Int32, pair3: ctm.Int32 +) -> ctm.Int32: + return cute_inline_ptx( + "{\n" + "\t.reg .b8 b0, b1, b2, b3;\n" + "\tcvt.rn.satfinite.e2m1x2.f16x2 b0, {$r0};\n" + "\tcvt.rn.satfinite.e2m1x2.f16x2 b1, {$r1};\n" + "\tcvt.rn.satfinite.e2m1x2.f16x2 b2, {$r2};\n" + "\tcvt.rn.satfinite.e2m1x2.f16x2 b3, {$r3};\n" + "\tmov.b32 {$w0}, {b0, b1, b2, b3};\n" + "\t}\n", + write_only_types=[ctm.Int32], + read_only_args=[pair0, pair1, pair2, pair3], + ) + + +@cute.jit +def _horizontal_pack_two_f16x2_sums(group_sum0: ctm.Int32, group_sum1: ctm.Int32) -> ctm.Int32: + return cute_inline_ptx( + "{\n" + "\t.reg .f16 g0lo, g0hi, g1lo, g1hi, sum0, sum1;\n" + "\tmov.b32 {g0lo, g0hi}, {$r0};\n" + "\tmov.b32 {g1lo, g1hi}, {$r1};\n" + "\tadd.rn.f16 sum0, g0lo, g0hi;\n" + "\tadd.rn.f16 sum1, g1lo, g1hi;\n" + "\tmov.b32 {$w0}, {sum0, sum1};\n" + "\t}\n", + write_only_types=[ctm.Int32], + read_only_args=[group_sum0, group_sum1], + ) + + +@cute.jit +def _pack_p4_16_scores_from_group( + group_scores, + group_base: ctm.Int32, + valid_k: ctm.Int32, + group_valid: ctm.Boolean, + mask_valid_k: ctm.Constexpr, + use_mixed_imlp: ctm.Constexpr, + softmax_scale_log2: ctm.Float32, + p_bias: ctm.Float32, + elem_base: ctm.Constexpr, +) -> tuple: + words = ctm.Array(ctm.Int32, 2, space=ctm.AddressSpace.rmem) + p_sum_pair = (ctm.Float32(0.0), ctm.Float32(0.0)) + for word_part in ctm.range_constexpr(2): + packed_p0 = ctm.Float32(0.0) + packed_p1 = ctm.Float32(0.0) + packed_p2 = ctm.Float32(0.0) + packed_p3 = ctm.Float32(0.0) + packed_p4 = ctm.Float32(0.0) + packed_p5 = ctm.Float32(0.0) + packed_p6 = ctm.Float32(0.0) + packed_p7 = ctm.Float32(0.0) + for pair_idx in ctm.range_constexpr(4): + elem0: ctm.Constexpr = elem_base + word_part * 8 + pair_idx * 2 + elem1: ctm.Constexpr = elem0 + 1 + p0 = ctm.Float32(0.0) + p1 = ctm.Float32(0.0) + should_compute = ctm.Boolean(True) + if cutlass.const_expr(mask_valid_k): + should_compute = group_valid + if should_compute: + score0 = group_scores[elem0] + score1 = group_scores[elem1] + if cutlass.const_expr(mask_valid_k): + global_k_idx0 = group_base + ctm.Int32(elem0) + global_k_idx1 = group_base + ctm.Int32(elem1) + score0 = _mask_invalid_score(score0, global_k_idx0, valid_k) + score1 = _mask_invalid_score(score1, global_k_idx1, valid_k) + p0, p1 = _softmax_exp2_pair(score0, score1, softmax_scale_log2, p_bias) + if cutlass.const_expr(word_part == 0 and pair_idx == 0): + p_sum_pair = (p0, p1) + else: + p_sum_pair = nvvm_add_packed_f32x2(p_sum_pair, (p0, p1)) + if cutlass.const_expr(pair_idx == 0): + packed_p0 = p0 + packed_p1 = p1 + elif cutlass.const_expr(pair_idx == 1): + packed_p2 = p0 + packed_p3 = p1 + elif cutlass.const_expr(pair_idx == 2): + packed_p4 = p0 + packed_p5 = p1 + else: + packed_p6 = p0 + packed_p7 = p1 + words[word_part] = _pack_e2m1x8( + packed_p0, + packed_p1, + packed_p2, + packed_p3, + packed_p4, + packed_p5, + packed_p6, + packed_p7, + ) + p_sum = p_sum_pair[0] + p_sum_pair[1] + return (words[0], words[1], p_sum) + + +@cute.jit +def _pack_p4_16_scores_from_group_packed_f16x2( + group_scores, + group_base: ctm.Int32, + valid_k: ctm.Int32, + group_valid: ctm.Boolean, + mask_valid_k: ctm.Constexpr, + softmax_scale_log2: ctm.Float32, + p_bias: ctm.Float32, + elem_base: ctm.Constexpr, +) -> tuple: + words = ctm.Array(ctm.Int32, 2, space=ctm.AddressSpace.rmem) + # One FP16x2 partial sum per 8-element word; P <= 6, so a 16-element + # group sum stays far inside FP16 range. + word_sums = ctm.Array(ctm.Int32, 2, space=ctm.AddressSpace.rmem) + for word_part in ctm.range_constexpr(2): + packed_pair0 = ctm.Int32(0) + packed_pair1 = ctm.Int32(0) + packed_pair2 = ctm.Int32(0) + packed_pair3 = ctm.Int32(0) + for pair_idx in ctm.range_constexpr(4): + elem0: ctm.Constexpr = elem_base + word_part * 8 + pair_idx * 2 + elem1: ctm.Constexpr = elem0 + 1 + p_h2 = ctm.Int32(0) + should_compute = ctm.Boolean(True) + if cutlass.const_expr(mask_valid_k): + should_compute = group_valid + if should_compute: + score0 = group_scores[elem0] + score1 = group_scores[elem1] + if cutlass.const_expr(mask_valid_k): + global_k_idx0 = group_base + ctm.Int32(elem0) + global_k_idx1 = group_base + ctm.Int32(elem1) + score0 = _mask_invalid_score(score0, global_k_idx0, valid_k) + score1 = _mask_invalid_score(score1, global_k_idx1, valid_k) + p_h2 = _softmax_exp2_pair_packed_f16x2(score0, score1, softmax_scale_log2, p_bias) + if cutlass.const_expr(pair_idx == 0): + word_sums[word_part] = p_h2 + else: + word_sums[word_part] = _add_f16x2(word_sums[word_part], p_h2) + if cutlass.const_expr(pair_idx == 0): + packed_pair0 = p_h2 + elif cutlass.const_expr(pair_idx == 1): + packed_pair1 = p_h2 + elif cutlass.const_expr(pair_idx == 2): + packed_pair2 = p_h2 + else: + packed_pair3 = p_h2 + words[word_part] = _pack_e2m1x8_from_f16x2( + packed_pair0, packed_pair1, packed_pair2, packed_pair3 + ) + p_sum0, p_sum1 = _unpack_f16x2_to_f32_pair(_add_f16x2(word_sums[0], word_sums[1])) + p_sum = p_sum0 + p_sum1 + return (words[0], words[1], p_sum) + + +@cute.jit +def _issue_mxf4nvf4_mma_tile_from_base( + sA, + sB, + sSFA, + sSFB, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + stage: ctm.Int32, + clear_accumulator: ctm.Boolean, + acc_col_offset: ctm.Int32 = SMEM_P4_QK_ACC_COL_OFFSET, + sfa_col_offset: ctm.Int32 = SMEM_P4_SCALE_A_COL_OFFSET, + sfb_col_offset: ctm.Constexpr = SMEM_P4_SCALE_B_COL_OFFSET, + n_dim: ctm.Constexpr = SMEM_P4_BMM2_N, + m_dim: ctm.Constexpr = SMEM_P4_CTA_GROUP_M, + mma_kblock_start: ctm.Constexpr = 0, + mma_kblocks: ctm.Constexpr = MMA_KBLOCKS_PER_TILE, + stage_sfa_from_smem: ctm.Constexpr = True, + stage_sfb_from_smem: ctm.Constexpr = True, + a_from_tmem: ctm.Constexpr = False, + a_tmem_col_offset: ctm.Int32 = SMEM_P4_QK_ACC_COL_OFFSET, + mma_kblock_dim: ctm.Constexpr = MMA_KBLOCK_DIM, + mma_kblock_idesc_k_dim: ctm.Constexpr = MMA_KBLOCK_IDESC_K_DIM, + mma_scale_chunks_per_kblock: ctm.Constexpr = MMA_SCALE_CHUNKS_PER_KBLOCK, + mma_scale_vec_elements: ctm.Constexpr = SF_VEC_SIZE, + mma_scale_vec_size: ctm.Constexpr = MMA_SCALE_VEC_SIZE, + a_tmem_kblock_cols: ctm.Constexpr = SMEM_P4_P_TMEM_KBLOCK_COLS, + sfb_mma_col_delta: ctm.Constexpr = 0, + sfb_mma_kblock_stride_cols: ctm.Constexpr = -1, + compact_sfb_words_per_row: ctm.Constexpr = 0, + compact_sfb_extra_row_stride: ctm.Constexpr = 0, + compact_sfb_extra_col_stride: ctm.Constexpr = -1, + compact_sfb_cp_group_one: ctm.Constexpr = False, + compact_sfb_warpx2_mode: ctm.Constexpr = 0, + use_split_operand_stages: ctm.Constexpr = False, + a_stage: ctm.Int32 = 0, + b_stage: ctm.Int32 = 0, + use_split_scale_stages: ctm.Constexpr = False, + sfa_stage: ctm.Constexpr = 0, + sfb_stage: ctm.Constexpr = 0, + b_major: ctm.Constexpr = 0, + a_stage_words: ctm.Constexpr = A_STAGE_BYTES // 4, + b_stage_words: ctm.Constexpr = B_STAGE_BYTES // 4, + b_mma_smem_word_delta: ctm.Constexpr = 0, + a_mma_leading_byte_offset: ctm.Constexpr = 16, + a_mma_stride_byte_offset: ctm.Constexpr = 1024, + a_mma_layout: ctm.Constexpr = 2, + b_mma_leading_byte_offset: ctm.Constexpr = 16, + b_mma_stride_byte_offset: ctm.Constexpr = 1024, + b_mma_layout: ctm.Constexpr = 2, + b_mma_k_segment_offset: ctm.Constexpr = 0, + sfa_stage_bytes: ctm.Constexpr = SFA_STAGE_BYTES, + sfb_stage_bytes: ctm.Constexpr = SFB_STAGE_BYTES, + a_stage_word_offset: ctm.Constexpr = -1, + b_stage_word_offset: ctm.Constexpr = -1, + sfa_stage_byte_offset: ctm.Constexpr = -1, + sfb_stage_byte_offset: ctm.Constexpr = -1, + use_a_stage_word_offset: ctm.Constexpr = False, + use_b_stage_word_offset: ctm.Constexpr = False, + use_sfa_stage_byte_offset: ctm.Constexpr = False, + use_sfb_stage_byte_offset: ctm.Constexpr = False, + sfb_s2t_stride_byte_offset: ctm.Constexpr = 128, + use_split_mma_kblock_start: ctm.Constexpr = False, + a_mma_kblock_start: ctm.Constexpr = 0, + b_mma_kblock_start: ctm.Constexpr = 0, + reuse_b_first_kblock: ctm.Constexpr = False, + first_order_residual_tail: ctm.Constexpr = False, + tmem_lane_offset: ctm.Constexpr = 0, + a_tmem_lane_offset: ctm.Constexpr = -1, + scale_smem_kblock_start: ctm.Constexpr = -1, + scale_mma_kblock_start: ctm.Constexpr = -1, + a_sf_layout: ctm.Constexpr = SMEM_P4_SFA_LAYOUT, + a_collector_op: ctm.Constexpr = None, +) -> None: + if cutlass.const_expr( + first_order_residual_tail + and ( + not use_split_mma_kblock_start + or mma_kblocks != QK_TAIL_EFFECTIVE_MMA_KBLOCKS + or scale_mma_kblock_start < 0 + ) + ): + raise ValueError("first-order tail requires three stage-local MMAs and explicit scales") + issue_this_cta = ctm.Boolean(True) + idesc = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=ctm.Float4E2M1FN, + b_dtype=ctm.Float4E2M1FN, + scale_format=MMA_SCALE_FORMAT, + n_dim=n_dim, + m_dim=m_dim, + k_dim=mma_kblock_idesc_k_dim, + b_major=b_major, + a_sf_layout=a_sf_layout, + ) + sfa_s2t_shape, _ = prims.S2TCopyMode.S2T_128x128b + sfb_s2t_shape, sfb_s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + acc_col_id = base_col_id + ctm.Int32(acc_col_offset) + sfa_col_id = base_col_id + ctm.Int32(sfa_col_offset) + sfb_col_id = base_col_id + ctm.Int32(sfb_col_offset) + operand_tmem_mma_row_id = base_row_id + ctm.Int32(tmem_lane_offset) + a_effective_tmem_lane_offset: ctm.Constexpr = ( + tmem_lane_offset if a_tmem_lane_offset < 0 else a_tmem_lane_offset + ) + a_tmem_mma_row_id = base_row_id + ctm.Int32(a_effective_tmem_lane_offset) + acc_tmem_addr = operand_tmem_mma_row_id << ctm.Int32(16) | acc_col_id + sfa_tmem_addr_base = operand_tmem_mma_row_id << ctm.Int32(16) | sfa_col_id + sfb_tmem_addr_base = base_row_id << ctm.Int32(16) | sfb_col_id + acc_tmem_ptr = ctm.inttoptr(acc_tmem_addr, 6, ctm.Int32) + if cutlass.const_expr(use_b_stage_word_offset): + sB_stage = sB.subview(b_stage_word_offset) + elif cutlass.const_expr(use_split_operand_stages): + sB_stage = sB.subview(b_stage_words * b_stage) + else: + sB_stage = sB.subview(b_stage_words * stage) + if cutlass.const_expr(stage_sfb_from_smem): + if cutlass.const_expr(use_sfb_stage_byte_offset): + sSFB_stage = sSFB.subview(sfb_stage_byte_offset) + elif cutlass.const_expr(use_split_scale_stages): + sSFB_stage = sSFB.subview(sfb_stage_bytes * sfb_stage) + else: + sSFB_stage = sSFB.subview(sfb_stage_bytes * stage) + if cutlass.const_expr(stage_sfa_from_smem): + if cutlass.const_expr(use_sfa_stage_byte_offset): + sSFA_stage = sSFA.subview(sfa_stage_byte_offset) + elif cutlass.const_expr(use_split_scale_stages): + sSFA_stage = sSFA.subview(sfa_stage_bytes * sfa_stage) + else: + sSFA_stage = sSFA.subview(sfa_stage_bytes * stage) + desc_a_s2t_base = prims.Tcgen05SmemDesc.build( + sSFA_stage, leading_byte_offset=16, stride_byte_offset=128, layout=0 + ) + if cutlass.const_expr(stage_sfb_from_smem): + desc_b_s2t_base = prims.Tcgen05SmemDesc.build( + sSFB_stage, + leading_byte_offset=16, + stride_byte_offset=sfb_s2t_stride_byte_offset, + layout=0, + ) + if cutlass.const_expr(not a_from_tmem): + if cutlass.const_expr(use_a_stage_word_offset): + sA_stage = sA.subview(a_stage_word_offset) + elif cutlass.const_expr(use_split_operand_stages): + sA_stage = sA.subview(a_stage_words * a_stage) + else: + sA_stage = sA.subview(a_stage_words * stage) + desc_a_mma_base = prims.Tcgen05SmemDesc.build( + sA_stage, + leading_byte_offset=a_mma_leading_byte_offset, + stride_byte_offset=a_mma_stride_byte_offset, + layout=a_mma_layout, + ) + desc_b_mma_base = prims.Tcgen05SmemDesc.build( + sB_stage.subview(b_mma_smem_word_delta), + leading_byte_offset=b_mma_leading_byte_offset, + stride_byte_offset=b_mma_stride_byte_offset, + layout=b_mma_layout, + k_segment_offset=b_mma_k_segment_offset, + ) + scale_source_kblock_start: ctm.Constexpr = ( + mma_kblock_start if scale_smem_kblock_start < 0 else scale_smem_kblock_start + ) + sfa_s2t_count: ctm.Constexpr = mma_kblocks * mma_scale_chunks_per_kblock + if cutlass.const_expr(stage_sfa_from_smem): + source_scale_begin: ctm.Constexpr = scale_source_kblock_start * mma_scale_chunks_per_kblock + source_scale_end: ctm.Constexpr = source_scale_begin + sfa_s2t_count + source_scale_atom_begin: ctm.Constexpr = ( + source_scale_begin + // SMEM_P4_SFA_UNIQUE_CHUNKS_PER_UTCCP + * SMEM_P4_SFA_UNIQUE_CHUNKS_PER_UTCCP + ) + source_scale_atom_count: ctm.Constexpr = _ceil_div( + source_scale_end - source_scale_atom_begin, + SMEM_P4_SFA_UNIQUE_CHUNKS_PER_UTCCP, + ) + for s2t_atom_idx in ctm.range_constexpr(source_scale_atom_count): + source_scale_atom_idx: ctm.Constexpr = ( + source_scale_atom_begin + s2t_atom_idx * SMEM_P4_SFA_UNIQUE_CHUNKS_PER_UTCCP + ) + sfa_tmem_ptr = ctm.inttoptr( + sfa_tmem_addr_base + ctm.Int32(source_scale_atom_idx), 6, ctm.Int32 + ) + if prims.elect_sync(): + prims.tcgen05_cp( + sfa_s2t_shape, + sfa_tmem_ptr, + desc_a_s2t_base + ctm.Int64(32 * source_scale_atom_idx), + group=prims.CTAGroup.CTA_2, + ) + mma_nblocks: ctm.Constexpr = _ceil_div(n_dim, 128) + if cutlass.const_expr(stage_sfb_from_smem): + compact_sfb_copy_cols_per_scale_chunk: ctm.Constexpr = ( + _compact_sfb_copy_cols_per_scale_chunk(m_dim, n_dim) + ) + burst_words_per_row: ctm.Constexpr = mma_kblocks * mma_scale_chunks_per_kblock + words_per_row: ctm.Constexpr = ( + compact_sfb_words_per_row if compact_sfb_words_per_row != 0 else burst_words_per_row + ) + if cutlass.const_expr(m_dim == 128): + if cutlass.const_expr(compact_sfb_warpx2_mode != 0): + if cutlass.const_expr(compact_sfb_warpx2_mode == 2): + sfb_s2t_shape_x2, sfb_s2t_multicast_x2 = ( + prims.S2TCopyMode.S2T_64x128b_WARPX2_02_13 + ) + else: + sfb_s2t_shape_x2, sfb_s2t_multicast_x2 = ( + prims.S2TCopyMode.S2T_64x128b_WARPX2_01_23 + ) + compact_sfb_mma_cols_per_scale_chunk: ctm.Constexpr = ( + _compact_sfb_mma_cols_per_scale_chunk(m_dim, n_dim) + ) + for s2t_idx in ctm.range_constexpr(burst_words_per_row): + source_word_idx: ctm.Constexpr = ( + scale_source_kblock_start * mma_scale_chunks_per_kblock + s2t_idx + ) + sfb_tmem_ptr = ctm.inttoptr( + sfb_tmem_addr_base + + ctm.Int32(s2t_idx * compact_sfb_mma_cols_per_scale_chunk), + 6, + ctm.Int32, + ) + increment_s2t: ctm.Constexpr = ( + source_word_idx * SMEM_P4_QK_SFB_N_CHUNKS * (512 // 16) + + SMEM_P4_QK_COMPACT_SFB_S2T_BYTE_OFFSET + ) + if prims.elect_sync(): + prims.tcgen05_cp( + sfb_s2t_shape_x2, + sfb_tmem_ptr, + desc_b_s2t_base + ctm.Int64(increment_s2t), + group=prims.CTAGroup.CTA_2, + multicast=sfb_s2t_multicast_x2, + ) + else: + compact_sfb_cols_per_nblock: ctm.Constexpr = ( + compact_sfb_copy_cols_per_scale_chunk // mma_nblocks + ) + for s2t_idx in ctm.range_constexpr(burst_words_per_row): + source_word_idx: ctm.Constexpr = ( + scale_source_kblock_start * mma_scale_chunks_per_kblock + s2t_idx + ) + for nblock_idx in ctm.range_constexpr(mma_nblocks): + sfb_tmem_ptr = ctm.inttoptr( + sfb_tmem_addr_base + + ctm.Int32( + s2t_idx * compact_sfb_copy_cols_per_scale_chunk + + nblock_idx * compact_sfb_cols_per_nblock + ), + 6, + ctm.Int32, + ) + increment_s2t: ctm.Constexpr = ( + source_word_idx * (512 // 16) + + nblock_idx * words_per_row * (512 // 16) + + SMEM_P4_QK_COMPACT_SFB_S2T_BYTE_OFFSET + ) + if prims.elect_sync(): + if cutlass.const_expr(compact_sfb_cp_group_one): + prims.tcgen05_cp( + sfb_s2t_shape, + sfb_tmem_ptr, + desc_b_s2t_base + ctm.Int64(increment_s2t), + group=prims.CTAGroup.CTA_1, + ) + else: + prims.tcgen05_cp( + sfb_s2t_shape, + sfb_tmem_ptr, + desc_b_s2t_base + ctm.Int64(increment_s2t), + group=prims.CTAGroup.CTA_2, + multicast=sfb_s2t_multicast, + ) + if cutlass.const_expr( + compact_sfb_extra_row_stride != 0 and compact_sfb_extra_col_stride >= 0 + ): + if cutlass.const_expr(nblock_idx != 0): + extra_sfb_tmem_addr = base_row_id + ctm.Int32( + nblock_idx * compact_sfb_extra_row_stride + ) << ctm.Int32(16) | sfb_col_id + ctm.Int32( + s2t_idx * compact_sfb_copy_cols_per_scale_chunk + + nblock_idx * compact_sfb_extra_col_stride + ) + extra_sfb_tmem_ptr = ctm.inttoptr(extra_sfb_tmem_addr, 6, ctm.Int32) + if prims.elect_sync(): + if cutlass.const_expr(compact_sfb_cp_group_one): + prims.tcgen05_cp( + sfb_s2t_shape, + extra_sfb_tmem_ptr, + desc_b_s2t_base + ctm.Int64(increment_s2t), + group=prims.CTAGroup.CTA_1, + ) + else: + prims.tcgen05_cp( + sfb_s2t_shape, + extra_sfb_tmem_ptr, + desc_b_s2t_base + ctm.Int64(increment_s2t), + group=prims.CTAGroup.CTA_2, + multicast=sfb_s2t_multicast, + ) + else: + compact_sfb_nblocks: ctm.Constexpr = _ceil_div(n_dim, SMEM_P4_SFB_ATOM_ROWS) + compact_sfb_cols_per_nblock: ctm.Constexpr = ( + compact_sfb_copy_cols_per_scale_chunk // compact_sfb_nblocks + ) + for s2t_idx in ctm.range_constexpr(burst_words_per_row): + source_word_idx: ctm.Constexpr = ( + mma_kblock_start * mma_scale_chunks_per_kblock + s2t_idx + ) + for nblock_idx in ctm.range_constexpr(compact_sfb_nblocks): + sfb_tmem_ptr = ctm.inttoptr( + sfb_tmem_addr_base + + ctm.Int32( + s2t_idx * compact_sfb_copy_cols_per_scale_chunk + + nblock_idx * compact_sfb_cols_per_nblock + ), + 6, + ctm.Int32, + ) + increment_s2t: ctm.Constexpr = ( + 32 * source_word_idx + 32 * words_per_row * nblock_idx + ) + if prims.elect_sync(): + prims.tcgen05_cp( + sfb_s2t_shape, + sfb_tmem_ptr, + desc_b_s2t_base + ctm.Int64(increment_s2t), + group=prims.CTAGroup.CTA_2, + multicast=sfb_s2t_multicast, + ) + if cutlass.const_expr(stage_sfa_from_smem or stage_sfb_from_smem): + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + for kblock_idx in ctm.range_constexpr(mma_kblocks): + source_kblock_idx: ctm.Constexpr = mma_kblock_start + kblock_idx + if cutlass.const_expr(first_order_residual_tail): + a_tail_kblock_idx: ctm.Constexpr = FIRST_ORDER_TAIL_A_LOCAL_KBLOCKS[kblock_idx] + b_tail_kblock_idx: ctm.Constexpr = FIRST_ORDER_TAIL_B_LOCAL_KBLOCKS[kblock_idx] + scale_mma_source_kblock_idx: ctm.Constexpr = scale_mma_kblock_start + a_tail_kblock_idx + else: + a_tail_kblock_idx: ctm.Constexpr = kblock_idx + b_tail_kblock_idx: ctm.Constexpr = kblock_idx + scale_mma_source_kblock_idx: ctm.Constexpr = ( + source_kblock_idx + if scale_mma_kblock_start < 0 + else scale_mma_kblock_start + kblock_idx + ) + b_sf_ids_per_scale_atom: ctm.Constexpr = _mma_sf_ids_per_scale_atom( + mma_kblock_dim, mma_scale_vec_elements + ) + b_sf_id: ctm.Constexpr = scale_mma_source_kblock_idx % b_sf_ids_per_scale_atom + a_sf_id: ctm.Constexpr = 0 + sfa_scale_col: ctm.Constexpr = scale_mma_source_kblock_idx * mma_scale_chunks_per_kblock + idesc_updated = idesc.set_sf_ids(a_sf_id=a_sf_id, b_sf_id=b_sf_id) + enable_input_d = ctm.Boolean(True) + if cutlass.const_expr(kblock_idx == 0): + enable_input_d = ~ctm.Boolean(clear_accumulator) + sfa_mma_addr = sfa_tmem_addr_base + ctm.Int32(sfa_scale_col) + sfa_tmem_ptr = ctm.inttoptr(sfa_mma_addr, 6, ctm.Int32) + sfb_kblock_stride_cols: ctm.Constexpr = ( + mma_scale_chunks_per_kblock * _compact_sfb_mma_cols_per_scale_chunk(m_dim, n_dim) + if sfb_mma_kblock_stride_cols < 0 + else sfb_mma_kblock_stride_cols + ) + sfb_tmem_ptr = ctm.inttoptr( + sfb_tmem_addr_base + + ctm.Int32(sfb_mma_col_delta) + + ctm.Int32(b_tail_kblock_idx * sfb_kblock_stride_cols), + 6, + ctm.Int32, + ) + if cutlass.const_expr(use_split_mma_kblock_start): + increment_mma_a: ctm.Constexpr = ( + mma_kblock_dim // 32 * (a_mma_kblock_start + a_tail_kblock_idx) + ) + b_local_kblock_idx: ctm.Constexpr = ( + b_tail_kblock_idx + if first_order_residual_tail + else 0 + if reuse_b_first_kblock + else kblock_idx + ) + increment_mma_b: ctm.Constexpr = ( + mma_kblock_dim // 32 * (b_mma_kblock_start + b_local_kblock_idx) + ) + else: + increment_mma_a: ctm.Constexpr = mma_kblock_dim // 32 * source_kblock_idx + increment_mma_b: ctm.Constexpr = mma_kblock_dim // 32 * source_kblock_idx + if cutlass.const_expr(a_from_tmem): + if cutlass.const_expr(use_split_mma_kblock_start): + a_tmem_kblock_idx: ctm.Constexpr = a_mma_kblock_start + a_tail_kblock_idx + else: + a_tmem_kblock_idx: ctm.Constexpr = source_kblock_idx + a_col_id = base_col_id + ctm.Int32( + a_tmem_col_offset + a_tmem_kblock_idx * a_tmem_kblock_cols + ) + a_tmem_addr = a_tmem_mma_row_id << ctm.Int32(16) | a_col_id + a_mma_operand = ctm.inttoptr(a_tmem_addr, 6, ctm.Int32) + else: + a_mma_operand = desc_a_mma_base + ctm.Int64(increment_mma_a) + if issue_this_cta: + if prims.elect_sync(): + prims.tcgen05_mma_block_scale( + MMA_BLOCK_SCALE_KIND, + prims.CTAGroup.CTA_2, + acc_tmem_ptr, + a_mma_operand, + desc_b_mma_base + ctm.Int64(increment_mma_b), + idesc_updated, + enable_input_d=enable_input_d, + scale_a=sfa_tmem_ptr, + scale_b=sfb_tmem_ptr, + scale_vec_size=mma_scale_vec_size, + collector_op=a_collector_op, + ) + + +@cute.jit +def _issue_smem_p4_qk_score_slot_from_base( + sQ, + sK, + sQSF, + sKSF, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + stage: ctm.Constexpr, + score_slot_idx: ctm.Int32, + clear_accumulator: ctm.Constexpr, + qk_kblock_idx: ctm.Constexpr, + qk_kblocks: ctm.Constexpr = 1, + acc_col_offset: ctm.Int32 = SMEM_P4_QK_ACC_COL_OFFSET, + sfa_col_offset: ctm.Constexpr = SMEM_P4_SCALE_A_COL_OFFSET, + sfb_col_offset: ctm.Constexpr = SMEM_P4_SCALE_B_COL_OFFSET, + stage_sfa_from_smem: ctm.Constexpr = True, + stage_sfb_from_smem: ctm.Constexpr = True, + use_split_operand_stages: ctm.Constexpr = False, + a_stage: ctm.Constexpr = 0, + b_stage: ctm.Int32 = 0, + use_stage_local_operand_kblock_start: ctm.Constexpr = False, + stage_local_operand_kblock_start: ctm.Constexpr = -1, + stage_local_scale_kblock_start: ctm.Constexpr = -1, + reuse_b_first_kblock: ctm.Constexpr = False, + first_order_residual_tail: ctm.Constexpr = False, + b_smem_is_compact_slot_base: ctm.Constexpr = False, +) -> None: + normalized_score_slot_idx = ctm.Int32(score_slot_idx) % ctm.Int32(SMEM_P4_QK_SCORE_SLOTS) + qk_stage_acc_col_offset = acc_col_offset + normalized_score_slot_idx * ctm.Int32( + SMEM_P4_QK_PIPELINE_SLOT_STRIDE + ) + if cutlass.const_expr(stage_local_operand_kblock_start >= 0): + qk_stage_local_kblock_idx: ctm.Constexpr = stage_local_operand_kblock_start + else: + qk_stage_local_kblock_idx: ctm.Constexpr = qk_kblock_idx % SMEM_P4_QK_DATA_STAGE_KBLOCKS + if cutlass.const_expr(qk_kblock_idx >= SMEM_P4_QK_TAIL_STAGE_KBLOCK_START): + qk_a_mma_leading_byte_offset: ctm.Constexpr = QK_TAIL_MMA_LEADING_BYTE_OFFSET + qk_a_mma_stride_byte_offset: ctm.Constexpr = QK_TAIL_MMA_STRIDE_BYTE_OFFSET + qk_a_mma_layout: ctm.Constexpr = QK_TAIL_MMA_LAYOUT + qk_b_mma_leading_byte_offset: ctm.Constexpr = QK_K_TAIL_MMA_LEADING_BYTE_OFFSET + qk_b_mma_stride_byte_offset: ctm.Constexpr = QK_K_TAIL_MMA_STRIDE_BYTE_OFFSET + qk_b_mma_layout: ctm.Constexpr = QK_K_TAIL_MMA_LAYOUT + else: + qk_a_mma_leading_byte_offset: ctm.Constexpr = QK_FULL_MMA_LEADING_BYTE_OFFSET + qk_a_mma_stride_byte_offset: ctm.Constexpr = QK_FULL_MMA_STRIDE_BYTE_OFFSET + qk_a_mma_layout: ctm.Constexpr = QK_FULL_MMA_LAYOUT + qk_b_mma_leading_byte_offset: ctm.Constexpr = SMEM_P4_QK_B_MMA_LEADING_BYTES + qk_b_mma_stride_byte_offset: ctm.Constexpr = SMEM_P4_QK_B_MMA_STRIDE_BYTES + qk_b_mma_layout: ctm.Constexpr = SMEM_P4_QK_B_MMA_LAYOUT + physical_a_stage: ctm.Constexpr = a_stage if use_split_operand_stages else stage + q_stage_word_offset: ctm.Constexpr = physical_a_stage * QK_A_STAGE_BYTES // 4 + qsf_stage_byte_offset: ctm.Constexpr = physical_a_stage * QK_SFA_STAGE_BYTES + if cutlass.const_expr(b_smem_is_compact_slot_base): + k_stage_word_offset: ctm.Constexpr = stage * QK_B_STAGE_BYTES // 4 + ksf_stage_byte_offset: ctm.Constexpr = stage * QK_SFB_STAGE_BYTES + else: + physical_b_stage: ctm.Constexpr = b_stage if use_split_operand_stages else stage + k_stage_word_offset: ctm.Constexpr = ( + _qk_compact_ring_stage_byte_offset( + physical_b_stage, QK_B_STAGE_BYTES, SMEM_P4_K_COMPACT_SLOT_BYTES + ) + // 4 + ) + ksf_stage_byte_offset: ctm.Constexpr = _qk_compact_ring_stage_byte_offset( + physical_b_stage, QK_SFB_STAGE_BYTES, SMEM_P4_KSF_COMPACT_SLOT_BYTES + ) + _issue_mxf4nvf4_mma_tile_from_base( + sQ, + sK, + sQSF, + sKSF, + base_col_id, + base_row_id, + stage, + clear_accumulator=clear_accumulator, + acc_col_offset=qk_stage_acc_col_offset, + sfa_col_offset=sfa_col_offset, + sfb_col_offset=sfb_col_offset, + n_dim=SMEM_P4_QK_SCORE_SLOT_N, + m_dim=SMEM_P4_CTA_GROUP_M, + mma_kblock_start=qk_kblock_idx, + mma_kblocks=qk_kblocks, + mma_kblock_dim=SMEM_P4_QK_MMA_KBLOCK_DIM, + mma_kblock_idesc_k_dim=SMEM_P4_QK_MMA_KBLOCK_IDESC_K_DIM, + mma_scale_chunks_per_kblock=SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK, + mma_scale_vec_size=MMA_SCALE_VEC_SIZE, + stage_sfa_from_smem=stage_sfa_from_smem, + stage_sfb_from_smem=stage_sfb_from_smem, + use_split_operand_stages=use_split_operand_stages, + a_stage=a_stage, + b_stage=b_stage, + a_stage_words=QK_A_STAGE_BYTES // 4, + b_stage_words=QK_B_STAGE_BYTES // 4, + sfa_stage_bytes=QK_SFA_STAGE_BYTES, + sfb_stage_bytes=QK_SFB_STAGE_BYTES, + a_stage_word_offset=q_stage_word_offset, + b_stage_word_offset=k_stage_word_offset, + sfa_stage_byte_offset=qsf_stage_byte_offset, + sfb_stage_byte_offset=ksf_stage_byte_offset, + use_a_stage_word_offset=True, + use_b_stage_word_offset=True, + use_sfa_stage_byte_offset=True, + use_sfb_stage_byte_offset=True, + compact_sfb_words_per_row=SMEM_P4_QK_MMA_KBLOCKS * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK, + compact_sfb_extra_row_stride=SMEM_P4_QK_SFB_EXTRA_ROW_STRIDE, + compact_sfb_extra_col_stride=SMEM_P4_QK_SFB_EXTRA_COL_STRIDE, + compact_sfb_cp_group_one=SMEM_P4_QK_SFB_CP_GROUP_ONE, + compact_sfb_warpx2_mode=SMEM_P4_QK_SFB_WARPX2_MODE, + sfb_mma_col_delta=SMEM_P4_QK_SFB_MMA_COL_DELTA, + use_split_mma_kblock_start=use_split_operand_stages or use_stage_local_operand_kblock_start, + a_mma_kblock_start=qk_stage_local_kblock_idx, + b_mma_kblock_start=qk_stage_local_kblock_idx, + reuse_b_first_kblock=reuse_b_first_kblock, + first_order_residual_tail=first_order_residual_tail, + scale_smem_kblock_start=stage_local_scale_kblock_start, + scale_mma_kblock_start=stage_local_scale_kblock_start, + a_mma_leading_byte_offset=qk_a_mma_leading_byte_offset, + a_mma_stride_byte_offset=qk_a_mma_stride_byte_offset, + a_mma_layout=qk_a_mma_layout, + b_mma_leading_byte_offset=qk_b_mma_leading_byte_offset, + b_mma_stride_byte_offset=qk_b_mma_stride_byte_offset, + b_mma_layout=qk_b_mma_layout, + b_mma_smem_word_delta=SMEM_P4_QK_B_MMA_SMEM_WORD_DELTA, + b_mma_k_segment_offset=SMEM_P4_QK_B_MMA_K_SEGMENT_OFFSET, + sfb_s2t_stride_byte_offset=SMEM_P4_QK_SFB_S2T_STRIDE_BYTES, + ) + + +@cute.jit +def _store_raw_trtllm_qsf_128dp_unique_owner_to_tmem_from_smem( + sQSF, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + cta_rank: ctm.Int32, + local_warp_idx: ctm.Int32, + lane_idx: ctm.Int32, + qk_kblocks: ctm.Constexpr = QK_MMA_KBLOCKS, +) -> None: + qsf_col_id = base_col_id + ctm.Int32(SMEM_P4_RESIDENT_Q_SF_COL_OFFSET) + qsf_subp = cta_rank * ctm.Int32(SMEM_P4_BMM1_M // 32) + (local_warp_idx & ctm.Int32(1)) + qsf_row_subp = qsf_subp + if cutlass.const_expr(qk_kblocks != QK_MMA_KBLOCKS): + raise ValueError("raw TRT QSF publication requires the complete resident QK tile") + for qk_kblock_idx in ctm.range_constexpr(QK_MMA_KBLOCKS): + if cutlass.const_expr(qk_kblock_idx < QK_MMA_KBLOCKS - 2): + qsf_stage: ctm.Constexpr = qk_kblock_idx // SMEM_P4_QK_DATA_STAGE_KBLOCKS + qsf_stage_chunk: ctm.Constexpr = qk_kblock_idx % SMEM_P4_QK_DATA_STAGE_KBLOCKS + sQSF_stage = sQSF.subview(QK_SFA_STAGE_BYTES * qsf_stage) + qsf_stage_ptr = ctm.Int64(sQSF_stage.data_ptr().toint()) + qsf_word_offset = ( + ctm.Int32(qsf_stage_chunk * 512) + + lane_idx * ctm.Int32(16) + + qsf_subp * ctm.Int32(4) + ) + qsf_word_ptr = ctm.inttoptr(qsf_stage_ptr + ctm.Int64(qsf_word_offset), 3, ctm.Int32) + qsf_word = qsf_word_ptr.load() + qsf_row_id = base_row_id + qsf_row_subp * ctm.Int32(32) + qsf_tmem_addr = qsf_row_id << ctm.Int32(16) | qsf_col_id + ctm.Int32(qk_kblock_idx) + qsf_tmem = ctm.inttoptr(qsf_tmem_addr, 6, ctm.Int32) + prims.tcgen05_st(prims.Tcgen05LdStShape.SHAPE_32X32B, qsf_tmem, qsf_word) + tail_stage: ctm.Constexpr = SMEM_P4_QK_FULL_DATA_STAGES + sQSF_tail = sQSF.subview(QK_SFA_STAGE_BYTES * tail_stage) + qsf_tail_ptr = ctm.Int64(sQSF_tail.data_ptr().toint()) + row_byte_offset = lane_idx * ctm.Int32(16) + qsf_subp * ctm.Int32(4) + raw_even_ptr = ctm.inttoptr(qsf_tail_ptr + ctm.Int64(row_byte_offset), 3, ctm.Int32) + raw_odd_ptr = ctm.inttoptr( + qsf_tail_ptr + ctm.Int64(ctm.Int32(512) + row_byte_offset), 3, ctm.Int32 + ) + raw_even_word = raw_even_ptr.load() + raw_odd_word = raw_odd_ptr.load() + logical_even_word = prims.prmt( + raw_even_word, + TRTLLM_QSF_EVEN_PRMT_SELECTOR, + prims.PermuteMode.DEFAULT, + hi=raw_odd_word, + ) + logical_odd_word = prims.prmt( + raw_even_word, + TRTLLM_QSF_ODD_PRMT_SELECTOR, + prims.PermuteMode.DEFAULT, + hi=raw_odd_word, + ) + qsf_row_id = base_row_id + qsf_row_subp * ctm.Int32(32) + for tail_kblock_idx in ctm.range_constexpr(2): + logical_word = logical_even_word + if cutlass.const_expr(tail_kblock_idx == 1): + logical_word = logical_odd_word + qsf_tmem_addr = qsf_row_id << ctm.Int32(16) | qsf_col_id + ctm.Int32( + QK_MMA_KBLOCKS - 2 + tail_kblock_idx + ) + qsf_tmem = ctm.inttoptr(qsf_tmem_addr, 6, ctm.Int32) + prims.tcgen05_st(prims.Tcgen05LdStShape.SHAPE_32X32B, qsf_tmem, logical_word) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + + +@cute.jit +def _make_smem_p4_qk_ksf_stage_desc( + sKSF, + sfb_stage_base: ctm.Int32, + qk_stage_idx: ctm.Constexpr, + sksf_is_compact_slot_base: ctm.Constexpr = False, + page_major_ksf: ctm.Constexpr = False, +): + if cutlass.const_expr(sksf_is_compact_slot_base): + ksf_stage_byte_offset: ctm.Constexpr = qk_stage_idx * QK_SFB_STAGE_BYTES + else: + physical_ksf_stage = sfb_stage_base + ctm.Int32(qk_stage_idx) + ksf_stage_byte_offset = _qk_compact_ring_stage_byte_offset( + physical_ksf_stage, QK_SFB_STAGE_BYTES, SMEM_P4_KSF_COMPACT_SLOT_BYTES + ) + page_scale_chunks: ctm.Constexpr = ( + SMEM_P4_QK_DATA_STAGE_SF_CHUNKS + if qk_stage_idx < SMEM_P4_QK_FULL_DATA_STAGES + else SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS + ) + s2t_stride_bytes: ctm.Constexpr = ( + SMEM_P4_QK_SFB_S2T_STRIDE_BYTES * page_scale_chunks + if page_major_ksf + else SMEM_P4_QK_SFB_S2T_STRIDE_BYTES + ) + return prims.Tcgen05SmemDesc.build( + sKSF.subview(ksf_stage_byte_offset), + leading_byte_offset=16, + stride_byte_offset=s2t_stride_bytes, + layout=0, + ) + + +@cute.jit +def _issue_smem_p4_qk_ksf_stage_from_desc( + desc_ksf_s2t, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + sfb_col_offset: ctm.Constexpr, + qk_stage_idx: ctm.Constexpr, + page_major_ksf: ctm.Constexpr = False, +) -> None: + if cutlass.const_expr(SMEM_P4_QK_SFB_WARPX2_MODE == 2): + sfb_s2t_shape, sfb_s2t_multicast = prims.S2TCopyMode.S2T_64x128b_WARPX2_02_13 + else: + sfb_s2t_shape, sfb_s2t_multicast = prims.S2TCopyMode.S2T_64x128b_WARPX2_01_23 + sfb_tmem_addr_base = base_row_id << ctm.Int32(16) | base_col_id + ctm.Int32(sfb_col_offset) + sfb_cols_per_scale_chunk: ctm.Constexpr = _compact_sfb_mma_cols_per_scale_chunk( + SMEM_P4_CTA_GROUP_M, SMEM_P4_QK_SCORE_SLOT_N + ) + stage_kblocks: ctm.Constexpr = ( + SMEM_P4_QK_DATA_STAGE_KBLOCKS + if qk_stage_idx < SMEM_P4_QK_FULL_DATA_STAGES + else QK_TAIL_DATA_STAGE_K_DIM // SMEM_P4_QK_MMA_KBLOCK_DIM + ) + for local_kblock_idx in ctm.range_constexpr(stage_kblocks): + for scale_chunk_idx in ctm.range_constexpr(SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK): + logical_local_scale_word_idx: ctm.Constexpr = ( + local_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK + scale_chunk_idx + ) + physical_local_scale_word_idx: ctm.Constexpr = logical_local_scale_word_idx + # One data stage at a time lives in the K scale ring. + sfb_tmem_ptr = ctm.inttoptr( + sfb_tmem_addr_base + + ctm.Int32(logical_local_scale_word_idx * sfb_cols_per_scale_chunk), + 6, + ctm.Int32, + ) + source_scale_word_stride: ctm.Constexpr = ( + 1 if page_major_ksf else SMEM_P4_QK_SFB_N_CHUNKS + ) + source_byte_offset: ctm.Constexpr = ( + physical_local_scale_word_idx * source_scale_word_stride * (512 // 16) + + SMEM_P4_QK_COMPACT_SFB_S2T_BYTE_OFFSET + ) + if prims.elect_sync(): + prims.tcgen05_cp( + sfb_s2t_shape, + sfb_tmem_ptr, + desc_ksf_s2t + ctm.Int64(source_byte_offset), + group=prims.CTAGroup.CTA_2, + multicast=sfb_s2t_multicast, + ) + + +@cute.jit +def _issue_smem_p4_qk_ksf_range_from_base( + sKSF, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + sfb_col_offset: ctm.Constexpr, + sfb_stage_base: ctm.Int32, + qk_stage_start: ctm.Constexpr, + qk_stage_count: ctm.Constexpr, + sksf_is_compact_slot_base: ctm.Constexpr = False, + page_major_ksf: ctm.Constexpr = False, +) -> None: + desc_ksf_s2t_root = _make_smem_p4_qk_ksf_stage_desc( + sKSF, sfb_stage_base, qk_stage_start, sksf_is_compact_slot_base, page_major_ksf + ) + for qk_stage_idx in ctm.range_constexpr(qk_stage_start, qk_stage_start + qk_stage_count): + desc_ksf_s2t = desc_ksf_s2t_root.advance_start_address( + (qk_stage_idx - qk_stage_start) * QK_SFB_STAGE_BYTES + ) + _issue_smem_p4_qk_ksf_stage_from_desc( + desc_ksf_s2t, + base_col_id, + base_row_id, + sfb_col_offset, + qk_stage_idx, + page_major_ksf, + ) + + +@cute.jit +def _issue_smem_p4_pv_vsf_ntile_from_base( + sVSF, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + v_stage: ctm.Int32, + n_tile_idx: ctm.Constexpr, +) -> None: + """Copy one PV N tile's V scales (all four kblocks) into the V scale columns.""" + sfb_s2t_shape, sfb_s2t_multicast = prims.S2TCopyMode.S2T_64x128b_WARPX2_01_23 + sVSF_stage = sVSF.subview(ctm.Int32(SMEM_P4_V_SFB_TMA_STAGE_BYTES) * v_stage) + desc_vsf_s2t = prims.Tcgen05SmemDesc.build( + sVSF_stage, + leading_byte_offset=16, + stride_byte_offset=SMEM_P4_PV_SFB_S2T_STRIDE_BYTES, + layout=0, + ) + vsf_tmem_addr_base = base_row_id << ctm.Int32(16) | base_col_id + ctm.Int32( + SMEM_P4_VSF_TMEM_COL_OFFSET + ) + total_n_chunks: ctm.Constexpr = SMEM_P4_PV_SFB_N_CHUNKS * SMEM_P4_N_OUT_TILES + for global_kblock_idx in ctm.range_constexpr(SMEM_P4_PV_KBLOCKS_PER_LI): + vsf_tmem_ptr = ctm.inttoptr( + vsf_tmem_addr_base + ctm.Int32(global_kblock_idx * SMEM_P4_PV_SFB_COLS_PER_KBLOCK), + 6, + ctm.Int32, + ) + source_byte_offset: ctm.Constexpr = ( + global_kblock_idx * total_n_chunks + n_tile_idx * SMEM_P4_PV_SFB_N_CHUNKS + ) * (512 // 16) + if prims.elect_sync(): + prims.tcgen05_cp( + sfb_s2t_shape, + vsf_tmem_ptr, + desc_vsf_s2t + ctm.Int64(source_byte_offset), + group=prims.CTAGroup.CTA_2, + multicast=sfb_s2t_multicast, + ) + + +@cute.jit +def _issue_smem_p4_qk_stage_burst_from_base( + sQ, + sK, + sQSF, + sKSF, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + score_slot_idx: ctm.Int32, + qk_kblock_start: ctm.Constexpr, + qk_kblocks: ctm.Constexpr, + acc_col_offset: ctm.Int32, + sfa_col_offset_base: ctm.Constexpr, + sfb_col_offset_base: ctm.Constexpr, + use_split_operand_stages: ctm.Constexpr = False, + use_resident_q_stages: ctm.Constexpr = False, + b_stage_base: ctm.Int32 = 0, + split_b_stage_base: ctm.Constexpr = NUM_QK_AB_STAGE, + fence_after: ctm.Constexpr = True, + b_smem_is_compact_slot_base: ctm.Constexpr = False, +) -> None: + qk_sfb_scale_stride_cols: ctm.Constexpr = SMEM_P4_QK_SFB_COLS_PER_KBLOCK + if cutlass.const_expr(use_resident_q_stages): + issue_stage_limit: ctm.Constexpr = SMEM_P4_QK_DATA_STAGES + issue_stage_start: ctm.Constexpr = qk_kblock_start // SMEM_P4_QK_DATA_STAGE_KBLOCKS + requested_stage_count: ctm.Constexpr = ( + qk_kblocks + SMEM_P4_QK_DATA_STAGE_KBLOCKS - 1 + ) // SMEM_P4_QK_DATA_STAGE_KBLOCKS + issue_stage_count: ctm.Constexpr = min( + requested_stage_count, issue_stage_limit - issue_stage_start + ) + for stage in ctm.range_constexpr(issue_stage_start, issue_stage_start + issue_stage_count): + stage_kblock_start: ctm.Constexpr = stage * SMEM_P4_QK_DATA_STAGE_KBLOCKS + stage_kblocks: ctm.Constexpr = ( + SMEM_P4_QK_DATA_STAGE_KBLOCKS + if stage < SMEM_P4_QK_FULL_DATA_STAGES + else QK_TAIL_EFFECTIVE_MMA_KBLOCKS + ) + first_order_residual_tail: ctm.Constexpr = stage == SMEM_P4_QK_FULL_DATA_STAGES + _issue_smem_p4_qk_score_slot_from_base( + sQ, + sK, + sQSF, + sKSF, + base_col_id, + base_row_id, + stage, + score_slot_idx=score_slot_idx, + clear_accumulator=stage_kblock_start == 0, + qk_kblock_idx=stage_kblock_start, + qk_kblocks=stage_kblocks, + acc_col_offset=acc_col_offset, + sfa_col_offset=sfa_col_offset_base, + sfb_col_offset=sfb_col_offset_base, + stage_sfa_from_smem=False, + stage_sfb_from_smem=False, + use_split_operand_stages=True, + a_stage=stage, + b_stage=b_stage_base + stage, + use_stage_local_operand_kblock_start=True, + stage_local_operand_kblock_start=0, + stage_local_scale_kblock_start=stage_kblock_start, + reuse_b_first_kblock=False, + first_order_residual_tail=first_order_residual_tail, + b_smem_is_compact_slot_base=b_smem_is_compact_slot_base, + ) + else: + for burst_kblock_idx in ctm.range_constexpr(qk_kblocks): + qk_kblock_idx: ctm.Constexpr = qk_kblock_start + burst_kblock_idx + stage: ctm.Constexpr = qk_kblock_idx % NUM_QK_AB_STAGE + if cutlass.const_expr(use_split_operand_stages): + a_stage: ctm.Constexpr = stage + b_stage: ctm.Constexpr = split_b_stage_base + stage + else: + a_stage: ctm.Constexpr = 0 + b_stage: ctm.Constexpr = 0 + _issue_smem_p4_qk_score_slot_from_base( + sQ, + sK, + sQSF, + sKSF, + base_col_id, + base_row_id, + stage, + score_slot_idx=score_slot_idx, + clear_accumulator=qk_kblock_idx == 0, + qk_kblock_idx=qk_kblock_idx, + acc_col_offset=acc_col_offset, + sfa_col_offset=sfa_col_offset_base, + sfb_col_offset=sfb_col_offset_base + burst_kblock_idx * qk_sfb_scale_stride_cols, + stage_sfa_from_smem=False, + stage_sfb_from_smem=False, + use_split_operand_stages=use_split_operand_stages, + a_stage=a_stage, + b_stage=b_stage, + ) + if cutlass.const_expr(fence_after): + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + + +SMEM_P4_N256_SCORE_HALVES = 2 +SMEM_P4_N256_SCORE_N_PER_COL_BAND = SMEM_P4_BMM1_N // SMEM_P4_TMEM_WARP_N +SMEM_P4_N256_SCORE_N_PER_PRODUCER = SMEM_P4_N256_SCORE_N_PER_COL_BAND // SMEM_P4_N256_SCORE_HALVES +SMEM_P4_N256_PSF_S2T_ROWS = SMEM_P4_BMM1_M +SMEM_P4_N256_PSF_S2T_COLS = SMEM_P4_TMEM_WARP_N * SMEM_P4_N256_SCORE_HALVES +SMEM_P4_N256_PSF_S2T_BANK_WORDS = SMEM_P4_N256_PSF_S2T_ROWS * SMEM_P4_N256_PSF_S2T_COLS +SMEM_P4_N256_PSF_S2T_BANKS = SMEM_P4_TMEM_BANKS +if ( + SMEM_P4_N256_PSF_S2T_ROWS != 64 + or SMEM_P4_N256_PSF_S2T_COLS != 4 + or SMEM_P4_N256_PSF_S2T_BANKS * SMEM_P4_N256_PSF_S2T_BANK_WORDS + != SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS +): + raise ValueError( + "row-major PSF S2T banks must exactly alias the existing 2 KiB sPSF allocation" + ) + + +@cute.jit +def _smem_p4_n256_psf_source_word_offset( + bank_idx: ctm.Int32, + score_half_idx: ctm.Int32, + local_warp_idx: ctm.Int32, + lane_idx: ctm.Int32, +) -> ctm.Int32: + row_band = local_warp_idx & ctm.Int32(1) + col_band = local_warp_idx >> ctm.Int32(1) + local_row = row_band * ctm.Int32(32) + lane_idx + return ( + bank_idx * ctm.Int32(SMEM_P4_N256_PSF_S2T_BANK_WORDS) + + local_row * ctm.Int32(SMEM_P4_N256_PSF_S2T_COLS) + + col_band * ctm.Int32(SMEM_P4_N256_SCORE_HALVES) + + score_half_idx + ) + + +@cute.jit +def _cross_arrive_smem_p4_cta2_mbar(mbar) -> None: + for peer_cta_rank in ctm.range_constexpr(CLUSTER_SHAPE_MNK[0]): + peer_mbar = _mapa_shared_cluster(mbar, peer_cta_rank) + _mbarrier_arrive_release_cta_shared_cluster(peer_mbar) + + +@cute.jit +def _arrive_smem_p4_mapped_leader_mbar(leader_mbar) -> None: + _mbarrier_arrive_release_cta_shared_cluster(leader_mbar) + + +@cute.jit +def _store_smem_p4_n256_rowsum_part( + sRowSumParts, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + producer_warp_base: ctm.Int32, + bank_idx: ctm.Int32, + score_half_idx: ctm.Int32, + value: ctm.Float32, +) -> None: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + row_band = local_warp_idx & ctm.Int32(1) + col_band = local_warp_idx >> ctm.Int32(1) + local_row = row_band * ctm.Int32(32) + lane_idx + slot_idx = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + score_half_idx + row_state_offset = slot_idx * ctm.Int32( + SMEM_P4_SCORE_SLOT_ROW_STATE_STRIDE + ) + col_band * ctm.Int32(SMEM_P4_BMM1_M) + sRowSumParts[row_state_offset + local_row] = value + + +@cute.jit +def _smem_p4_qk_score_group_col( + score_slot_idx: ctm.Constexpr, col_band: ctm.Int32, local_group_idx: ctm.Constexpr +) -> ctm.Int32: + return ctm.Int32(score_slot_idx * SMEM_P4_QK_SCORE_SLOT_STRIDE + local_group_idx * SF_VEC_SIZE) + + +@cute.jit +def _smem_p4_p4_materialize_group_col( + score_slot_idx: ctm.Constexpr, col_band: ctm.Int32, local_group_idx: ctm.Constexpr +) -> ctm.Int32: + group_col = _smem_p4_qk_score_group_col(score_slot_idx, col_band, local_group_idx) + if cutlass.const_expr(local_group_idx == 1): + group_col = ctm.Int32(SMEM_P4_P4_SCORE_READ_GROUP1_COL) + elif cutlass.const_expr(local_group_idx == 2): + group_col = ctm.Int32(SMEM_P4_P4_SCORE_READ_GROUP2_COL) + return group_col + + +@cute.jit +def _float_to_ordered_u32_for_atomic_max(value: ctm.Float32) -> ctm.Uint32: + bits = prims.mov_b32(value, target_type=ctm.Int32) + sign_mask = bits >> ctm.Int32(31) | ctm.Int32(2147483648) + encoded = bits ^ sign_mask + return prims.mov_b32(encoded, target_type=ctm.Uint32) + + +@cute.jit +def _ordered_u32_to_float_after_atomic_max(value: ctm.Uint32) -> ctm.Float32: + encoded = prims.mov_b32(value, target_type=ctm.Int32) + sign_mask = ~(encoded >> ctm.Int32(31)) | ctm.Int32(2147483648) + bits = encoded ^ sign_mask + return prims.mov_b32(bits, target_type=ctm.Float32) + + +@cute.jit +def _smem_atomic_max_ordered_u32(pointer, value: ctm.Uint32) -> None: + prims.atomicrmw( + prims.AtomicOp.MAX, + pointer, + value, + syncscope=prims.MemScope.CTA, + space=prims.SharedSpace.shared_cta, + ) + + +@cute.jit +def _tcgen05_ld_red_32x32b_x16_max_f32(tmem) -> tuple: + tmem_addr = tmem.toint(ctm.Int32) + return cute_inline_ptx( + "tcgen05.ld.red.sync.aligned.32x32b.x16.max.f32 " + "{{$w0}, {$w1}, {$w2}, {$w3}, {$w4}, {$w5}, {$w6}, {$w7}, " + "{$w8}, {$w9}, {$w10}, {$w11}, {$w12}, {$w13}, {$w14}, {$w15}}, " + "{$w16}, [{$r0}];", + write_only_types=[ctm.Int32] * 17, + read_only_args=[tmem_addr], + ) + + +@cute.jit +def _publish_p4_n256_atomic_rowmax( + sAtomicRunningRowMax, + tile_row_max: ctm.Float32, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + producer_warp_base: ctm.Int32, +) -> None: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + row_band = local_warp_idx & ctm.Int32(1) + local_row = row_band * ctm.Int32(32) + lane_idx + _smem_atomic_max_ordered_u32( + sAtomicRunningRowMax.data_ptr() + local_row, + _float_to_ordered_u32_for_atomic_max(tile_row_max), + ) + + +@cute.jit +def _prepare_p4_n256_score_half_tmem_addresses( + warp_idx: ctm.Int32, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + score_bank_col_offset: ctm.Int32, +) -> tuple: + local_warp_idx = warp_idx & ctm.Int32(3) + row_band = local_warp_idx & ctm.Int32(1) + col_band = local_warp_idx >> ctm.Int32(1) + score_half_idx = warp_idx >> ctm.Int32(2) + row_id_with_warp_offset = base_row_id + col_band * ctm.Int32(64) + row_band * ctm.Int32(32) + score_half_tmem_cols: ctm.Constexpr = SMEM_P4_BMM1_ACC_TMEM_COLS // SMEM_P4_N256_SCORE_HALVES + score_half_groups_per_col_band: ctm.Constexpr = ( + SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N + ) + score_half_col_offset = ctm.Int32(score_bank_col_offset) + score_half_idx * ctm.Int32( + score_half_tmem_cols + ) + score_tmem_addresses = [] + for local_group_idx in ctm.range_constexpr(score_half_groups_per_col_band): + group_col = _smem_p4_p4_materialize_group_col(0, col_band, local_group_idx) + col_id = base_col_id + score_half_col_offset + ctm.Int32(group_col) + score_tmem_addresses.append(row_id_with_warp_offset << ctm.Int32(16) | col_id) + return tuple(score_tmem_addresses) + + +@cute.jit +def _load_p4_n256_score_half_from_tmem( + sAtomicRunningRowMax, + qk_full_mbar, + score_tmem_addresses, + producer_col_band: ctm.Int32, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + kv_tile_idx: ctm.Int32, + valid_k: ctm.Int32, + mask_valid_k: ctm.Constexpr, + qk_phase: ctm.Int32, + wait_for_qk_full: ctm.Constexpr, + producer_warp_base: ctm.Int32, + score_half_idx: ctm.Int32, +) -> tuple: + col_band = producer_col_band + if cutlass.const_expr(wait_for_qk_full): + while not prims.mbarrier_try_wait_parity(qk_full_mbar, qk_phase, time_limit=10000000): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + score_half_n: ctm.Constexpr = SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES + score_half_groups_per_col_band: ctm.Constexpr = ( + score_half_n // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N + ) + tile_base = kv_tile_idx * ctm.Int32(SMEM_P4_BMM1_N) + pending_score_groups = [] + pending_group_stats = [] + for local_group_idx in ctm.range_constexpr(score_half_groups_per_col_band): + tmem = ctm.inttoptr(score_tmem_addresses[local_group_idx], 6, ctm.Float32) + if cutlass.const_expr(mask_valid_k): + group_scores = prims.tcgen05_ld( + prims.Tcgen05LdStShape.SHAPE_32X32B, tmem, num=SF_VEC_SIZE + ) + pending_score_groups.append(group_scores) + pending_group_stats.append(None) + else: + regs = _tcgen05_ld_red_32x32b_x16_max_f32(tmem) + pending_score_groups.append(regs) + pending_group_stats.append(regs[SF_VEC_SIZE]) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + score_values = [] + group_max_values = [] + group_valid_values = [] + tile_row_max = ctm.Float32(-ctm.Float32.inf) + for local_group_idx in ctm.range_constexpr(score_half_groups_per_col_band): + pending_scores = pending_score_groups[local_group_idx] + if cutlass.const_expr(mask_valid_k): + group_scores = pending_scores + group_max = ctm.Float32(-ctm.Float32.inf) + group_base = ( + tile_base + + col_band * ctm.Int32(SMEM_P4_N256_SCORE_N_PER_COL_BAND) + + score_half_idx * ctm.Int32(SMEM_P4_N256_SCORE_N_PER_PRODUCER) + + ctm.Int32(local_group_idx * SF_VEC_SIZE) + ) + group_valid = ctm.Boolean(group_base < valid_k) + for elem in ctm.range_constexpr(SF_VEC_SIZE): + global_k_idx = group_base + ctm.Int32(elem) + score = _mask_invalid_score(group_scores[elem], global_k_idx, valid_k) + score_values.append(score) + group_max = cute.arch.fmax(group_max, score) + else: + group_scores = ctm.Vector.from_elements( + tuple((pending_scores[elem].bitcast(ctm.Float32) for elem in range(SF_VEC_SIZE))), + dtype=ctm.Float32, + ) + group_max = pending_group_stats[local_group_idx].bitcast(ctm.Float32) + for elem in ctm.range_constexpr(SF_VEC_SIZE): + score_values.append(group_scores[elem]) + group_valid = ctm.Boolean(True) + group_max_values.append(group_max) + group_valid_values.append(group_valid) + tile_row_max = cute.arch.fmax(tile_row_max, group_max) + if cutlass.const_expr(mask_valid_k): + if not ctm.Boolean(tile_base < valid_k): + tile_row_max = ctm.Float32(-ctm.Float32.inf) + _publish_p4_n256_atomic_rowmax( + sAtomicRunningRowMax, tile_row_max, warp_idx, tidx, producer_warp_base + ) + dsl_hints.cfence() + return ( + ctm.Vector.from_elements(tuple(score_values), dtype=ctm.Float32), + ctm.Vector.from_elements(tuple(group_max_values), dtype=ctm.Float32), + ctm.Vector.from_elements(tuple(group_valid_values), dtype=ctm.Boolean), + ) + + +@cute.jit +def _runtime_load_final_p4_n256_score_half_from_tmem( + sAtomicRunningRowMax, + qk_full_mbar, + score_tmem_addresses, + producer_col_band: ctm.Int32, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + kv_tile_idx: ctm.Int32, + valid_k: ctm.Int32, + tail_needs_mask: ctm.Boolean, + qk_phase: ctm.Int32, + producer_warp_base: ctm.Int32, + score_half_idx: ctm.Int32, +) -> tuple: + score_half_n: ctm.Constexpr = SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES + score_half_groups_per_col_band: ctm.Constexpr = ( + score_half_n // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N + ) + score_value_count: ctm.Constexpr = score_half_groups_per_col_band * SF_VEC_SIZE + score_values = ctm.Vector.from_elements( + tuple((ctm.Float32(0.0) for _ in range(score_value_count))), dtype=ctm.Float32 + ) + group_max_values = ctm.Vector.from_elements( + tuple((ctm.Float32(-ctm.Float32.inf) for _ in range(score_half_groups_per_col_band))), + dtype=ctm.Float32, + ) + group_valid_values = ctm.Vector.from_elements( + tuple((ctm.Boolean(False) for _ in range(score_half_groups_per_col_band))), + dtype=ctm.Boolean, + ) + if tail_needs_mask: + score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( + sAtomicRunningRowMax, + qk_full_mbar, + score_tmem_addresses, + producer_col_band, + warp_idx, + tidx, + kv_tile_idx=kv_tile_idx, + valid_k=valid_k, + mask_valid_k=True, + qk_phase=qk_phase, + wait_for_qk_full=True, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + ) + else: + score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( + sAtomicRunningRowMax, + qk_full_mbar, + score_tmem_addresses, + producer_col_band, + warp_idx, + tidx, + kv_tile_idx=kv_tile_idx, + valid_k=valid_k, + mask_valid_k=False, + qk_phase=qk_phase, + wait_for_qk_full=True, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + ) + return (score_values, group_max_values, group_valid_values) + + +@cute.jit +def _load_p4_n256_atomic_running_rowmax( + sAtomicRunningRowMax, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + producer_warp_base: ctm.Int32, +) -> ctm.Float32: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + row_band = local_warp_idx & ctm.Int32(1) + local_row = row_band * ctm.Int32(32) + lane_idx + return _ordered_u32_to_float_after_atomic_max(sAtomicRunningRowMax[local_row]) + + +@cute.jit +def _load_p4_n256_atomic_running_rowmax_with_validity( + sAtomicRunningRowMax, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + producer_warp_base: ctm.Int32, +) -> tuple: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + row_band = local_warp_idx & ctm.Int32(1) + local_row = row_band * ctm.Int32(32) + lane_idx + ordered_rowmax = sAtomicRunningRowMax[local_row] + new_row_valid = ordered_rowmax != ctm.Uint32(8388607) + new_row_max = _ordered_u32_to_float_after_atomic_max(ordered_rowmax) + return (new_row_max, new_row_valid) + + +@cute.jit +def _select_smem_p4_lazy_anchor_with_validity( + new_row_max: ctm.Float32, + new_row_valid: ctm.Boolean, + running_row_anchor: ctm.Float32, + softmax_scale_log2: ctm.Float32, +) -> tuple: + new_row_anchor = running_row_anchor + lane_rebase = ctm.Boolean(False) + anchor_delta_log2 = ctm.Float32(0.0) + if new_row_valid: + negative_infinity = ctm.Float32(-ctm.Float32.inf) + if running_row_anchor == negative_infinity: + new_row_anchor = new_row_max + else: + anchor_delta_log2 = (new_row_max - running_row_anchor) * softmax_scale_log2 + if anchor_delta_log2 > ctm.Float32(SMEM_P4_LAZY_ANCHOR_REBASE_LOG2): + new_row_anchor = new_row_max + lane_rebase = ctm.Boolean(True) + return (new_row_anchor, lane_rebase, anchor_delta_log2) + + +@cute.jit +def _select_smem_p4_lazy_anchor_runtime( + new_row_max: ctm.Float32, + running_row_anchor: ctm.Float32, + softmax_scale_log2: ctm.Float32, + has_previous_tile: ctm.Boolean, +) -> tuple: + anchor_delta_log2 = (new_row_max - running_row_anchor) * softmax_scale_log2 + lane_rebase = has_previous_tile & ( + anchor_delta_log2 > ctm.Float32(SMEM_P4_LAZY_ANCHOR_REBASE_LOG2) + ) + new_row_anchor = running_row_anchor + if ~has_previous_tile | lane_rebase: + new_row_anchor = new_row_max + return (new_row_anchor, lane_rebase, anchor_delta_log2) + + +@cute.jit +def _compute_smem_p4_lazy_rebase_scale( + lane_rebase: ctm.Boolean, anchor_delta_log2: ctm.Float32 +) -> ctm.Float32: + rebase_scale = ctm.Float32(1.0) + if lane_rebase: + rebase_scale = cute.exp2(-anchor_delta_log2, fastmath=True) + return rebase_scale + + +@cute.jit +def _smem_p4_direct_a_word_offset( + p_stage_base: ctm.Int32, + score_half_idx: ctm.Int32, + local_row: ctm.Int32, + col_band: ctm.Int32, + local_chunk_idx: ctm.Constexpr, +) -> ctm.Int32: + physical_k64_band = score_half_idx + raw_word = ( + p_stage_base * ctm.Int32(SMEM_P4_P4_SMEM_STAGE_WORDS) + + col_band * ctm.Int32(SMEM_P4_P4_SMEM_STAGE_WORDS) + + local_row * ctm.Int32(SMEM_P4_P4_SMEM_ROW_STRIDE_BYTES // 4) + + physical_k64_band * ctm.Int32(SMEM_P4_BMM2_STAGE_K // 16) + + ctm.Int32(local_chunk_idx * 4) + ) + raw_byte = raw_word << ctm.Int32(2) + swizzle = (raw_byte >> ctm.Int32(7) & ctm.Int32(3)) << ctm.Int32(4) + return (raw_byte ^ swizzle) >> ctm.Int32(2) + + +@cute.jit +def _store_p4_n256_score_half_to_direct_a( + sP, + p_stage_base: ctm.Int32, + score_half_idx: ctm.Int32, + local_row: ctm.Int32, + col_band: ctm.Int32, + p4_chunk0_word0: ctm.Int32, + p4_chunk0_word1: ctm.Int32, + p4_chunk0_word2: ctm.Int32, + p4_chunk0_word3: ctm.Int32, + p4_chunk1_word0: ctm.Int32, + p4_chunk1_word1: ctm.Int32, + p4_chunk1_word2: ctm.Int32, + p4_chunk1_word3: ctm.Int32, +) -> None: + word_base0 = _smem_p4_direct_a_word_offset(p_stage_base, score_half_idx, local_row, col_band, 0) + sP[word_base0 + ctm.Int32(0)] = p4_chunk0_word0 + sP[word_base0 + ctm.Int32(1)] = p4_chunk0_word1 + sP[word_base0 + ctm.Int32(2)] = p4_chunk0_word2 + sP[word_base0 + ctm.Int32(3)] = p4_chunk0_word3 + word_base1 = word_base0 ^ ctm.Int32(4) + sP[word_base1 + ctm.Int32(0)] = p4_chunk1_word0 + sP[word_base1 + ctm.Int32(1)] = p4_chunk1_word1 + sP[word_base1 + ctm.Int32(2)] = p4_chunk1_word2 + sP[word_base1 + ctm.Int32(3)] = p4_chunk1_word3 + + +@cute.jit +def _prepack_p4_n256_score_half_before_rowmax( + score_values, + group_max_values, + group_valid_values, + warp_idx: ctm.Int32, + softmax_scale_log2: ctm.Float32, + kv_tile_idx: ctm.Int32, + valid_k: ctm.Int32, + producer_warp_base: ctm.Int32, + score_half_idx: ctm.Int32, + mask_valid_k: ctm.Constexpr = False, + use_mixed_imlp: ctm.Constexpr = False, +) -> tuple: + local_warp_idx = warp_idx - producer_warp_base + col_band = local_warp_idx >> ctm.Int32(1) + score_half_n: ctm.Constexpr = SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES + score_half_groups_per_col_band: ctm.Constexpr = ( + score_half_n // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N + ) + tile_base = kv_tile_idx * ctm.Int32(SMEM_P4_BMM1_N) + packed_p4_words = [] + group_p_sums = [] + for local_group_idx in ctm.range_constexpr(score_half_groups_per_col_band): + score_value_base: ctm.Constexpr = local_group_idx * SF_VEC_SIZE + group_scores = ctm.Vector.from_elements( + tuple((score_values[score_value_base + elem] for elem in range(SF_VEC_SIZE))), + dtype=ctm.Float32, + ) + group_base = ( + tile_base + + col_band * ctm.Int32(SMEM_P4_N256_SCORE_N_PER_COL_BAND) + + score_half_idx * ctm.Int32(SMEM_P4_N256_SCORE_N_PER_PRODUCER) + + ctm.Int32(local_group_idx * SF_VEC_SIZE) + ) + group_valid = group_valid_values[local_group_idx] + group_max = group_max_values[local_group_idx] + dsl_hints.warp_switch() + if cutlass.const_expr(use_mixed_imlp): + p_bias = ctm.Float32(LOG2_6) - group_max * softmax_scale_log2 + word0, word1, p_sum = _pack_p4_16_scores_from_group_packed_f16x2( + group_scores, + group_base, + valid_k, + group_valid, + mask_valid_k, + softmax_scale_log2, + p_bias, + elem_base=0, + ) + else: + p_bias = ctm.Float32(LOG2_6) - group_max * softmax_scale_log2 + word0, word1, p_sum = _pack_p4_16_scores_from_group( + group_scores, + group_base, + valid_k, + group_valid, + mask_valid_k, + use_mixed_imlp, + softmax_scale_log2, + p_bias, + elem_base=0, + ) + packed_p4_words.append(word0) + packed_p4_words.append(word1) + group_p_sums.append(p_sum) + return ( + ctm.Vector.from_elements(tuple(packed_p4_words), dtype=ctm.Int32), + ctm.Vector.from_elements(tuple(group_p_sums), dtype=ctm.Float32), + ) + + +@cute.jit +def _store_prepacked_p4_n256_score_half_to_direct_a( + sP, + packed_p4_words, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + producer_warp_base: ctm.Int32, + p_stage_base: ctm.Int32, + score_half_idx: ctm.Int32, +) -> None: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + row_band = local_warp_idx & ctm.Int32(1) + col_band = local_warp_idx >> ctm.Int32(1) + source_row = row_band * ctm.Int32(32) + lane_idx + _store_p4_n256_score_half_to_direct_a( + sP, + p_stage_base, + score_half_idx, + source_row, + col_band, + packed_p4_words[0], + packed_p4_words[1], + packed_p4_words[2], + packed_p4_words[3], + packed_p4_words[4], + packed_p4_words[5], + packed_p4_words[6], + packed_p4_words[7], + ) + + +@cute.jit +def _finalize_p4_n256_score_half_psf_source( + sPSF, + group_max_values, + group_p_sums, + new_row_max: ctm.Float32, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + kv_tile_idx: ctm.Int32, + valid_k: ctm.Int32, + mask_valid_k: ctm.Constexpr, + producer_warp_base: ctm.Int32, + score_half_idx: ctm.Int32, + bank_idx: ctm.Int32, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + col_band = local_warp_idx >> ctm.Int32(1) + score_half_n: ctm.Constexpr = SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES + score_half_groups_per_col_band: ctm.Constexpr = ( + score_half_n // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N + ) + tile_base = kv_tile_idx * ctm.Int32(SMEM_P4_BMM1_N) + sf_bias = -new_row_max * softmax_scale_log2 - ctm.Float32(LOG2_6) + if cutlass.const_expr(score_half_groups_per_col_band != 4): + raise ValueError("packed P_SFA conversion requires four groups per col band") + prequant_row_sum = ctm.Float32(0.0) + if cutlass.const_expr(mask_valid_k): + sf_exact_0 = ctm.Float32(0.0) + sf_exact_1 = ctm.Float32(0.0) + sf_exact_2 = ctm.Float32(0.0) + sf_exact_3 = ctm.Float32(0.0) + for local_group_idx in ctm.range_constexpr(score_half_groups_per_col_band): + group_base = ( + tile_base + + col_band * ctm.Int32(SMEM_P4_N256_SCORE_N_PER_COL_BAND) + + score_half_idx * ctm.Int32(SMEM_P4_N256_SCORE_N_PER_PRODUCER) + + ctm.Int32(local_group_idx * SF_VEC_SIZE) + ) + group_valid = ctm.Boolean(group_base < valid_k) + sf_exact = ctm.Float32(0.0) + if group_valid: + sf_exact = cute.exp2( + group_max_values[local_group_idx] * softmax_scale_log2 + sf_bias, + fastmath=True, + ) + prequant_row_sum += sf_exact * group_p_sums[local_group_idx] + if cutlass.const_expr(local_group_idx == 0): + sf_exact_0 = sf_exact + elif cutlass.const_expr(local_group_idx == 1): + sf_exact_1 = sf_exact + elif cutlass.const_expr(local_group_idx == 2): + sf_exact_2 = sf_exact + else: + sf_exact_3 = sf_exact + denominator_sf_word = _pack_e4m3x4_natural(sf_exact_0, sf_exact_1, sf_exact_2, sf_exact_3) + pv_sf_word = _pack_e4m3x4_natural( + sf_exact_0 * pv_psf_rescale, + sf_exact_1 * pv_psf_rescale, + sf_exact_2 * pv_psf_rescale, + sf_exact_3 * pv_psf_rescale, + ) + elif cutlass.const_expr(use_mixed_imlp): + sf_h2_01 = _softmax_exp2_pair_packed_f16x2( + group_max_values[0], group_max_values[1], softmax_scale_log2, sf_bias + ) + sf_h2_23 = _softmax_exp2_pair_packed_f16x2( + group_max_values[2], group_max_values[3], softmax_scale_log2, sf_bias + ) + denominator_sf_word = _pack_e4m3x4_natural_from_f16x2(sf_h2_01, sf_h2_23) + pv_sf_word = _scale_pack_e4m3x4_natural_from_f16x2(sf_h2_01, sf_h2_23, pv_psf_rescale) + prequant_row_sum = _p4_n256_prequant_rowsum_from_f16x2_psf(sf_h2_01, sf_h2_23, group_p_sums) + else: + sf_args_01 = nvvm_fma_packed_f32x2( + (group_max_values[0], group_max_values[1]), + (softmax_scale_log2, softmax_scale_log2), + (sf_bias, sf_bias), + ) + sf_args_23 = nvvm_fma_packed_f32x2( + (group_max_values[2], group_max_values[3]), + (softmax_scale_log2, softmax_scale_log2), + (sf_bias, sf_bias), + ) + sf_exact_0 = cute.exp2(sf_args_01[0], fastmath=True) + sf_exact_1 = cute.exp2(sf_args_01[1], fastmath=True) + sf_exact_2 = cute.exp2(sf_args_23[0], fastmath=True) + sf_exact_3 = cute.exp2(sf_args_23[1], fastmath=True) + denominator_sf_word = _pack_e4m3x4_natural(sf_exact_0, sf_exact_1, sf_exact_2, sf_exact_3) + pv_sf_word = _pack_e4m3x4_natural( + sf_exact_0 * pv_psf_rescale, + sf_exact_1 * pv_psf_rescale, + sf_exact_2 * pv_psf_rescale, + sf_exact_3 * pv_psf_rescale, + ) + row_sum_parts01 = nvvm_mul_packed_f32x2( + (sf_exact_0, sf_exact_1), (group_p_sums[0], group_p_sums[1]) + ) + row_sum_parts23 = nvvm_mul_packed_f32x2( + (sf_exact_2, sf_exact_3), (group_p_sums[2], group_p_sums[3]) + ) + prequant_row_sum = _sum_p4_n256_psf_row_parts( + row_sum_parts01[0], + row_sum_parts01[1], + row_sum_parts23[0], + row_sum_parts23[1], + ) + psf_source_word_offset = _smem_p4_n256_psf_source_word_offset( + bank_idx, score_half_idx, local_warp_idx, lane_idx + ) + sPSF[psf_source_word_offset] = pv_sf_word + dsl_hints.cfence() + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + dsl_hints.cfence() + return (denominator_sf_word, prequant_row_sum) + + +@cute.jit +def _runtime_finalize_final_p4_n256_score_half_psf_source( + sPSF, + group_max_values, + group_valid_values, + group_p_sums, + new_row_max: ctm.Float32, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + tail_needs_mask: ctm.Boolean, + producer_warp_base: ctm.Int32, + score_half_idx: ctm.Int32, + bank_idx: ctm.Int32, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + score_half_n: ctm.Constexpr = SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES + score_half_groups_per_col_band: ctm.Constexpr = ( + score_half_n // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N + ) + sf_bias = -new_row_max * softmax_scale_log2 - ctm.Float32(LOG2_6) + if cutlass.const_expr(score_half_groups_per_col_band != 4): + raise ValueError("packed P_SFA conversion requires four groups per col band") + denominator_sf_word = ctm.Int32(0) + pv_sf_word = ctm.Int32(0) + prequant_row_sum = ctm.Float32(0.0) + if tail_needs_mask: + sf_exact_0 = ctm.Float32(0.0) + sf_exact_1 = ctm.Float32(0.0) + sf_exact_2 = ctm.Float32(0.0) + sf_exact_3 = ctm.Float32(0.0) + for local_group_idx in ctm.range_constexpr(score_half_groups_per_col_band): + group_valid = group_valid_values[local_group_idx] + sf_exact = ctm.Float32(0.0) + if group_valid: + sf_exact = cute.exp2( + group_max_values[local_group_idx] * softmax_scale_log2 + sf_bias, + fastmath=True, + ) + prequant_row_sum += sf_exact * group_p_sums[local_group_idx] + if cutlass.const_expr(local_group_idx == 0): + sf_exact_0 = sf_exact + elif cutlass.const_expr(local_group_idx == 1): + sf_exact_1 = sf_exact + elif cutlass.const_expr(local_group_idx == 2): + sf_exact_2 = sf_exact + else: + sf_exact_3 = sf_exact + denominator_sf_word = _pack_e4m3x4_natural(sf_exact_0, sf_exact_1, sf_exact_2, sf_exact_3) + pv_sf_word = _pack_e4m3x4_natural( + sf_exact_0 * pv_psf_rescale, + sf_exact_1 * pv_psf_rescale, + sf_exact_2 * pv_psf_rescale, + sf_exact_3 * pv_psf_rescale, + ) + elif cutlass.const_expr(use_mixed_imlp): + sf_h2_01 = _softmax_exp2_pair_packed_f16x2( + group_max_values[0], group_max_values[1], softmax_scale_log2, sf_bias + ) + sf_h2_23 = _softmax_exp2_pair_packed_f16x2( + group_max_values[2], group_max_values[3], softmax_scale_log2, sf_bias + ) + denominator_sf_word = _pack_e4m3x4_natural_from_f16x2(sf_h2_01, sf_h2_23) + pv_sf_word = _scale_pack_e4m3x4_natural_from_f16x2(sf_h2_01, sf_h2_23, pv_psf_rescale) + prequant_row_sum = _p4_n256_prequant_rowsum_from_f16x2_psf(sf_h2_01, sf_h2_23, group_p_sums) + else: + sf_args_01 = nvvm_fma_packed_f32x2( + (group_max_values[0], group_max_values[1]), + (softmax_scale_log2, softmax_scale_log2), + (sf_bias, sf_bias), + ) + sf_args_23 = nvvm_fma_packed_f32x2( + (group_max_values[2], group_max_values[3]), + (softmax_scale_log2, softmax_scale_log2), + (sf_bias, sf_bias), + ) + sf_exact_0 = cute.exp2(sf_args_01[0], fastmath=True) + sf_exact_1 = cute.exp2(sf_args_01[1], fastmath=True) + sf_exact_2 = cute.exp2(sf_args_23[0], fastmath=True) + sf_exact_3 = cute.exp2(sf_args_23[1], fastmath=True) + denominator_sf_word = _pack_e4m3x4_natural(sf_exact_0, sf_exact_1, sf_exact_2, sf_exact_3) + pv_sf_word = _pack_e4m3x4_natural( + sf_exact_0 * pv_psf_rescale, + sf_exact_1 * pv_psf_rescale, + sf_exact_2 * pv_psf_rescale, + sf_exact_3 * pv_psf_rescale, + ) + row_sum_parts01 = nvvm_mul_packed_f32x2( + (sf_exact_0, sf_exact_1), (group_p_sums[0], group_p_sums[1]) + ) + row_sum_parts23 = nvvm_mul_packed_f32x2( + (sf_exact_2, sf_exact_3), (group_p_sums[2], group_p_sums[3]) + ) + prequant_row_sum = _sum_p4_n256_psf_row_parts( + row_sum_parts01[0], + row_sum_parts01[1], + row_sum_parts23[0], + row_sum_parts23[1], + ) + psf_source_word_offset = _smem_p4_n256_psf_source_word_offset( + bank_idx, score_half_idx, local_warp_idx, lane_idx + ) + sPSF[psf_source_word_offset] = pv_sf_word + dsl_hints.cfence() + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + dsl_hints.cfence() + return (denominator_sf_word, prequant_row_sum) + + +@cute.jit +def _reload_p4_n256_owned_psf_word_after_ready( + sPSF, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + producer_warp_base: ctm.Int32, + score_half_idx: ctm.Int32, + bank_idx: ctm.Int32, +) -> ctm.Int32: + local_warp_idx = warp_idx - producer_warp_base + lane_idx = tidx & ctm.Int32(31) + word_offset = _smem_p4_n256_psf_source_word_offset( + bank_idx, score_half_idx, local_warp_idx, lane_idx + ) + return sPSF[word_offset] + + +@cute.jit +def _make_p4_n256_psf_bank_to_tmem_desc( + sPSF, base_col_id: ctm.Int32, base_row_id: ctm.Int32, bank_idx: ctm.Int32 +) -> tuple: + sPSF_bank = sPSF.subview(bank_idx * ctm.Int32(SMEM_P4_N256_PSF_S2T_BANK_WORDS)) + desc_psf_s2t = prims.Tcgen05SmemDesc.build( + sPSF_bank, leading_byte_offset=16, stride_byte_offset=128, layout=0 + ) + p_sfa_col_id = ( + base_col_id + + ctm.Int32(SMEM_P4_PSF_TMEM_COL_OFFSET) + + bank_idx * ctm.Int32(SMEM_P4_P_SCALE_FEED_STAGE_COLS) + ) + p_sfa_tmem_addr = base_row_id << ctm.Int32(16) | p_sfa_col_id + p_sfa_tmem = ctm.inttoptr(p_sfa_tmem_addr, 6, ctm.Int32) + return (desc_psf_s2t, p_sfa_tmem) + + +@cute.jit +def _issue_p4_n256_psf_bank_to_tmem_from_desc(desc_psf_s2t, p_sfa_tmem) -> None: + psf_s2t_shape, psf_s2t_multicast = prims.S2TCopyMode.S2T_64x128b_WARPX2_02_13 + if prims.elect_sync(): + prims.tcgen05_cp( + psf_s2t_shape, + p_sfa_tmem, + desc_psf_s2t, + group=prims.CTAGroup.CTA_2, + multicast=psf_s2t_multicast, + ) + + +@cute.jit +def _sum_p4_n256_psf_row_parts( + row_sum_part0: ctm.Float32, + row_sum_part1: ctm.Float32, + row_sum_part2: ctm.Float32, + row_sum_part3: ctm.Float32, +) -> ctm.Float32: + row_sum = row_sum_part0 + row_sum_part1 + row_sum = row_sum + row_sum_part2 + return row_sum + row_sum_part3 + + +@cute.jit +def _p4_n256_rowsum_from_packed_psf_word(sf_word: ctm.Int32, packed_group_sums) -> ctm.Float32: + return cute_inline_ptx( + "{\n" + "\t.reg .b16 sf01, sf23;\n" + "\t.reg .b32 sf_h2_01, sf_h2_23, weighted_h2;\n" + "\t.reg .f16 weighted_lo, weighted_hi;\n" + "\t.reg .f32 weighted_lo_f32, weighted_hi_f32;\n" + "\tmov.b32 {sf01, sf23}, {$r0};\n" + "\tcvt.rn.f16x2.e4m3x2 sf_h2_01, sf01;\n" + "\tcvt.rn.f16x2.e4m3x2 sf_h2_23, sf23;\n" + "\tmul.f16x2 weighted_h2, {$r1}, sf_h2_01;\n" + "\tfma.rn.f16x2 weighted_h2, {$r2}, sf_h2_23, weighted_h2;\n" + "\tmov.b32 {weighted_lo, weighted_hi}, weighted_h2;\n" + "\tcvt.f32.f16 weighted_lo_f32, weighted_lo;\n" + "\tcvt.f32.f16 weighted_hi_f32, weighted_hi;\n" + "\tadd.f32 {$w0}, weighted_lo_f32, weighted_hi_f32;\n" + "\t}\n", + write_only_types=[ctm.Float32], + read_only_args=[sf_word, packed_group_sums[0], packed_group_sums[1]], + ) + + +@cute.jit +def _issue_smem_p4_pv_score_slot_from_base( + sP, + sV, + sPSF, + sVSF, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + p_stage: ctm.Int32, + is_first_kv_tile: ctm.Boolean, + score_slot_idx: ctm.Constexpr = -1, + v_stage: ctm.Int32 = 0, + acc_col_offset: ctm.Constexpr = SMEM_P4_O_ACC_COL_OFFSET, + p_tmem_col_offset: ctm.Int32 = -1, + n_tile_idx: ctm.Constexpr = 0, + p_layout_score_slot_idx: ctm.Constexpr = -1, + pv_kblock_start_delta: ctm.Constexpr = 0, + pv_kblocks: ctm.Constexpr = -1, + clear_accumulator_override: ctm.Constexpr = -1, + bank_idx: ctm.Int32 = -1, + p_half_idx: ctm.Constexpr = -1, + local_kblock_idx: ctm.Constexpr = -1, + global_kblock_idx: ctm.Constexpr = -1, + p_sfa_col_offset: ctm.Int32 = -1, + a_collector_op: ctm.Constexpr = None, +) -> None: + pv_stage = v_stage + p_layout_half_idx: ctm.Constexpr = ( + score_slot_idx if p_layout_score_slot_idx < 0 else p_layout_score_slot_idx + ) + physical_p_half_idx: ctm.Constexpr = p_layout_half_idx if p_half_idx < 0 else p_half_idx + physical_local_kblock_idx: ctm.Constexpr = ( + pv_kblock_start_delta if local_kblock_idx < 0 else local_kblock_idx + ) + physical_global_kblock_idx: ctm.Constexpr = ( + score_slot_idx * SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT + physical_local_kblock_idx + if global_kblock_idx < 0 + else global_kblock_idx + ) + p_bank_col_offset = p_tmem_col_offset + default_clear_pv_accumulator = ctm.Boolean(False) + if cutlass.const_expr(physical_global_kblock_idx == 0): + default_clear_pv_accumulator = is_first_kv_tile + clear_pv_accumulator = default_clear_pv_accumulator + if cutlass.const_expr(clear_accumulator_override >= 0): + clear_pv_accumulator = ctm.Boolean(clear_accumulator_override != 0) + sfb_col_offset: ctm.Constexpr = ( + SMEM_P4_VSF_TMEM_COL_OFFSET + physical_global_kblock_idx * SMEM_P4_PV_SFB_COLS_PER_KBLOCK + ) + p_stage_slot_col_offset = p_bank_col_offset + ctm.Int32( + physical_p_half_idx * SMEM_P4_P_TMEM_SLOT_STRIDE + ) + effective_p_sfa_col_offset = p_sfa_col_offset + pv_mma_kblocks: ctm.Constexpr = ( + pv_kblocks if pv_kblocks > 0 else SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT + ) + pv_mma_kblock_start: ctm.Constexpr = physical_global_kblock_idx + pv_a_tmem_col_offset = p_stage_slot_col_offset + pv_m_dim: ctm.Constexpr = SMEM_P4_CTA_GROUP_M + pv_b_major: ctm.Constexpr = 0 + pv_page_idx: ctm.Constexpr = physical_global_kblock_idx // (TRTLLM_PAGE_SIZE // MMA_KBLOCK_DIM) + pv_page_local_kblock_idx: ctm.Constexpr = physical_global_kblock_idx % ( + TRTLLM_PAGE_SIZE // MMA_KBLOCK_DIM + ) + pv_b_mma_smem_word_delta: ctm.Constexpr = pv_page_idx * (B_STAGE_BYTES // 4) + n_tile_idx * ( + SMEM_P4_V_PAGE_N_TILE_BYTES // 4 + ) + pv_b_mma_kblock_start: ctm.Constexpr = pv_page_local_kblock_idx + pv_b_mma_leading_bytes: ctm.Constexpr = SMEM_P4_PV_PAGE_B_MMA_LEADING_BYTES + pv_b_mma_stride_bytes: ctm.Constexpr = SMEM_P4_PV_PAGE_B_MMA_STRIDE_BYTES + pv_b_mma_layout: ctm.Constexpr = SMEM_P4_PV_PAGE_B_MMA_LAYOUT + _issue_mxf4nvf4_mma_tile_from_base( + sP, + sV, + sPSF, + sVSF, + base_col_id, + base_row_id, + pv_stage, + clear_accumulator=clear_pv_accumulator, + acc_col_offset=acc_col_offset, + sfa_col_offset=effective_p_sfa_col_offset, + sfb_col_offset=sfb_col_offset, + n_dim=SMEM_P4_BMM2_N, + m_dim=pv_m_dim, + mma_kblock_start=pv_mma_kblock_start, + mma_kblocks=pv_mma_kblocks, + mma_kblock_dim=MMA_KBLOCK_DIM, + mma_kblock_idesc_k_dim=MMA_KBLOCK_IDESC_K_DIM, + mma_scale_chunks_per_kblock=MMA_SCALE_CHUNKS_PER_KBLOCK, + mma_scale_vec_size=MMA_SCALE_VEC_SIZE, + a_tmem_kblock_cols=SMEM_P4_P_TMEM_KBLOCK_COLS, + stage_sfa_from_smem=False, + stage_sfb_from_smem=False, + compact_sfb_words_per_row=MMA_SCALE_CHUNKS_PER_TILE, + compact_sfb_extra_row_stride=SMEM_P4_PV_SFB_EXTRA_ROW_STRIDE, + compact_sfb_extra_col_stride=SMEM_P4_PV_SFB_EXTRA_COL_STRIDE, + compact_sfb_cp_group_one=SMEM_P4_PV_SFB_CP_GROUP_ONE, + compact_sfb_warpx2_mode=SMEM_P4_PV_SFB_WARPX2_MODE, + sfb_mma_col_delta=SMEM_P4_PV_SFB_MMA_COL_DELTA, + sfb_mma_kblock_stride_cols=SMEM_P4_PV_SFB_COLS_PER_KBLOCK, + use_split_operand_stages=True, + a_stage=p_stage, + b_stage=pv_stage, + a_stage_words=SMEM_P4_P4_SMEM_STAGE_WORDS, + b_stage_words=SMEM_P4_V_DATA_STAGE_BYTES // 4, + sfb_stage_bytes=SMEM_P4_V_SFB_TMA_STAGE_BYTES, + a_from_tmem=False, + a_tmem_col_offset=pv_a_tmem_col_offset, + b_major=pv_b_major, + b_mma_smem_word_delta=pv_b_mma_smem_word_delta, + b_mma_leading_byte_offset=pv_b_mma_leading_bytes, + b_mma_stride_byte_offset=pv_b_mma_stride_bytes, + b_mma_layout=pv_b_mma_layout, + b_mma_k_segment_offset=SMEM_P4_PV_B_MMA_K_SEGMENT_OFFSET, + use_split_mma_kblock_start=True, + a_mma_kblock_start=physical_local_kblock_idx, + b_mma_kblock_start=pv_b_mma_kblock_start, + a_mma_leading_byte_offset=SMEM_P4_P4_SMEM_MMA_LEADING_BYTE_OFFSET, + a_mma_stride_byte_offset=SMEM_P4_P4_SMEM_MMA_STRIDE_BYTE_OFFSET, + a_mma_layout=SMEM_P4_P4_SMEM_MMA_LAYOUT, + scale_mma_kblock_start=physical_global_kblock_idx, + sfb_s2t_stride_byte_offset=SMEM_P4_PV_SFB_S2T_STRIDE_BYTES, + a_sf_layout=SMEM_P4_P_SFA_LAYOUT, + a_collector_op=a_collector_op, + ) + + +@cute.jit +def _rescale_smem_p4_resident_o_in_tmem(o_acc_tmem, old_scale: ctm.Float32) -> None: + scaled_accum = ctm.Array(ctm.Float32, 128, space=ctm.AddressSpace.rmem) + c_rmem = prims.tcgen05_ld(prims.Tcgen05LdStShape.SHAPE_32X32B, o_acc_tmem, num=128) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + old_scale_f32x2 = (old_scale, old_scale) + for pair_idx in ctm.range_constexpr(64): + elem_idx: ctm.Constexpr = pair_idx * 2 + scaled_pair = nvvm_mul_packed_f32x2( + (c_rmem[elem_idx], c_rmem[elem_idx + 1]), old_scale_f32x2 + ) + scaled_accum[elem_idx] = scaled_pair[0] + scaled_accum[elem_idx + 1] = scaled_pair[1] + prims.tcgen05_st( + prims.Tcgen05LdStShape.SHAPE_32X32B, + o_acc_tmem, + scaled_accum.load(0, 128, alignment=32), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + + +@cute.jit +def _store_final_o_from_tmem( + mC_mnl: cute.Tensor, + mAccum_mnl: cute.Tensor, + mRowMax_ml: cute.Tensor, + mRowSum_ml: cute.Tensor, + tmem_holding_buf, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + bidx: ctm.Int32, + bidy: ctm.Int32, + bidz: ctm.Int32, + m: ctm.Int32, + n: ctm.Int32, + final_row_max: ctm.Float32, + final_row_sum: ctm.Float32, + final_stat_scale: ctm.Float32, + output_normalizer: ctm.Float32, + valid_m: ctm.Int32, + c_batch_stride: ctm.Int32, + c_batch_idx: ctm.Int32, + write_softmax_stats: ctm.Constexpr = False, + producer_warp_base: ctm.Constexpr = 0, + store_split: ctm.Constexpr = False, + direct_single_task: ctm.Constexpr = False, +) -> None: + # Split mode exports the FP32 partial plus row state; direct mode stores the + # normalized output. A single-task request in split mode does both halves + # that matter: direct output and row state (for the Helix stats reduction). + store_accum: ctm.Constexpr = store_split and not direct_single_task + store_row_state: ctm.Constexpr = store_split + store_direct: ctm.Constexpr = (not store_split) or direct_single_task + # A 16-bit partial buffer takes the normalized split output (see + # SPLIT_PARTIAL_DTYPES), stored like the direct output. + normalized_accum: ctm.Constexpr = store_accum and mAccum_mnl.element_type.width == 16 + gC_arr = ctm.make_array_view(mC_mnl) + gAccum_arr = ctm.make_array_view(mAccum_mnl) + vsize: ctm.Constexpr = 16 + tmem_raw_addr = tmem_holding_buf.load() + base_col_id = tmem_raw_addr & ctm.Int32(65535) + base_row_id = tmem_raw_addr >> ctm.Int32(16) + cta_rank = cute.arch.block_idx_in_cluster() + local_warp_idx = warp_idx - ctm.Int32(producer_warp_base) + row_band = local_warp_idx & ctm.Int32(1) + col_band = local_warp_idx >> ctm.Int32(1) + local_row = row_band * ctm.Int32(32) + (tidx & ctm.Int32(31)) + row_id_with_warp_offset = base_row_id + row_band * ctm.Int32(64) + col_band * ctm.Int32(32) + m_tile_idx = bidx // ctm.Int32(CLUSTER_SHAPE_MNK[0]) + row = ( + m_tile_idx * ctm.Int32(SMEM_P4_CTA_GROUP_M) + + cta_rank * ctm.Int32(SMEM_P4_BMM1_M) + + local_row + ) + # Padded packed rows (row >= valid_m) are compute lanes only. + store_row = row < valid_m + if cutlass.const_expr(store_row_state): + # Both column bands have the same row state; export each row once. + if col_band == ctm.Int32(0): + mRowMax_ml[row, bidz] = final_row_max + mRowSum_ml[row, bidz] = final_row_sum + batch_offset = bidz * m * n + c_batch_offset = c_batch_idx * c_batch_stride + if cutlass.const_expr(write_softmax_stats): + if col_band == ctm.Int32(0) and bidy == ctm.Int32(0) and store_row: + mRowMax_ml[row, bidz] = final_row_max * final_stat_scale + mRowSum_ml[row, bidz] = final_row_sum + n_tile_total = n // ctm.Int32(OUT_DIM) + subtile_cols: ctm.Constexpr = 32 + subtiles_per_n_tile: ctm.Constexpr = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N // subtile_cols + output_subtile_total = n_tile_total * ctm.Int32(subtiles_per_n_tile) + for output_subtile_idx in cutlass.range(0, output_subtile_total, 1, unroll=1): + n_tile_idx = output_subtile_idx // ctm.Int32(subtiles_per_n_tile) + subtile_in_n_tile = output_subtile_idx - n_tile_idx * ctm.Int32(subtiles_per_n_tile) + col_chunk = subtile_in_n_tile * ctm.Int32(subtile_cols) + o_acc_col_offset = ctm.Int32(SMEM_P4_O_ACC_COL_OFFSETS[0]) + n_tile_idx * ctm.Int32( + SMEM_P4_O_ACC_TILE_COLS + ) + col_base = ( + bidy * ctm.Int32(MMA_TILER_MNK[1]) + + n_tile_idx * ctm.Int32(OUT_DIM) + + col_band * ctm.Int32(SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N) + + col_chunk + ) + col_id = base_col_id + o_acc_col_offset + col_chunk + tmem_addr = row_id_with_warp_offset << ctm.Int32(16) | col_id + tmem = ctm.inttoptr(tmem_addr, 6, ctm.Float32) + c_rmem = prims.tcgen05_ld(prims.Tcgen05LdStShape.SHAPE_32X32B, tmem, num=subtile_cols) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + if cutlass.const_expr(normalized_accum): + gAccum_arr.store( + (c_rmem[0:subtile_cols] * output_normalizer).to(mAccum_mnl.element_type), + idx=row * n + col_base + batch_offset, + vector_size=subtile_cols, + alignment=32, + ) + else: + for vec_idx in ctm.range_constexpr(subtile_cols // vsize): + linear_idx = row * n + col_base + vec_idx * vsize + batch_offset + vec_f32 = c_rmem[vec_idx * vsize : vec_idx * vsize + vsize] + if cutlass.const_expr(store_accum): + gAccum_arr.store( + vec_f32 * final_stat_scale, + idx=linear_idx, + vector_size=vsize, + alignment=32, + ) + if cutlass.const_expr(store_direct): + if store_row: + output_linear_idx = row * n + col_base + c_batch_offset + output_vec_f32 = c_rmem[0:subtile_cols] + output_vec = (output_vec_f32 * output_normalizer).to(mC_mnl.element_type) + # The decode output is consumed after this kernel and has no in-kernel + # reuse. Avoid allocating its streaming write in L1; on GR100 this + # preserves the 256-bit vector store while lowering it to + # STG.E.NA.ENL2.256. + gC_arr.subview(output_linear_idx).data_ptr().nvvm_store_ext( + output_vec, + evict=prims.EvictPriority.NOALLOCATE, + ) + + +@cute.jit +def _dealloc_tmem_cluster( + tmem_dealloc_mbar, + tmem_holding_buf, + warp_idx: ctm.Int32, + cta_rank: ctm.Int32, + num_tmem_alloc_cols: ctm.Constexpr = NUM_TMEM_ALLOC_COLS, + dealloc_warp_id: ctm.Constexpr = 0, +) -> None: + if warp_idx == ctm.Int32(dealloc_warp_id): + acc_tmem_ptr = ctm.inttoptr(tmem_holding_buf.load(), 6, ctm.Float32) + prims.tcgen05_relinquish_alloc_permit(group=prims.CTAGroup.CTA_2) + peer_cta_rank = cta_rank ^ 1 + peer_mbar = _mapa_shared_cluster(tmem_dealloc_mbar, peer_cta_rank) + _mbarrier_arrive_release_cta_shared_cluster(peer_mbar) + while not prims.mbarrier_try_wait_parity(tmem_dealloc_mbar, 0, time_limit=10000000): + pass + prims.tcgen05_dealloc( + acc_tmem_ptr, + num_tmem_alloc_cols, + is_exclusive=SMEM_P4_TMEM_ALLOC_EXCLUSIVE, + group=prims.CTAGroup.CTA_2, + ) + + +@cute.jit +def _runtime_issue_pv_tile_slot( + sP, + sV, + sPSF, + sVSF, + tmem_base_col_id: ctm.Int32, + tmem_base_row_id: ctm.Int32, + pair_overlap_sync_mbars, + p4_smem_slot0_mbars, + pv_full_mbar, + v_smem_mbars, + stream_li_idx: ctm.Int32, + stream_li_total: ctm.Int32, + bank_idx: ctm.Int32, + initial_tile: ctm.Boolean, +) -> None: + v_stage = stream_li_idx % ctm.Int32(SMEM_P4_V_PIPELINE_STAGES) + # QK of this tile is complete: its producers read the scores before arriving + # on P_SOURCE_READY below. Waiting on qk_full here instead could miss a phase, + # because the QK two tiles ahead already targets the same bank. + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _issue_smem_p4_pv_vsf_ntile_from_base(sVSF, tmem_base_col_id, tmem_base_row_id, v_stage, 0) + prefetched_psf_desc, prefetched_psf_tmem = _make_p4_n256_psf_bank_to_tmem_desc( + sPSF, tmem_base_col_id, tmem_base_row_id, bank_idx + ) + pair_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + p_source_ready_mbar = pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + while not prims.mbarrier_try_wait_parity(p_source_ready_mbar, pair_phase, time_limit=10000000): + pass + _issue_p4_n256_psf_bank_to_tmem_from_desc(prefetched_psf_desc, prefetched_psf_tmem) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + p_archive_col_offset = ctm.Int32(SMEM_P4_P_ARCHIVE_TMEM_COLS[0]) + bank_idx * ctm.Int32( + SMEM_P4_P_ARCHIVE_SLOT_STRIDE + ) + p_half0_smem_stage = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + p_sfa_col_offset = ctm.Int32(SMEM_P4_PSF_TMEM_COL_OFFSET) + bank_idx * ctm.Int32( + SMEM_P4_P_SCALE_FEED_STAGE_COLS + ) + # N-tile major: the V scale columns hold one N tile at a time, refilled in + # issue order after the previous N tile's MMAs. + for n_tile in cutlass.range_constexpr(SMEM_P4_N_OUT_TILES): + if cutlass.const_expr(n_tile > 0): + _issue_smem_p4_pv_vsf_ntile_from_base( + sVSF, tmem_base_col_id, tmem_base_row_id, v_stage, n_tile + ) + for p_half_idx in cutlass.range_constexpr(SMEM_P4_SCORE_SLOTS_PER_KV_TILE): + for local_kblock_idx in ctm.range_constexpr(SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT): + _issue_smem_p4_pv_score_slot_from_base( + sP, + sV, + sPSF, + sVSF, + tmem_base_col_id, + tmem_base_row_id, + p_half0_smem_stage + ctm.Int32(p_half_idx), + initial_tile, + v_stage=v_stage, + acc_col_offset=SMEM_P4_O_ACC_COL_OFFSETS[n_tile], + n_tile_idx=n_tile, + p_tmem_col_offset=p_archive_col_offset, + bank_idx=bank_idx, + p_half_idx=p_half_idx, + local_kblock_idx=local_kblock_idx, + global_kblock_idx=p_half_idx * SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT + + local_kblock_idx, + pv_kblocks=1, + p_sfa_col_offset=p_sfa_col_offset, + ) + if stream_li_idx > ctm.Int32(0) and stream_li_idx + ctm.Int32(1) == stream_li_total: + dsl_hints.cfence() + if prims.elect_sync(): + _cross_arrive_smem_p4_cta2_mbar( + pair_overlap_sync_mbars.subview(SMEM_P4_PAIR_FINAL_PV_ISSUED_MBAR_OFFSET + bank_idx) + ) + if prims.elect_sync(): + prims.tcgen05_commit( + pv_full_mbar.subview(bank_idx), + multicast_mask=ctm.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + prims.tcgen05_commit( + p4_smem_slot0_mbars.subview(SMEM_P4_P4_PIPELINE_STAGES + bank_idx), + multicast_mask=ctm.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + prims.tcgen05_commit( + v_smem_mbars.subview(SMEM_P4_V_PIPELINE_STAGES + v_stage), + multicast_mask=ctm.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + + +@cute.jit +def _runtime_issue_pv_tile_dispatch( + sP, + sV, + sPSF, + sVSF, + tmem_base_col_id: ctm.Int32, + tmem_base_row_id: ctm.Int32, + pair_overlap_sync_mbars, + p4_smem_slot0_mbars, + pv_full_mbar, + v_smem_mbars, + stream_li_idx: ctm.Int32, + stream_li_total: ctm.Int32, + initial_tile: ctm.Boolean, +) -> None: + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + _runtime_issue_pv_tile_slot( + sP, + sV, + sPSF, + sVSF, + tmem_base_col_id, + tmem_base_row_id, + pair_overlap_sync_mbars, + p4_smem_slot0_mbars, + pv_full_mbar, + v_smem_mbars, + stream_li_idx, + stream_li_total, + bank_idx, + initial_tile, + ) + + +@cute.jit +def _runtime_refill_qk_tile_state( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id: ctm.Int32, + tmem_base_row_id: ctm.Int32, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_next_li_idx: ctm.Int32, + qk_bulk_slot_idx: ctm.Int32, + qk_bulk_phase: ctm.Int32, + qk_consumer_phase: ctm.Int32, + qk_score_ring_idx: ctm.Int32, + split_qk15_stage1: ctm.Boolean, + page_major_ksf: ctm.Constexpr, +) -> None: + # The score bank's previous tile was read before its scores were consumed, + # which this warp waited on before issuing this QK. + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + while not prims.mbarrier_try_wait_parity( + qk_bulk_ready_mbars.subview(qk_bulk_slot_idx), qk_bulk_phase, time_limit=10000000 + ): + pass + sK_slot = sB.subview(qk_bulk_slot_idx * ctm.Int32(SMEM_P4_K_COMPACT_SLOT_BYTES // 4)) + sKSF_slot = sSFB.subview(qk_bulk_slot_idx * ctm.Int32(SMEM_P4_KSF_COMPACT_SLOT_BYTES)) + desc_ksf_s2t_root = _make_smem_p4_qk_ksf_stage_desc( + sKSF_slot, ctm.Int32(0), 0, True, page_major_ksf + ) + # Each data stage copies its K scales into the ring right before its MMAs; the + # issue order keeps a stage's copy behind the previous stage's MMAs. + _issue_smem_p4_qk_ksf_stage_from_desc( + desc_ksf_s2t_root, + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + 0, + page_major_ksf, + ) + _issue_smem_p4_qk_stage_burst_from_base( + sA, + sK_slot, + sSFA, + sKSF_slot, + tmem_base_col_id, + tmem_base_row_id, + score_slot_idx=qk_score_ring_idx, + qk_kblock_start=0, + qk_kblocks=SMEM_P4_QK_DATA_STAGE_KBLOCKS, + acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, + sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, + use_resident_q_stages=True, + b_stage_base=ctm.Int32(0), + fence_after=False, + b_smem_is_compact_slot_base=True, + ) + if split_qk15_stage1: + qk15_stage1_ready_mbar = pair_overlap_sync_mbars.subview(SMEM_P4_QK15_STAGE1_MBAR_OFFSET) + while not prims.mbarrier_try_wait_parity(qk15_stage1_ready_mbar, 0, time_limit=10000000): + pass + _issue_smem_p4_qk_ksf_stage_from_desc( + desc_ksf_s2t_root.advance_start_address(QK_SFB_STAGE_BYTES), + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + 1, + page_major_ksf, + ) + _issue_smem_p4_qk_stage_burst_from_base( + sA, + sK_slot, + sSFA, + sKSF_slot, + tmem_base_col_id, + tmem_base_row_id, + score_slot_idx=qk_score_ring_idx, + qk_kblock_start=SMEM_P4_QK_DATA_STAGE_KBLOCKS, + qk_kblocks=SMEM_P4_QK_DATA_STAGE_KBLOCKS, + acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, + sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, + use_resident_q_stages=True, + b_stage_base=ctm.Int32(0), + fence_after=False, + b_smem_is_compact_slot_base=True, + ) + qk_handle = qk_smem_consumer.wait() + qk_smem_consumer.advance() + _issue_smem_p4_qk_ksf_range_from_base( + sKSF_slot, + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + ctm.Int32(0), + SMEM_P4_QK_FULL_DATA_STAGES, + 1, + True, + page_major_ksf, + ) + _issue_smem_p4_qk_stage_burst_from_base( + sA, + sK_slot, + sSFA, + sKSF_slot, + tmem_base_col_id, + tmem_base_row_id, + score_slot_idx=qk_score_ring_idx, + qk_kblock_start=SMEM_P4_QK_TAIL_STAGE_KBLOCK_START, + qk_kblocks=QK_TAIL_DATA_STAGE_K_DIM // SMEM_P4_QK_MMA_KBLOCK_DIM, + acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, + sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, + use_resident_q_stages=True, + b_stage_base=ctm.Int32(0), + fence_after=True, + b_smem_is_compact_slot_base=True, + ) + if prims.elect_sync(): + prims.tcgen05_commit( + qk_full_mbar.subview(qk_score_ring_idx), + multicast_mask=ctm.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + qk_handle.release() + + +@cute.jit +def _runtime_refill_qk_tile_dispatch( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id: ctm.Int32, + tmem_base_row_id: ctm.Int32, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_next_li_idx: ctm.Int32, + qk_next_stage_idx: ctm.Int32, + qk_bulk_phase: ctm.Int32, + qk_consumer_phase: ctm.Int32, + split_qk15_stage1: ctm.Boolean, + page_major_ksf: ctm.Constexpr, +) -> None: + _runtime_refill_qk_tile_state( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id, + tmem_base_row_id, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_next_li_idx, + qk_next_stage_idx, + qk_bulk_phase, + qk_consumer_phase, + qk_next_li_idx % ctm.Int32(SMEM_P4_QK_SCORE_RING_SLOTS), + split_qk15_stage1, + page_major_ksf, + ) + + +@cute.jit +def _runtime_refill_qk_tile_at( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id: ctm.Int32, + tmem_base_row_id: ctm.Int32, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_li_idx: ctm.Int32, + stream_li_total: ctm.Int32, +) -> None: + """Issue QK for KV tile ``qk_li_idx`` from its K ring slot into its score bank.""" + qk_stage_idx = qk_li_idx % ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS) + qk_consumer_phase = ( + qk_li_idx + // ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + # Tile 0 is loaded on the pipeline barrier, so slot 0's bulk barrier starts at tile 3. + qk_bulk_phase = qk_consumer_phase ^ ctm.Int32(qk_stage_idx == ctm.Int32(0)) + split_qk15_stage1 = ctm.Boolean(stream_li_total == ctm.Int32(16)) & ctm.Boolean( + qk_li_idx == ctm.Int32(15) + ) + _runtime_refill_qk_tile_dispatch( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id, + tmem_base_row_id, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_li_idx, + qk_stage_idx, + qk_bulk_phase, + qk_consumer_phase, + split_qk15_stage1, + False, + ) + + +@cute.jit +def _runtime_t336_prepare_p_source( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_rowmax_read_done_mbar, + leader_score_consumed_mbar, + p4_smem_slot0_producer, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + p4_tmem_base_col_id: ctm.Int32, + p4_tmem_base_row_id: ctm.Int32, + valid_k_for_l: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + running_row_anchor: ctm.Float32, + stream_li_idx: ctm.Int32, + has_previous_tile: ctm.Boolean, + mask_valid_k: ctm.Constexpr, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + """Prepare one score half and return only scalar state across tail dispatch. + + Keeping score vectors, packed P words, and packed group sums inside this + constexpr helper avoids runtime PHIs with ambiguous packed operand types. + The caller dispatches this helper with a CTA-uniform tail predicate and + merges only FP32/Boolean scalar state. + """ + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + qk_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + qk_score_chunk_col_offset = ctm.Int32(SMEM_P4_QK_ACC_COL_OFFSET) + bank_idx * ctm.Int32( + SMEM_P4_QK_PIPELINE_SLOT_STRIDE + ) + p_stage_base = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + score_half_idx = warp_idx >> ctm.Int32(2) + producer_warp_base = score_half_idx << ctm.Int32(2) + producer_local_warp_idx = warp_idx & ctm.Int32(3) + producer_col_band = producer_local_warp_idx >> ctm.Int32(1) + is_compute_warp = score_half_idx == ctm.Int32(0) + score_tmem_addresses = _prepare_p4_n256_score_half_tmem_addresses( + warp_idx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + qk_score_chunk_col_offset, + ) + score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( + sAtomicRunningRowMax, + qk_full_mbar.subview(bank_idx), + score_tmem_addresses, + producer_col_band, + warp_idx, + tidx, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + mask_valid_k=mask_valid_k, + qk_phase=qk_phase, + wait_for_qk_full=True, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + ) + # The scores are in registers now: hand the bank back to the MMA warp. + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_score_consumed_mbar) + packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( + score_values, + group_max_values, + group_valid_values, + warp_idx, + softmax_scale_log2, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + mask_valid_k=mask_valid_k, + use_mixed_imlp=use_mixed_imlp, + ) + if is_compute_warp: + p4_smem_slot0_producer.acquire() + prims.barrier( + barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, + number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, + ) + if cutlass.const_expr(mask_valid_k): + new_row_max, new_row_valid = _load_p4_n256_atomic_running_rowmax_with_validity( + sAtomicRunningRowMax, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + ) + next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( + _select_smem_p4_lazy_anchor_with_validity( + new_row_max, + new_row_valid, + running_row_anchor, + softmax_scale_log2, + ) + ) + else: + new_row_max = _load_p4_n256_atomic_running_rowmax( + sAtomicRunningRowMax, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + ) + next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( + _select_smem_p4_lazy_anchor_runtime( + new_row_max, + running_row_anchor, + softmax_scale_log2, + has_previous_tile, + ) + ) + if is_compute_warp: + dsl_hints.cfence() + if prims.elect_sync(): + prims.mbarrier_arrive(pair_rowmax_read_done_mbar) + _store_prepacked_p4_n256_score_half_to_direct_a( + sP, + packed_p4_words, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + p_stage_base=p_stage_base, + score_half_idx=score_half_idx, + ) + if is_compute_warp: + dsl_hints.cfence() + owned_psf_word, prequant_row_sum = _finalize_p4_n256_score_half_psf_source( + sPSF, + group_max_values, + group_p_sums, + next_running_row_anchor, + warp_idx, + tidx, + softmax_scale_log2, + pv_psf_rescale, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + mask_valid_k=mask_valid_k, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + bank_idx=bank_idx, + use_mixed_imlp=use_mixed_imlp, + ) + row_sum = prequant_row_sum + return ( + new_row_max, + next_running_row_anchor, + lane_rebase, + anchor_delta_log2, + row_sum, + ) + + +@cute.jit +def _runtime_t336_producer_tile( + sAtomicRunningRowMax, + sP, + sPSF, + sScoreSlotRowMax, + sRowScaleRing, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + p4_smem_slot0_mbar_ptr, + tmem_score_slot0_consumer, + tmem_score_slot0_mbar_ptr, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + p4_tmem_base_col_id: ctm.Int32, + p4_tmem_base_row_id: ctm.Int32, + valid_k_for_l: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + running_row_max: ctm.Float32, + running_row_anchor: ctm.Float32, + carried_dq1024: ctm.Int32, + carried_old_scale: ctm.Float32, + stream_li_idx: ctm.Int32, + has_previous_tile: ctm.Constexpr, + tail_needs_mask: ctm.Boolean, + is_final_tile: ctm.Constexpr, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + qk_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + rowsum_full_mbar = pair_overlap_sync_mbars.subview(SMEM_P4_PAIR_P4_READY_MBAR_OFFSET + bank_idx) + pair_rowmax_read_done_mbar = pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + bank_idx + ) + direct_p_ready_mbar = pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_mbar = leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_dsmem_mbar = leader_pair_overlap_sync_mbars_dsmem.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + qk_score_chunk_col_offset = ctm.Int32(SMEM_P4_QK_ACC_COL_OFFSET) + bank_idx * ctm.Int32( + SMEM_P4_QK_PIPELINE_SLOT_STRIDE + ) + p_stage_base = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + score_half_idx = warp_idx >> ctm.Int32(2) + producer_warp_base = score_half_idx << ctm.Int32(2) + producer_local_warp_idx = warp_idx & ctm.Int32(3) + producer_row_band = producer_local_warp_idx & ctm.Int32(1) + producer_col_band = producer_local_warp_idx >> ctm.Int32(1) + is_compute_warp = score_half_idx == ctm.Int32(0) + score_tmem_addresses = _prepare_p4_n256_score_half_tmem_addresses( + warp_idx, p4_tmem_base_col_id, p4_tmem_base_row_id, qk_score_chunk_col_offset + ) + next_running_row_max = running_row_max + next_running_row_anchor = running_row_anchor + next_carried_dq1024 = carried_dq1024 + next_carried_old_scale = carried_old_scale + if is_compute_warp: + tmem_score_slot0_consumer.wait() + if cutlass.const_expr(is_final_tile): + score_values, group_max_values, group_valid_values = ( + _runtime_load_final_p4_n256_score_half_from_tmem( + sAtomicRunningRowMax, + qk_full_mbar.subview(bank_idx), + score_tmem_addresses, + producer_col_band, + warp_idx, + tidx, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + tail_needs_mask=tail_needs_mask, + qk_phase=qk_phase, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + ) + ) + else: + score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( + sAtomicRunningRowMax, + qk_full_mbar.subview(bank_idx), + score_tmem_addresses, + producer_col_band, + warp_idx, + tidx, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + mask_valid_k=False, + qk_phase=qk_phase, + wait_for_qk_full=True, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + ) + packed_p4_words = ctm.Vector.from_elements( + tuple((ctm.Int32(0) for _ in range(8))), dtype=ctm.Int32 + ) + group_p_sums = ctm.Vector.from_elements( + tuple((ctm.Float32(0.0) for _ in range(4))), dtype=ctm.Float32 + ) + if cutlass.const_expr(is_final_tile): + # ``tail_needs_mask`` is CTA-uniform. Keep the NaN-safe validity path + # for a genuinely partial final tile, but do not execute its per-SF16 + # predicates for a full final tile. Both cases remain in one runtime-KV + # binary; no request-dependent valid_k value enters the JIT key. + if tail_needs_mask: + packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( + score_values, + group_max_values, + group_valid_values, + warp_idx, + softmax_scale_log2, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + mask_valid_k=True, + use_mixed_imlp=use_mixed_imlp, + ) + else: + packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( + score_values, + group_max_values, + group_valid_values, + warp_idx, + softmax_scale_log2, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + mask_valid_k=False, + use_mixed_imlp=use_mixed_imlp, + ) + else: + packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( + score_values, + group_max_values, + group_valid_values, + warp_idx, + softmax_scale_log2, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + mask_valid_k=False, + use_mixed_imlp=use_mixed_imlp, + ) + if is_compute_warp: + p4_smem_slot0_producer.acquire() + prims.barrier( + barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, + number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, + ) + new_row_max = running_row_max + new_row_valid = ctm.Boolean(True) + if cutlass.const_expr(is_final_tile): + if tail_needs_mask: + new_row_max, new_row_valid = _load_p4_n256_atomic_running_rowmax_with_validity( + sAtomicRunningRowMax, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + ) + else: + new_row_max = _load_p4_n256_atomic_running_rowmax( + sAtomicRunningRowMax, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + ) + else: + new_row_max = _load_p4_n256_atomic_running_rowmax( + sAtomicRunningRowMax, warp_idx, tidx, producer_warp_base=producer_warp_base + ) + if is_compute_warp: + dsl_hints.cfence() + if prims.elect_sync(): + prims.mbarrier_arrive(pair_rowmax_read_done_mbar) + lane_rebase = ctm.Boolean(False) + anchor_delta_log2 = ctm.Float32(0.0) + if cutlass.const_expr(is_final_tile): + if tail_needs_mask: + next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( + _select_smem_p4_lazy_anchor_with_validity( + new_row_max, + new_row_valid, + running_row_anchor, + softmax_scale_log2, + ) + ) + else: + next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( + _select_smem_p4_lazy_anchor_runtime( + new_row_max, + running_row_anchor, + softmax_scale_log2, + has_previous_tile, + ) + ) + _store_prepacked_p4_n256_score_half_to_direct_a( + sP, + packed_p4_words, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + p_stage_base=p_stage_base, + score_half_idx=score_half_idx, + ) + if prims.elect_sync(): + prims.mbarrier_arrive(direct_p_ready_mbar) + if is_compute_warp: + dsl_hints.cfence() + if cutlass.const_expr(not is_final_tile): + next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( + _select_smem_p4_lazy_anchor_runtime( + new_row_max, running_row_anchor, softmax_scale_log2, has_previous_tile + ) + ) + current_old_scale = ctm.Float32(1.0) + warp_rebase = cute.arch.vote_any_sync(lane_rebase) + next_running_row_max = new_row_max + if cutlass.const_expr(is_final_tile): + owned_psf_word, prequant_row_sum = _runtime_finalize_final_p4_n256_score_half_psf_source( + sPSF, + group_max_values, + group_valid_values, + group_p_sums, + next_running_row_anchor, + warp_idx, + tidx, + softmax_scale_log2, + pv_psf_rescale, + tail_needs_mask=tail_needs_mask, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + bank_idx=bank_idx, + use_mixed_imlp=use_mixed_imlp, + ) + else: + owned_psf_word, prequant_row_sum = _finalize_p4_n256_score_half_psf_source( + sPSF, + group_max_values, + group_p_sums, + next_running_row_anchor, + warp_idx, + tidx, + softmax_scale_log2, + pv_psf_rescale, + kv_tile_idx=stream_li_idx, + valid_k=valid_k_for_l, + mask_valid_k=False, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + bank_idx=bank_idx, + use_mixed_imlp=use_mixed_imlp, + ) + if has_previous_tile: + common_source_ready_pred = prims.elect_sync() & ~warp_rebase + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) + if warp_rebase: + pv_done_li_idx = stream_li_idx - ctm.Int32(1) + pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + pv_done_phase = ( + pv_done_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pv_full_mbar.subview(pv_done_slot), pv_done_phase, time_limit=10000000 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + current_old_scale = _compute_smem_p4_lazy_rebase_scale(lane_rebase, anchor_delta_log2) + _rescale_smem_p4_resident_o_in_tmem(corr_o_acc_tmem, current_old_scale) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + elif prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + if is_compute_warp: + if cutlass.const_expr(not is_final_tile): + while not prims.mbarrier_try_wait_parity( + direct_p_ready_mbar, qk_phase, time_limit=10000000 + ): + pass + cute.arch.mbarrier_arrive( + p4_smem_slot0_mbar_ptr + p4_smem_slot0_producer.current_handle().index, + cutlass.Int32(0), + ) + p4_smem_slot0_producer.advance() + cute.arch.mbarrier_arrive( + tmem_score_slot0_mbar_ptr + + SMEM_P4_TMEM_SCORE_PIPELINE_STAGES + + tmem_score_slot0_consumer.current_handle().index, + cutlass.Int32(0), + ) + tmem_score_slot0_consumer.advance() + if is_compute_warp: + prims.barrier( + barrier_id=SMEM_P4_P4_SLOT0_BAR_ID, + number_of_threads=SMEM_P4_SCORE_SLOT_PRODUCER_THREADS, + ) + else: + prims.barrier( + barrier_id=SMEM_P4_P4_SLOT1_BAR_ID, + number_of_threads=SMEM_P4_CORRECTION_THREADS, + ) + row_sum = prequant_row_sum + rowsum_empty_mbar = pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + bank_idx + ) + if stream_li_idx >= ctm.Int32(SMEM_P4_TMEM_BANKS): + previous_rowsum_phase = ( + (stream_li_idx - ctm.Int32(SMEM_P4_TMEM_BANKS)) + // ctm.Int32(SMEM_P4_TMEM_BANKS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + rowsum_empty_mbar, previous_rowsum_phase, time_limit=10000000 + ): + pass + _store_smem_p4_n256_rowsum_part( + sScoreSlotRowMax, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + bank_idx=bank_idx, + score_half_idx=score_half_idx, + value=row_sum, + ) + if is_compute_warp & (producer_col_band == ctm.Int32(0)): + row_state_local_row = producer_row_band * ctm.Int32(32) + (tidx & ctm.Int32(31)) + sRowScaleRing[bank_idx * ctm.Int32(SMEM_P4_BMM1_M) + row_state_local_row] = ( + current_old_scale + ) + if prims.elect_sync(): + prims.mbarrier_arrive(rowsum_full_mbar) + if not is_compute_warp: + while not prims.mbarrier_try_wait_parity( + pair_rowmax_read_done_mbar, qk_phase, time_limit=10000000 + ): + pass + return ( + next_running_row_max, + next_running_row_anchor, + next_carried_dq1024, + next_carried_old_scale, + ) + + +@cute.jit +def _issue_pipelined_score_loads( + qk_full_mbar, stream_li_idx: ctm.Int32, warp_idx, base_col_id, base_row_id +): + """Wait for tile ``stream_li_idx``'s scores and issue their TMEM loads without waiting on them.""" + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + qk_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + score_tmem_addresses = _prepare_p4_n256_score_half_tmem_addresses( + warp_idx, + base_col_id, + base_row_id, + ctm.Int32(SMEM_P4_QK_ACC_COL_OFFSET) + + bank_idx * ctm.Int32(SMEM_P4_QK_PIPELINE_SLOT_STRIDE), + ) + while not prims.mbarrier_try_wait_parity( + qk_full_mbar.subview(bank_idx), qk_phase, time_limit=10000000 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + pending = [] + for local_group_idx in ctm.range_constexpr(4): + tmem = ctm.inttoptr(score_tmem_addresses[local_group_idx], 6, ctm.Float32) + pending.append(_tcgen05_ld_red_32x32b_x16_max_f32(tmem)) + return pending + + +@cute.jit +def _finish_pipelined_score_loads( + pending, + sPipelinedRowMax, + leader_pair_overlap_sync_mbars, + stream_li_idx: ctm.Int32, + warp_idx: ctm.Int32, + tidx: ctm.Int32, +): + """Wait for the score loads, publish the tile row max by tile parity and release the score bank.""" + producer_warp_base = (warp_idx >> ctm.Int32(2)) << ctm.Int32(2) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + score_values = [] + group_max_values = [] + tile_row_max = ctm.Float32(-ctm.Float32.inf) + for local_group_idx in ctm.range_constexpr(4): + regs = pending[local_group_idx] + for elem in ctm.range_constexpr(SF_VEC_SIZE): + score_values.append(regs[elem].bitcast(ctm.Float32)) + group_max = regs[SF_VEC_SIZE].bitcast(ctm.Float32) + group_max_values.append(group_max) + tile_row_max = cute.arch.fmax(tile_row_max, group_max) + parity = stream_li_idx & ctm.Int32(1) + _publish_p4_n256_atomic_rowmax( + sPipelinedRowMax.subview(parity * ctm.Int32(SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS)), + tile_row_max, + warp_idx, + tidx, + producer_warp_base, + ) + dsl_hints.cfence() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar( + leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + + stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + ) + ) + return ( + ctm.Vector.from_elements(tuple(score_values), dtype=ctm.Float32), + ctm.Vector.from_elements(tuple(group_max_values), dtype=ctm.Float32), + ) + + +@cute.jit +def _prepack_pipelined_tile( + score_values, + group_max_values, + warp_idx, + softmax_scale_log2, + stream_li_idx, + valid_k, + use_mixed_imlp: ctm.Constexpr, +): + """Pack an unmasked tile's FP4 P words and group sums ahead of its row-max exchange. + + Returns 16 scalars (8 P words, 4 group maxima, 4 group sums): the tile loop carries them to + the next iteration, and the DSL cannot rebind a list inside a dynamic loop. + """ + producer_warp_base = (warp_idx >> ctm.Int32(2)) << ctm.Int32(2) + score_half_idx = warp_idx >> ctm.Int32(2) + group_valid_values = ctm.Vector.from_elements( + tuple((ctm.Boolean(True) for _ in range(4))), dtype=ctm.Boolean + ) + packed, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( + score_values, + group_max_values, + group_valid_values, + warp_idx, + softmax_scale_log2, + kv_tile_idx=stream_li_idx, + valid_k=valid_k, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + mask_valid_k=False, + use_mixed_imlp=use_mixed_imlp, + ) + return ( + packed[0], + packed[1], + packed[2], + packed[3], + packed[4], + packed[5], + packed[6], + packed[7], + group_max_values[0], + group_max_values[1], + group_max_values[2], + group_max_values[3], + group_p_sums[0], + group_p_sums[1], + group_p_sums[2], + group_p_sums[3], + ) + + +@cute.jit +def _finalize_pipelined_tile( + sPipelinedRowMax, + sP, + sPSF, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + carried, + valid_k: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + running_row_max: ctm.Float32, + running_row_anchor: ctm.Float32, + running_part_sum: ctm.Float32, + stream_li_idx: ctm.Int32, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + """Finish tile ``stream_li_idx`` from its prepacked P words. + + Same steps as ``_runtime_t336_prepare_p_source`` from the dual-half barrier on followed by the + tail of ``_runtime_t336_producer_tile_v23``: row anchor, P and PSF stores, lazy O rescale and + the P_SOURCE_READY arrive. + """ + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + p_stage_base = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + score_half_idx = warp_idx >> ctm.Int32(2) + producer_warp_base = score_half_idx << ctm.Int32(2) + is_compute_warp = score_half_idx == ctm.Int32(0) + has_previous_tile = stream_li_idx > ctm.Int32(0) + packed_p4_words = ctm.Vector.from_elements(tuple(carried[0:8]), dtype=ctm.Int32) + group_max_values = ctm.Vector.from_elements(tuple(carried[8:12]), dtype=ctm.Float32) + group_p_sums = ctm.Vector.from_elements(tuple(carried[12:16]), dtype=ctm.Float32) + pair_rowmax_read_done_mbar = pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_mbar = leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_dsmem_mbar = leader_pair_overlap_sync_mbars_dsmem.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + if is_compute_warp: + p4_smem_slot0_producer.acquire() + prims.barrier( + barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, + number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, + ) + parity = stream_li_idx & ctm.Int32(1) + new_row_max = cute.arch.fmax( + running_row_max, + _load_p4_n256_atomic_running_rowmax( + sPipelinedRowMax.subview(parity * ctm.Int32(SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS)), + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + ), + ) + next_running_row_anchor, lane_rebase, anchor_delta_log2 = _select_smem_p4_lazy_anchor_runtime( + new_row_max, + running_row_anchor, + softmax_scale_log2, + has_previous_tile, + ) + if is_compute_warp: + dsl_hints.cfence() + if prims.elect_sync(): + prims.mbarrier_arrive(pair_rowmax_read_done_mbar) + _store_prepacked_p4_n256_score_half_to_direct_a( + sP, + packed_p4_words, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + p_stage_base=p_stage_base, + score_half_idx=score_half_idx, + ) + if is_compute_warp: + dsl_hints.cfence() + owned_psf_word, row_sum = _finalize_p4_n256_score_half_psf_source( + sPSF, + group_max_values, + group_p_sums, + next_running_row_anchor, + warp_idx, + tidx, + softmax_scale_log2, + pv_psf_rescale, + kv_tile_idx=stream_li_idx, + valid_k=valid_k, + mask_valid_k=False, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + bank_idx=bank_idx, + use_mixed_imlp=use_mixed_imlp, + ) + current_old_scale = ctm.Float32(1.0) + warp_rebase = cute.arch.vote_any_sync(lane_rebase) + if has_previous_tile: + common_source_ready_pred = prims.elect_sync() & ~warp_rebase + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) + if warp_rebase: + pv_done_li_idx = stream_li_idx - ctm.Int32(1) + pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + pv_done_phase = ( + pv_done_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pv_full_mbar.subview(pv_done_slot), pv_done_phase, time_limit=10000000 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + current_old_scale = _compute_smem_p4_lazy_rebase_scale(lane_rebase, anchor_delta_log2) + _rescale_smem_p4_resident_o_in_tmem(corr_o_acc_tmem, current_old_scale) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + elif prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + if is_compute_warp: + p4_smem_slot0_producer.advance() + next_running_part_sum = running_part_sum * current_old_scale + row_sum + return (new_row_max, next_running_row_anchor, next_running_part_sum) + + +@cute.jit +def _run_pipelined_leading_tiles( + sAtomicRunningRowMax, + sPipelinedRowMax, + sP, + sPSF, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + valid_k: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + pipelined_tiles: ctm.Int32, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + """Run producer tiles [0, pipelined_tiles), which need no causal mask, pipelined by one tile. + + Tile t + 1's score loads are in flight while tile t is finished, so the TMEM load latency + and the row-max exchange of one tile overlap the P stores of the previous one. + """ + pending = _issue_pipelined_score_loads( + qk_full_mbar, ctm.Int32(0), warp_idx, base_col_id, base_row_id + ) + scores, group_max = _finish_pipelined_score_loads( + pending, sPipelinedRowMax, leader_pair_overlap_sync_mbars, ctm.Int32(0), warp_idx, tidx + ) + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15) = ( + _prepack_pipelined_tile( + scores, group_max, warp_idx, softmax_scale_log2, ctm.Int32(0), valid_k, use_mixed_imlp + ) + ) + running_row_max = ctm.Float32(-ctm.Float32.inf) + running_row_anchor = ctm.Float32(-ctm.Float32.inf) + running_part_sum = ctm.Float32(0.0) + for stream_li_idx in cutlass.range(0, pipelined_tiles - ctm.Int32(1), 1, unroll=1): + next_li_idx = stream_li_idx + ctm.Int32(1) + next_pending = _issue_pipelined_score_loads( + qk_full_mbar, next_li_idx, warp_idx, base_col_id, base_row_id + ) + running_row_max, running_row_anchor, running_part_sum = _finalize_pipelined_tile( + sPipelinedRowMax, + sP, + sPSF, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15), + valid_k, + softmax_scale_log2, + pv_psf_rescale, + running_row_max, + running_row_anchor, + running_part_sum, + stream_li_idx, + use_mixed_imlp, + ) + next_scores, next_group_max = _finish_pipelined_score_loads( + next_pending, + sPipelinedRowMax, + leader_pair_overlap_sync_mbars, + next_li_idx, + warp_idx, + tidx, + ) + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15) = ( + _prepack_pipelined_tile( + next_scores, + next_group_max, + warp_idx, + softmax_scale_log2, + next_li_idx, + valid_k, + use_mixed_imlp, + ) + ) + running_row_max, running_row_anchor, running_part_sum = _finalize_pipelined_tile( + sPipelinedRowMax, + sP, + sPSF, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15), + valid_k, + softmax_scale_log2, + pv_psf_rescale, + running_row_max, + running_row_anchor, + running_part_sum, + pipelined_tiles - ctm.Int32(1), + use_mixed_imlp, + ) + # The per-tile loop that follows reads the running row max from sAtomicRunningRowMax. + _publish_p4_n256_atomic_rowmax( + sAtomicRunningRowMax, + running_row_max, + warp_idx, + tidx, + (warp_idx >> ctm.Int32(2)) << ctm.Int32(2), + ) + return (running_row_max, running_row_anchor, running_part_sum) + + +@cute.jit +def _runtime_t336_producer_tile_v23( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + p4_tmem_base_col_id: ctm.Int32, + p4_tmem_base_row_id: ctm.Int32, + valid_k_for_l: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + running_row_max: ctm.Float32, + running_row_anchor: ctm.Float32, + carried_dq1024: ctm.Int32, + carried_old_scale: ctm.Float32, + running_part_sum: ctm.Float32, + stream_li_idx: ctm.Int32, + has_previous_tile: ctm.Boolean, + tail_needs_mask: ctm.Boolean, + is_final_tile: ctm.Constexpr, + use_mixed_imlp: ctm.Constexpr, + causal_mask_possible: ctm.Constexpr = False, +) -> tuple: + """Run one producer tile with a scalar-only runtime final-tail dispatch.""" + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + qk_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + pair_rowmax_read_done_mbar = pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_mbar = leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_dsmem_mbar = leader_pair_overlap_sync_mbars_dsmem.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + leader_score_consumed_mbar = leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + bank_idx + ) + score_half_idx = warp_idx >> ctm.Int32(2) + is_compute_warp = score_half_idx == ctm.Int32(0) + + next_running_row_max = running_row_max + next_running_row_anchor = running_row_anchor + next_carried_dq1024 = carried_dq1024 + next_carried_old_scale = carried_old_scale + lane_rebase = ctm.Boolean(False) + anchor_delta_log2 = ctm.Float32(0.0) + row_sum = ctm.Float32(0.0) + if cutlass.const_expr(is_final_tile or causal_mask_possible): + if tail_needs_mask: + ( + next_running_row_max, + next_running_row_anchor, + lane_rebase, + anchor_delta_log2, + row_sum, + ) = _runtime_t336_prepare_p_source( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_rowmax_read_done_mbar, + leader_score_consumed_mbar, + p4_smem_slot0_producer, + warp_idx, + tidx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + valid_k_for_l, + softmax_scale_log2, + pv_psf_rescale, + running_row_anchor, + stream_li_idx, + has_previous_tile, + True, + use_mixed_imlp, + ) + else: + ( + next_running_row_max, + next_running_row_anchor, + lane_rebase, + anchor_delta_log2, + row_sum, + ) = _runtime_t336_prepare_p_source( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_rowmax_read_done_mbar, + leader_score_consumed_mbar, + p4_smem_slot0_producer, + warp_idx, + tidx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + valid_k_for_l, + softmax_scale_log2, + pv_psf_rescale, + running_row_anchor, + stream_li_idx, + has_previous_tile, + False, + use_mixed_imlp, + ) + else: + ( + next_running_row_max, + next_running_row_anchor, + lane_rebase, + anchor_delta_log2, + row_sum, + ) = _runtime_t336_prepare_p_source( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_rowmax_read_done_mbar, + leader_score_consumed_mbar, + p4_smem_slot0_producer, + warp_idx, + tidx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + valid_k_for_l, + softmax_scale_log2, + pv_psf_rescale, + running_row_anchor, + stream_li_idx, + has_previous_tile, + False, + use_mixed_imlp, + ) + + current_old_scale = ctm.Float32(1.0) + warp_rebase = cute.arch.vote_any_sync(lane_rebase) + if has_previous_tile: + common_source_ready_pred = prims.elect_sync() & ~warp_rebase + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) + if warp_rebase: + pv_done_li_idx = stream_li_idx - ctm.Int32(1) + pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + pv_done_phase = ( + pv_done_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pv_full_mbar.subview(pv_done_slot), pv_done_phase, time_limit=10000000 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + current_old_scale = _compute_smem_p4_lazy_rebase_scale(lane_rebase, anchor_delta_log2) + _rescale_smem_p4_resident_o_in_tmem(corr_o_acc_tmem, current_old_scale) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + elif prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + # P_SOURCE_READY is the only P-side signal the MMA warp needs: it implies + # this tile's P/PSF are in smem and its score bank has been read. + if is_compute_warp: + p4_smem_slot0_producer.advance() + # Each warp keeps the denominator of its own 64-column slice in a register, + # rebased with the row anchor; the slices are combined once after the loop. + next_running_part_sum = running_part_sum * current_old_scale + row_sum + if not is_compute_warp: + while not prims.mbarrier_try_wait_parity( + pair_rowmax_read_done_mbar, qk_phase, time_limit=10000000 + ): + pass + return ( + next_running_row_max, + next_running_row_anchor, + next_carried_dq1024, + next_carried_old_scale, + next_running_part_sum, + ) + + +@cute.jit +def _stage_smem_p4_runtime_scale_pair( + mQGlobalScale: cute.Tensor, + mKvGlobalScale: cute.Tensor, + sRuntimeScalePair, + softmax_scale_log2: ctm.Float32, +) -> None: + q_global_scale = mQGlobalScale[0] + kv_global_scale = mKvGlobalScale[0] + runtime_scale_pair = ctm.Vector.from_elements( + ( + softmax_scale_log2 / (q_global_scale * kv_global_scale), + ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) / kv_global_scale, + ), + dtype=ctm.Float32, + ) + if prims.elect_sync(): + sRuntimeScalePair.data_ptr().store(runtime_scale_pair, alignment=8) + + +@ctm.dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [], + "griddepcontrol.launch_dependents;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@ctm.dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [], + "griddepcontrol.wait;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _stage_smem_p4_page_id_plan( + mPageTable_pl: cute.Tensor, + sPageIdPlan, + tidx: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, + planned_page_count: ctm.Int32, +) -> None: + logical_page_idx = tidx + while logical_page_idx < planned_page_count: + physical_page = _lookup_physical_page( + mPageTable_pl, logical_page_idx, ctm.Int32(0), page_begin, page_count + ) + (sPageIdPlan.data_ptr() + logical_page_idx).store(physical_page) + logical_page_idx += ctm.Int32(THREADS_PER_CTA) + + +@cute.jit +def _packed_query_bound( + mHelixKvBounds_l: cute.Tensor, + request_idx: ctm.Int32, + query_idx: ctm.Int32, + request_valid_k: ctm.Int32, + query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, +) -> ctm.Int32: + """Visible KV count of one query token of a packed request. + + Helix supplies a per-token rank-local bound; otherwise the linear MTP + window ends ``query_len - 1 - query_idx`` tokens before the request's KV end. + """ + bound = ctm.Int32(0) + if cutlass.const_expr(use_helix_kv_bounds): + bound = ctm.Int32(mHelixKvBounds_l[request_idx * ctm.Int32(query_len_per_seq) + query_idx]) + else: + bound = ctm.max( + request_valid_k - (ctm.Int32(query_len_per_seq - 1) - query_idx), ctm.Int32(0) + ) + return bound + + +@cute.jit +def _packed_request_tiles( + mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, + request_idx: ctm.Int32, + query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, +) -> ctm.Int32: + """KV tiles a packed request spans (its last query token's bound), at least one.""" + bound = _packed_query_bound( + mHelixKvBounds_l, + request_idx, + ctm.Int32(query_len_per_seq - 1), + ctm.Int32(mValidK_l[request_idx]), + query_len_per_seq, + use_helix_kv_bounds, + ) + return ctm.max((bound + ctm.Int32(KV_TILE - 1)) // ctm.Int32(KV_TILE), ctm.Int32(1)) + + +@cute.jit +def _run_mla_decode_body( + mPageTable_pl: cute.Tensor, + mPageIndptr_s: cute.Tensor, + mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, + mPlan_i: cute.Tensor, + mQGlobalScale: cute.Tensor, + mKvGlobalScale: cute.Tensor, + tma_q_ptr, + tma_k_page_ptr, + tma_q_tail_ptr, + tma_k_tail_page_ptr, + tma_v_tile_ptr, + tma_qsf_ptr, + tma_qsf_tail_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + tma_vsf_ptr, + tma_vsf_stage_ptr, + tma_v_pair_ptr, + tma_sfb_ptr, + mC_mnl: cute.Tensor, + mAccum_mnl: cute.Tensor, + mRowMax_ml: cute.Tensor, + mRowSum_ml: cute.Tensor, + problem_size: tuple, + runtime_m: ctm.Int32, + num_requests: ctm.Int32, + num_splits: ctm.Int32, + valid_m: ctm.Int32, + c_batch_stride: ctm.Int32, + v_page_offset: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_output_scale: ctm.Float32, + page_size: ctm.Constexpr, + query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, + use_smem_page_plan: ctm.Constexpr, + use_mixed_imlp: ctm.Constexpr = False, + use_consecutive_page_pair: ctm.Constexpr = False, + use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, + store_split: ctm.Constexpr = False, + folded_query_heads: ctm.Constexpr = 0, + packed_query_heads: ctm.Constexpr = 0, + pdl_wait: ctm.Constexpr = False, + launch_dependents: ctm.Constexpr = False, +) -> None: + pv_psf_rescale = ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) * pv_output_scale + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + n, k = problem_size + m = runtime_m + page_batch = ctm.Int32(0) + valid_k_for_l = ctm.Int32(0) + q_batch = bidz + # Packed rows: grid.z is the request when num_splits == 1, else a slot of the plan's + # tables, in grid order (packed_plan_kernel). Rows of this M tile are [query token, + # head] pairs of one request; the split owns KV tiles [split_tile_begin, + # split_tile_begin + split_tiles) and every bound below is relative to split_base_k. + split_tile_begin = ctm.Int32(0) + split_base_k = ctm.Int32(0) + first_valid_k = ctm.Int32(0) + packed_tile_row0 = ctm.Int32(0) + packed_request_valid_k = ctm.Int32(0) + packed_single_task = ctm.Boolean(False) + c_batch_idx = bidz + # Partials and row stats go to begins[request] + split, where the reduction reads them. + store_task_idx = bidz + if cutlass.const_expr(packed_query_heads > 0): + # Balanced plan (one task per request when num_splits == 1, else the + # plan kernel's tables): the request and KV split this task owns. + tiles_per_split = ctm.Int32(1 << 14) + task_request = bidz + task_split = ctm.Int32(0) + task_count = ctm.Int32(1) + if num_splits > ctm.Int32(1): + task_capacity = num_requests * num_splits + plan_tables = ctm.Int32(3) + num_requests + tiles_per_split = ctm.Int32(mPlan_i[0]) + task_request = ctm.Int32(mPlan_i[plan_tables + bidz]) + task_split = ctm.Int32(mPlan_i[plan_tables + task_capacity + bidz]) + task_count = ctm.Int32(mPlan_i[plan_tables + task_capacity + task_capacity + bidz]) + store_task_idx = ctm.Int32(mPlan_i[ctm.Int32(2) + task_request]) + task_split + store_task_idx = cute.arch.make_warp_uniform(store_task_idx) + tiles_per_split = cute.arch.make_warp_uniform(tiles_per_split) + page_batch = cute.arch.make_warp_uniform(task_request) + split_idx = cute.arch.make_warp_uniform(task_split) + # A request that owns a single task writes its output directly even in + # split mode; the reduction kernel then only derives its Helix stats. + packed_single_task = cute.arch.make_warp_uniform(task_count) == ctm.Int32(1) + q_batch = page_batch + c_batch_idx = page_batch + packed_tile_row0 = (bidx // ctm.Int32(CLUSTER_SHAPE_MNK[0])) * ctm.Int32( + SMEM_P4_CTA_GROUP_M + ) + packed_request_valid_k = cute.arch.make_warp_uniform(ctm.Int32(mValidK_l[page_batch])) + packed_last_query = ctm.Int32(query_len_per_seq - 1) + tile_first_query = ctm.min( + packed_tile_row0 // ctm.Int32(packed_query_heads), packed_last_query + ) + tile_last_query = ctm.min( + (packed_tile_row0 + ctm.Int32(SMEM_P4_CTA_GROUP_M - 1)) + // ctm.Int32(packed_query_heads), + packed_last_query, + ) + tile_max_bound = _packed_query_bound( + mHelixKvBounds_l, + page_batch, + tile_last_query, + packed_request_valid_k, + query_len_per_seq, + use_helix_kv_bounds, + ) + tile_min_bound = _packed_query_bound( + mHelixKvBounds_l, + page_batch, + tile_first_query, + packed_request_valid_k, + query_len_per_seq, + use_helix_kv_bounds, + ) + k_tile_total = (tile_max_bound + ctm.Int32(KV_TILE - 1)) // ctm.Int32(KV_TILE) + split_tile_begin = ctm.min(split_idx * tiles_per_split, k_tile_total) + split_tile_end = ctm.min(k_tile_total, split_tile_begin + tiles_per_split) + split_tiles = ctm.max(split_tile_end - split_tile_begin, ctm.Int32(0)) + split_base_k = split_tile_begin * ctm.Int32(KV_TILE) + valid_k_for_l = ctm.min( + ctm.max(tile_max_bound - split_base_k, ctm.Int32(0)), + split_tiles * ctm.Int32(KV_TILE), + ) + first_valid_k = ctm.min(ctm.max(tile_min_bound - split_base_k, ctm.Int32(0)), valid_k_for_l) + else: + page_batch = cute.arch.make_warp_uniform(bidz // ctm.Int32(query_len_per_seq)) + if cutlass.const_expr(folded_query_heads > 0): + valid_k_for_l = mValidK_l[page_batch] + elif cutlass.const_expr(use_helix_kv_bounds): + valid_k_for_l = mHelixKvBounds_l[bidz] + else: + query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) + valid_k_for_l = ctm.max( + mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), + ctm.Int32(0), + ) + valid_k_for_l = cute.arch.make_warp_uniform(valid_k_for_l) + first_valid_k = cute.arch.make_warp_uniform(first_valid_k) + folded_first_valid_k = ctm.Int32(0) + if cutlass.const_expr(folded_query_heads > 0): + folded_first_valid_k = cute.arch.make_warp_uniform(mHelixKvBounds_l[bidz]) + csr_page_begin = ctm.Int32(0) + csr_page_count = ctm.Int32(0) + csr_page_begin = cute.arch.make_warp_uniform(ctm.Int32(mPageIndptr_s[page_batch])) + csr_page_count = cute.arch.make_warp_uniform( + ctm.Int32(mPageIndptr_s[page_batch + ctm.Int32(1)]) - csr_page_begin + ) + if cutlass.const_expr(packed_query_heads > 0): + # The split's page window starts at its first KV tile; lookups past + # the request's last page clamp inside _lookup_physical_page. + split_page_begin = split_tile_begin * ctm.Int32(SMEM_P4_PAGES_PER_KV_TILE) + csr_page_begin = csr_page_begin + split_page_begin + csr_page_count = ctm.max(csr_page_count - split_page_begin, ctm.Int32(0)) + raw_kv_tiles = (valid_k_for_l + ctm.Int32(KV_TILE - 1)) // ctm.Int32(KV_TILE) + kv_tiles = ctm.max(raw_kv_tiles, ctm.Int32(1)) + planned_page_count = kv_tiles * ctm.Int32(SMEM_P4_PAGES_PER_KV_TILE) + if cutlass.const_expr(page_size != TRTLLM_PAGE_SIZE): + raise ValueError( + f"FP4 MLA decode requires TensorRT-LLM page_size={TRTLLM_PAGE_SIZE}, got {page_size}" + ) + physical_sfb_blocks_per_l = k // ctm.Int32(SMEM_P4_SFB_ATOM_ROWS) + num_tmem_alloc_cols: ctm.Constexpr = NUM_TMEM_ALLOC_COLS + q_tma_bytes: ctm.Constexpr = SMEM_P4_Q_ONLY_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + qk_tma_bytes: ctm.Constexpr = SMEM_P4_QK_KONLY_TILE_BYTES * CLUSTER_SHAPE_MNK[0] + v_tma_bytes: ctm.Constexpr = SMEM_P4_V_TILE_TMA_BYTES * CLUSTER_SHAPE_MNK[0] + cta_rank = cute.arch.block_idx_in_cluster() + is_leader_cta = cta_rank == 0 + qk_full_mbar = ctm.Array(ctm.Int64, SMEM_P4_QK_COMPLETION_MBARS, space=ctm.AddressSpace.smem) + pv_full_mbar = ctm.Array( + ctm.Int64, SMEM_P4_QK_PIPELINE_SLOTS, space=ctm.AddressSpace.smem, alignment=8 + ) + tmem_dealloc_mbar = ctm.Array(ctm.Int64, 1, space=ctm.AddressSpace.smem) + tmem_holding_buf = ctm.Array(ctm.Int32, 1, space=ctm.AddressSpace.smem) + q_smem_mbars = ctm.Array( + ctm.Int64, + NUM_QK_Q_PIPELINE_STAGE * 2, + space=ctm.AddressSpace.smem, + alignment=8, + ) + qk_smem_mbars = ctm.Array( + ctm.Int64, + SMEM_P4_QK_BARRIER_STAGES * 2, + space=ctm.AddressSpace.smem, + alignment=8, + ) + v_smem_mbars = ctm.Array( + ctm.Int64, + SMEM_P4_V_PIPELINE_STAGES * 2, + space=ctm.AddressSpace.smem, + alignment=8, + ) + p4_smem_slot0_mbars = ctm.Array( + ctm.Int64, + SMEM_P4_P4_CONTROL_MBAR_STAGES * 2, + space=ctm.AddressSpace.smem, + alignment=8, + ) + pair_overlap_sync_mbars = ctm.Array( + ctm.Int64, + SMEM_P4_PAIR_OVERLAP_MBAR_COUNT, + space=ctm.AddressSpace.smem, + alignment=8, + ) + qsf_owner_ready_mbar = ctm.Array(ctm.Int64, 1, space=ctm.AddressSpace.smem, alignment=8) + qsf_owner_done_mbar = ctm.Array(ctm.Int64, 1, space=ctm.AddressSpace.smem, alignment=8) + qk_bulk_ready_mbars = ctm.Array( + ctm.Int64, + SMEM_P4_QK_BULK_READY_MBARS, + space=ctm.AddressSpace.smem, + alignment=8, + ) + sA = ctm.Array( + ctm.Int32, + SMEM_P4_Q_COMPACT_TILE_BYTES // 4, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sB = ctm.Array( + ctm.Int32, + SMEM_P4_K_COMPACT_RING_BYTES // 4, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sSFA = ctm.Array( + ctm.Uint8, + SMEM_P4_QSF_COMPACT_TILE_BYTES, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sSFB = ctm.Array( + ctm.Uint8, + SMEM_P4_KSF_COMPACT_RING_BYTES, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sP = ctm.Array(ctm.Int32, SMEM_P4_P4_SMEM_WORDS, space=ctm.AddressSpace.smem, alignment=128) + sV = ctm.Array( + ctm.Int32, + SMEM_P4_V_PIPELINE_STAGES * SMEM_P4_V_DATA_STAGE_BYTES // 4, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sPSF = ctm.Array( + ctm.Int32, + SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sVSF = ctm.Array( + ctm.Uint8, + SMEM_P4_V_PIPELINE_STAGES * SMEM_P4_V_SFB_TMA_STAGE_BYTES, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sScoreSlotRowMax = ctm.Array( + ctm.Float32, + SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sAtomicRunningRowMax = ctm.Array( + ctm.Uint32, + SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS, + space=ctm.AddressSpace.smem, + alignment=128, + ) + # Pipelined producer tiles publish their row maxima by tile parity. + sPipelinedRowMax = ctm.Array( + ctm.Uint32, + 2 * SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sPageIdPlan = ctm.Array( + ctm.Int32, + SMEM_P4_PAGE_ID_PLAN_INTS, + space=ctm.AddressSpace.smem, + alignment=128, + ) + sRuntimeScalePair = ctm.Array( + ctm.Float32, + SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS, + space=ctm.AddressSpace.smem, + alignment=8, + ) + cta_layout_vmnk = cute.make_layout((CLUSTER_SHAPE_MNK[0], 1, 1, 1)) + tma_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + umma_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + p4_stage_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + SMEM_P4_SCORE_SLOT_PRODUCER_THREADS * CLUSTER_SHAPE_MNK[0], + ) + q_smem_mbar_ptr = cute.make_ptr( + ctm.Int64, q_smem_mbars.data_ptr().toint(), cute.AddressSpace.smem + ) + qk_smem_mbar_ptr = cute.make_ptr( + ctm.Int64, qk_smem_mbars.data_ptr().toint(), cute.AddressSpace.smem + ) + v_smem_mbar_ptr = cute.make_ptr( + ctm.Int64, v_smem_mbars.data_ptr().toint(), cute.AddressSpace.smem + ) + p4_smem_slot0_mbar_ptr = cute.make_ptr( + ctm.Int64, p4_smem_slot0_mbars.data_ptr().toint(), cute.AddressSpace.smem + ) + qk_bulk_ready_mbar_ptr = cute.make_ptr( + ctm.Int64, qk_bulk_ready_mbars.data_ptr().toint(), cute.AddressSpace.smem + ) + q_smem_producer, q_smem_consumer = pipeline.PipelineTmaUmma.create( + num_stages=NUM_QK_Q_PIPELINE_STAGE, + producer_group=tma_group, + consumer_group=umma_group, + tx_count=q_tma_bytes, + barrier_storage=q_smem_mbar_ptr, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + name="fp4_q_smem", + ).make_participants() + qk_smem_producer, qk_smem_consumer = pipeline.PipelineTmaUmma.create( + num_stages=SMEM_P4_QK_BARRIER_STAGES, + producer_group=tma_group, + consumer_group=umma_group, + tx_count=qk_tma_bytes, + barrier_storage=qk_smem_mbar_ptr, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + name="fp4_qk_smem", + ).make_participants() + v_smem_producer, v_smem_consumer = pipeline.PipelineTmaUmma.create( + num_stages=SMEM_P4_V_PIPELINE_STAGES, + producer_group=tma_group, + consumer_group=umma_group, + tx_count=v_tma_bytes, + barrier_storage=v_smem_mbar_ptr, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + name="fp4_v_smem", + ).make_participants() + p4_smem_slot0_producer, _ = pipeline.PipelineAsyncUmma.create( + num_stages=SMEM_P4_P4_PIPELINE_STAGES, + producer_group=p4_stage_group, + consumer_group=umma_group, + barrier_storage=p4_smem_slot0_mbar_ptr, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + name="fp4_p4_smem_slot0", + ).make_participants() + if warp_idx == 0: + if prims.elect_sync(): + for qk_completion_mbar_idx in ctm.range_constexpr(SMEM_P4_QK_COMPLETION_MBARS): + prims.mbarrier_init(qk_full_mbar.subview(qk_completion_mbar_idx), 1) + for pv_full_mbar_idx in ctm.range_constexpr(SMEM_P4_QK_PIPELINE_SLOTS): + prims.mbarrier_init(pv_full_mbar.subview(pv_full_mbar_idx), 1) + for pair_slot_idx in ctm.range_constexpr(SMEM_P4_QK_PIPELINE_SLOTS): + prims.mbarrier_init( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_P4_READY_MBAR_OFFSET + pair_slot_idx + ), + SMEM_P4_TOTAL_PRODUCER_WARPS, + ) + prims.mbarrier_init( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + pair_slot_idx + ), + 1, + ) + prims.mbarrier_init( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_FINAL_PV_ISSUED_MBAR_OFFSET + pair_slot_idx + ), + 1, + ) + prims.mbarrier_init( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + pair_slot_idx + ), + SMEM_P4_COMPUTE_WARPS, + ) + prims.mbarrier_init( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + pair_slot_idx + ), + SMEM_P4_TOTAL_PRODUCER_WARPS * CLUSTER_SHAPE_MNK[0], + ) + prims.mbarrier_init( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + pair_slot_idx + ), + SMEM_P4_TOTAL_PRODUCER_WARPS * CLUSTER_SHAPE_MNK[0], + ) + prims.mbarrier_init(qsf_owner_ready_mbar, 1) + prims.mbarrier_init(qsf_owner_done_mbar, CLUSTER_SHAPE_MNK[0]) + for qk_bulk_slot_idx in ctm.range_constexpr(SMEM_P4_QK_BULK_READY_MBARS): + prims.mbarrier_init(qk_bulk_ready_mbars.subview(qk_bulk_slot_idx), 1) + prims.mbarrier_init(pair_overlap_sync_mbars.subview(SMEM_P4_QK15_STAGE1_MBAR_OFFSET), 1) + prims.mbarrier_init(tmem_dealloc_mbar, 32) + prims.fence_mbarrier_init() + if warp_idx == ctm.Int32(TMA_QK_WARP_ID): + _stage_smem_p4_runtime_scale_pair( + mQGlobalScale, mKvGlobalScale, sRuntimeScalePair, softmax_scale_log2 + ) + if prims.elect_sync(): + # The per-tile mask decision reloads this bound from smem instead + # of holding it in a register for the whole KV loop (see below). + (sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD)).store( + _bitcast_i32_to_f32(first_valid_k) + ) + if cutlass.const_expr(use_smem_page_plan): + _stage_smem_p4_page_id_plan( + mPageTable_pl, sPageIdPlan, tidx, csr_page_begin, csr_page_count, planned_page_count + ) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + runtime_scale_pair = sRuntimeScalePair.data_ptr().load(count=2, alignment=8) + # Both scales are CTA-wide constants used in every KV tile. The producer + # warps run at their register ceiling, so without the uniform hint the + # allocator spills this pair and reloads it from local memory per tile. + softmax_scale_log2 = _warp_uniform_f32(runtime_scale_pair[0]) + pv_psf_rescale = _warp_uniform_f32(runtime_scale_pair[1]) + # Both exact and mixed-IMLP modes feed the amplified PV-only PSF to MMA. + # Keep the denominator in its original coordinate and cancel the PV + # amplification exactly once at the final output. + pv_output_scale = ctm.Float32(1.0 / FP4_MLA_P_GLOBAL_SCALE) + if warp_idx >= ctm.Int32(MMA_WARP_ID): + prims.setmaxregister(SMEM_P4_AUX_REGISTER_BUDGET, prims.SetMaxRegisterAction.DECREASE) + if warp_idx == MMA_WARP_ID: + prims.tcgen05_alloc( + tmem_holding_buf, + num_tmem_alloc_cols, + is_exclusive=SMEM_P4_TMEM_ALLOC_EXCLUSIVE, + group=prims.CTAGroup.CTA_2, + ) + prims.barrier(barrier_id=TMEM_BAR_ID, number_of_threads=TMEM_BAR_THREADS) + if warp_idx == TMA_QK_WARP_ID: + if cutlass.const_expr(not pdl_wait): + for q_stage_idx in ctm.range_constexpr(NUM_QK_Q_STAGE): + q_stage_kblock_idx: ctm.Constexpr = q_stage_idx * SMEM_P4_QK_DATA_STAGE_KBLOCKS + if cutlass.const_expr(q_stage_kblock_idx >= SMEM_P4_QK_TAIL_STAGE_KBLOCK_START): + q_handle = q_smem_producer.acquire_and_advance( + expected_tx=SMEM_P4_Q_ONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + ) + else: + q_handle = q_smem_producer.acquire_and_advance() + q_tma_mbar = ctm.Array(q_handle.barrier, shape=1) + _load_qk_qonly_kblock_stage( + tma_q_ptr, + tma_q_tail_ptr, + tma_qsf_ptr, + tma_qsf_tail_ptr, + q_tma_mbar, + sA, + sSFA, + tidx, + bidx, + q_batch, + cta_rank, + qk_kblock_idx=q_stage_kblock_idx, + stage=q_stage_idx, + q_tma_phase=q_stage_idx, + manage_mbarrier=False, + ) + qk0_handle = qk_smem_producer.acquire_and_advance() + qk0_tma_mbar = ctm.Array(qk0_handle.barrier, shape=1) + if is_leader_cta: + if cutlass.const_expr(use_smem_page_plan): + qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(0) + ) + else: + qk0_page0, qk0_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, ctm.Int32(0), csr_page_begin, csr_page_count + ) + _load_runtime_t336_qk_tile( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk0_tma_mbar, + qk_bulk_ready_mbar_ptr, + pair_overlap_sync_mbars, + sB, + sSFB, + tidx, + bidz, + physical_sfb_blocks_per_l, + csr_page_begin, + csr_page_count, + qk0_page0, + qk0_page1, + ctm.Int32(0), + ctm.Int32(0), + page_size, + initial_tile=True, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + if cutlass.const_expr(pdl_wait): + # Programmatic launch: the loads above read the plan, the page table and the + # first KV tile, which no grid still in flight writes + # (run_trtllm_fp4_mla_decode_packed's pdl_wait contract); the query comes from + # the previous grid and goes ahead of the next KV tiles. + _griddepcontrol_wait() + for q_stage_idx in ctm.range_constexpr(NUM_QK_Q_STAGE): + q_stage_kblock_idx: ctm.Constexpr = q_stage_idx * SMEM_P4_QK_DATA_STAGE_KBLOCKS + if cutlass.const_expr(q_stage_kblock_idx >= SMEM_P4_QK_TAIL_STAGE_KBLOCK_START): + q_handle = q_smem_producer.acquire_and_advance( + expected_tx=SMEM_P4_Q_ONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + ) + else: + q_handle = q_smem_producer.acquire_and_advance() + q_tma_mbar = ctm.Array(q_handle.barrier, shape=1) + _load_qk_qonly_kblock_stage( + tma_q_ptr, + tma_q_tail_ptr, + tma_qsf_ptr, + tma_qsf_tail_ptr, + q_tma_mbar, + sA, + sSFA, + tidx, + bidx, + q_batch, + cta_rank, + qk_kblock_idx=q_stage_kblock_idx, + stage=q_stage_idx, + q_tma_phase=q_stage_idx, + manage_mbarrier=False, + ) + qk_prefix_end = min(kv_tiles, ctm.Int32(3)) + for qk_prefix_li_idx in cutlass.range(1, qk_prefix_end, 1, unroll=1): + qk_prefix_handle = qk_smem_producer.acquire_and_advance( + expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + ) + qk_prefix_tma_mbar = ctm.Array(qk_prefix_handle.barrier, shape=1) + if is_leader_cta: + if cutlass.const_expr(use_smem_page_plan): + qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) + else: + qk_prefix_page0, qk_prefix_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count + ) + _load_runtime_t336_qk_tile( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_prefix_tma_mbar, + qk_bulk_ready_mbar_ptr, + pair_overlap_sync_mbars, + sB, + sSFB, + tidx, + bidz, + physical_sfb_blocks_per_l, + csr_page_begin, + csr_page_count, + qk_prefix_page0, + qk_prefix_page1, + qk_prefix_li_idx, + qk_prefix_li_idx, + page_size, + split_prefix_owner=True, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + qk_steady_end = kv_tiles + qk15_final_enabled = kv_tiles == ctm.Int32(16) + if qk15_final_enabled: + qk_steady_end = ctm.Int32(15) + for kv_tile_idx in cutlass.range(3, qk_steady_end, 1, unroll=1): + qk_handle = qk_smem_producer.acquire_and_advance( + _sleep_until_producer_slot_free(qk_smem_producer), + expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0], + ) + qk_tma_mbar = ctm.Array(qk_handle.barrier, shape=1) + if is_leader_cta: + if cutlass.const_expr(use_smem_page_plan): + qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) + else: + qk_steady_page0, qk_steady_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count + ) + _load_runtime_t336_qk_tile( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_tma_mbar, + qk_bulk_ready_mbar_ptr, + pair_overlap_sync_mbars, + sB, + sSFB, + tidx, + bidz, + physical_sfb_blocks_per_l, + csr_page_begin, + csr_page_count, + qk_steady_page0, + qk_steady_page1, + kv_tile_idx, + kv_tile_idx % ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS), + page_size, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + else: + if cutlass.const_expr(use_smem_page_plan): + qk_rank1_page0, qk_rank1_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) + else: + qk_rank1_page0, qk_rank1_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count + ) + _load_runtime_t336_qk_tile_rank1( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_handle.barrier, + qk_bulk_ready_mbar_ptr, + sB, + sSFB, + qk_rank1_page0, + qk_rank1_page1, + kv_tile_idx, + kv_tile_idx % ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS), + page_size, + ) + if qk15_final_enabled: + qk15_handle = qk_smem_producer.acquire_and_advance( + expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + ) + qk15_tma_mbar = ctm.Array(qk15_handle.barrier, shape=1) + if is_leader_cta: + if cutlass.const_expr(use_smem_page_plan): + qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(15) + ) + else: + qk15_page0, qk15_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, ctm.Int32(15), csr_page_begin, csr_page_count + ) + _load_runtime_t336_qk_tile( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk15_tma_mbar, + qk_bulk_ready_mbar_ptr, + pair_overlap_sync_mbars, + sB, + sSFB, + tidx, + bidz, + physical_sfb_blocks_per_l, + csr_page_begin, + csr_page_count, + qk15_page0, + qk15_page1, + ctm.Int32(15), + ctm.Int32(0), + page_size, + qk15_final=True, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + ) + cute.arch.cp_async_bulk_wait_group(0, read=True) + q_smem_producer.tail() + if warp_idx == TMA_V_WARP_ID: + v_physical_page0 = ctm.Int32(0) + v_physical_page1 = ctm.Int32(0) + if cutlass.const_expr(use_smem_page_plan): + v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(0) + ) + for kv_tile_idx in cutlass.range(0, kv_tiles, 1, unroll=1): + v_handle = v_smem_producer.acquire_and_advance( + _sleep_until_producer_slot_free(v_smem_producer) + ) + v_tma_mbar = ctm.Array(v_handle.barrier, shape=1) + v_stage = kv_tile_idx % ctm.Int32(SMEM_P4_V_PIPELINE_STAGES) + _load_v_tile_stage( + mPageTable_pl, + sPageIdPlan, + tma_v_tile_ptr, + tma_v_pair_ptr, + tma_vsf_ptr, + tma_vsf_stage_ptr, + v_tma_mbar, + sV, + sVSF, + tidx, + bidy, + bidz, + cta_rank, + csr_page_begin, + csr_page_count, + v_physical_page0, + v_physical_page1, + kv_tile_idx=kv_tile_idx, + v_page_offset=v_page_offset, + page_size=page_size, + v_tma_phase=v_stage, + stage=v_stage, + manage_mbarrier=False, + use_smem_page_plan=use_smem_page_plan, + use_consecutive_page_pair=use_consecutive_page_pair, + ) + if cutlass.const_expr(use_smem_page_plan): + if kv_tile_idx + ctm.Int32(1) < kv_tiles: + v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ctm.Int32(1) + ) + v_smem_producer.tail() + if warp_idx == PV_ISSUE_WARP_ID: + prims.barrier(barrier_id=TMEM_BAR_ID, number_of_threads=TMEM_BAR_THREADS) + if is_leader_cta: + qk_prefix_end = min(kv_tiles, ctm.Int32(3)) + for qk_prefix_li_idx in cutlass.range(1, qk_prefix_end, 1, unroll=1): + if cutlass.const_expr(use_smem_page_plan): + rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) + else: + rowmeta_page0, rowmeta_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count + ) + _issue_runtime_t336_qk_prefix_rank1( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_bulk_ready_mbar_ptr, + sB, + sSFB, + tidx, + bidz, + physical_sfb_blocks_per_l, + csr_page_begin, + csr_page_count, + rowmeta_page0, + rowmeta_page1, + qk_prefix_li_idx, + qk_prefix_li_idx, + page_size, + ) + # PV issue runs on this warp (SMSP 3) in parallel with the QK issue of + # the MMA warp (SMSP 0). + pv_tmem_raw_addr = tmem_holding_buf.load() + pv_tmem_base_col_id = pv_tmem_raw_addr & ctm.Int32(65535) + pv_tmem_base_row_id = pv_tmem_raw_addr >> ctm.Int32(16) + for pv_li_idx in cutlass.range(0, kv_tiles, 1, unroll=1): + v_smem_consumer.wait_and_advance() + _runtime_issue_pv_tile_dispatch( + sP, + sV, + sPSF, + sVSF, + pv_tmem_base_col_id, + pv_tmem_base_row_id, + pair_overlap_sync_mbars, + p4_smem_slot0_mbars, + pv_full_mbar, + v_smem_mbars, + pv_li_idx, + kv_tiles, + ctm.Boolean(pv_li_idx == ctm.Int32(0)), + ) + if warp_idx == MMA_WARP_ID: + if is_leader_cta: + tmem_raw_addr = tmem_holding_buf.load() + tmem_base_col_id = tmem_raw_addr & ctm.Int32(65535) + tmem_base_row_id = tmem_raw_addr >> ctm.Int32(16) + q_res_handle0 = q_smem_consumer.wait() + q_smem_consumer.advance() + q_res_handle1 = q_smem_consumer.wait() + q_smem_consumer.advance() + q_res_handle2 = q_smem_consumer.wait() + q_smem_consumer.advance() + _fence_async_shared_cta() + if prims.elect_sync(): + for qsf_ready_cta_rank in ctm.range_constexpr(CLUSTER_SHAPE_MNK[0]): + qsf_ready_peer_mbar = _mapa_shared_cluster( + qsf_owner_ready_mbar, qsf_ready_cta_rank + ) + _mbarrier_arrive_release_cta_shared_cluster(qsf_ready_peer_mbar) + prims.barrier( + barrier_id=SMEM_P4_QSF_OWNER_BAR_ID, + number_of_threads=SMEM_P4_QSF_OWNER_BAR_THREADS, + ) + prims.barrier( + barrier_id=SMEM_P4_QSF_OWNER_BAR_ID, + number_of_threads=SMEM_P4_QSF_OWNER_BAR_THREADS, + ) + while not prims.mbarrier_try_wait_parity( + qsf_owner_done_mbar, 0, time_limit=10000000 + ): + pass + stream_li_total = kv_tiles + dsl_hints.cfence() + qk0_handle = qk_smem_consumer.wait() + qk_smem_consumer.advance() + # Tile 0 sits in K slot 0; its stages go through the K scale ring + # like every refill. + qk0_desc_ksf = _make_smem_p4_qk_ksf_stage_desc(sSFB, ctm.Int32(0), 0, True, False) + for qk0_stage in ctm.range_constexpr(SMEM_P4_QK_DATA_STAGES): + if cutlass.const_expr(qk0_stage < SMEM_P4_QK_FULL_DATA_STAGES): + _issue_smem_p4_qk_ksf_stage_from_desc( + qk0_desc_ksf.advance_start_address(qk0_stage * QK_SFB_STAGE_BYTES), + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + qk0_stage, + False, + ) + else: + _issue_smem_p4_qk_ksf_range_from_base( + sSFB, + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + ctm.Int32(0), + SMEM_P4_QK_FULL_DATA_STAGES, + 1, + True, + False, + ) + _issue_smem_p4_qk_stage_burst_from_base( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id, + tmem_base_row_id, + score_slot_idx=0, + qk_kblock_start=qk0_stage * SMEM_P4_QK_DATA_STAGE_KBLOCKS, + qk_kblocks=SMEM_P4_QK_DATA_STAGE_KBLOCKS + if qk0_stage < SMEM_P4_QK_FULL_DATA_STAGES + else QK_TAIL_DATA_STAGE_K_DIM // SMEM_P4_QK_MMA_KBLOCK_DIM, + acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, + sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, + use_resident_q_stages=True, + b_stage_base=ctm.Int32(0), + fence_after=qk0_stage == SMEM_P4_QK_DATA_STAGES - 1, + b_smem_is_compact_slot_base=True, + ) + if prims.elect_sync(): + prims.tcgen05_commit( + qk_full_mbar, + multicast_mask=ctm.Int16(3), + group=prims.CTAGroup.CTA_2, + ) + qk0_handle.release() + # QK runs two tiles ahead of PV: QK(t) refills score bank t % 2 as + # soon as the producers have read tile t - 2 from it instead of + # waiting behind PV(t - 2), which needs that tile's P. + for qk_ahead_li_idx in cutlass.range( + 1, min(stream_li_total, ctm.Int32(2)), 1, unroll=1 + ): + _runtime_refill_qk_tile_at( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id, + tmem_base_row_id, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_ahead_li_idx, + stream_li_total, + ) + for stream_li_idx in cutlass.range(0, stream_li_total, 1, unroll=1): + qk_ahead_li_idx = stream_li_idx + ctm.Int32(2) + if qk_ahead_li_idx < stream_li_total: + score_bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + score_consumed_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_TMEM_BANKS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + score_bank_idx + ), + score_consumed_phase, + time_limit=10000000, + ): + pass + _runtime_refill_qk_tile_at( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id, + tmem_base_row_id, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_ahead_li_idx, + stream_li_total, + ) + q_res_handle0.release() + q_res_handle1.release() + q_res_handle2.release() + if warp_idx == MMA_WARP_ID and (not is_leader_cta): + while not prims.mbarrier_try_wait_parity(qsf_owner_ready_mbar, 0, time_limit=10000000): + pass + _fence_async_shared_cta() + prims.barrier( + barrier_id=SMEM_P4_QSF_OWNER_BAR_ID, + number_of_threads=SMEM_P4_QSF_OWNER_BAR_THREADS, + ) + prims.barrier( + barrier_id=SMEM_P4_QSF_OWNER_BAR_ID, + number_of_threads=SMEM_P4_QSF_OWNER_BAR_THREADS, + ) + if warp_idx < MMA_WARP_ID: + prims.setmaxregister(SMEM_P4_ALU_REGISTER_BUDGET, prims.SetMaxRegisterAction.INCREASE) + if warp_idx < ctm.Int32(SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS // 32): + atomic_rowmax_row = warp_idx * ctm.Int32(32) + (tidx & ctm.Int32(31)) + sAtomicRunningRowMax[atomic_rowmax_row] = _float_to_ordered_u32_for_atomic_max( + ctm.Float32(-ctm.Float32.inf) + ) + if warp_idx < ctm.Int32(2 * SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS // 32): + sPipelinedRowMax[warp_idx * ctm.Int32(32) + (tidx & ctm.Int32(31))] = ( + _float_to_ordered_u32_for_atomic_max(ctm.Float32(-ctm.Float32.inf)) + ) + prims.barrier(barrier_id=TMEM_BAR_ID, number_of_threads=TMEM_BAR_THREADS) + is_correction_warp = warp_idx >= ctm.Int32(SMEM_P4_CORRECTION_WARP_ID_BEGIN) + if warp_idx < ctm.Int32(SMEM_P4_COMPUTE_WARPS): + qsf_tmem_raw_addr = tmem_holding_buf.load() + qsf_tmem_base_col_id = qsf_tmem_raw_addr & ctm.Int32(65535) + qsf_tmem_base_row_id = qsf_tmem_raw_addr >> ctm.Int32(16) + prims.barrier( + barrier_id=SMEM_P4_QSF_OWNER_BAR_ID, + number_of_threads=SMEM_P4_QSF_OWNER_BAR_THREADS, + ) + if warp_idx < ctm.Int32(SMEM_P4_BMM1_M // 32): + _permute_raw_trtllm_q_tail_in_smem(sA, warp_idx, tidx & ctm.Int32(31)) + _fence_async_shared_cta() + _store_raw_trtllm_qsf_128dp_unique_owner_to_tmem_from_smem( + sSFA, + qsf_tmem_base_col_id, + qsf_tmem_base_row_id, + cta_rank, + warp_idx, + tidx & ctm.Int32(31), + ) + prims.barrier( + barrier_id=SMEM_P4_QSF_OWNER_BAR_ID, + number_of_threads=SMEM_P4_QSF_OWNER_BAR_THREADS, + ) + if warp_idx == 0: + if prims.elect_sync(): + qsf_done_peer_mbar = _mapa_shared_cluster(qsf_owner_done_mbar, ctm.Int32(0)) + _mbarrier_arrive_release_cta_shared_cluster(qsf_done_peer_mbar) + stream_li_total = kv_tiles + producer_valid_k = valid_k_for_l + if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): + # Both producer column bands own the same score row. Each CTA owns + # 64 rows of the request's [query token, head] rows. + producer_row = ( + cta_rank * ctm.Int32(SMEM_P4_BMM1_M) + + (warp_idx & ctm.Int32(1)) * ctm.Int32(32) + + (tidx & ctm.Int32(31)) + ) + producer_query = ctm.min( + (packed_tile_row0 + producer_row) // ctm.Int32(packed_query_heads), + ctm.Int32(query_len_per_seq - 1), + ) + producer_bound = _packed_query_bound( + mHelixKvBounds_l, + page_batch, + producer_query, + packed_request_valid_k, + query_len_per_seq, + use_helix_kv_bounds, + ) + producer_valid_k = ctm.min( + ctm.max(producer_bound - split_base_k, ctm.Int32(0)), valid_k_for_l + ) + if cutlass.const_expr(folded_query_heads > 0): + # Both producer column bands own the same score row. Each CTA owns + # 64 head rows; packed rows are [logical query, logical head]. + producer_row = ( + cta_rank * ctm.Int32(SMEM_P4_BMM1_M) + + (warp_idx & ctm.Int32(1)) * ctm.Int32(32) + + (tidx & ctm.Int32(31)) + ) + logical_query = producer_row // ctm.Int32(folded_query_heads) + producer_valid_k = ctm.min( + valid_k_for_l, ctm.max(folded_first_valid_k + logical_query, ctm.Int32(0)) + ) + running_row_max = ctm.Float32(-ctm.Float32.inf) + running_row_anchor = ctm.Float32(-ctm.Float32.inf) + p4_tmem_raw_addr = tmem_holding_buf.load() + p4_tmem_base_col_id = p4_tmem_raw_addr & ctm.Int32(65535) + p4_tmem_base_row_id = p4_tmem_raw_addr >> ctm.Int32(16) + corr_local_warp_idx = warp_idx & ctm.Int32(3) + corr_row_band = corr_local_warp_idx & ctm.Int32(1) + corr_col_band = corr_local_warp_idx >> ctm.Int32(1) + corr_row_id = ( + p4_tmem_base_row_id + corr_row_band * ctm.Int32(64) + corr_col_band * ctm.Int32(32) + ) + corr_n_tile_idx = ctm.Int32(1) - (warp_idx >> ctm.Int32(2)) + corr_col_id = ( + p4_tmem_base_col_id + + ctm.Int32(SMEM_P4_O_ACC_COL_OFFSETS[0]) + + corr_n_tile_idx * ctm.Int32(SMEM_P4_O_ACC_TILE_COLS) + ) + corr_tmem_addr = corr_row_id << ctm.Int32(16) | corr_col_id + corr_o_acc_tmem = ctm.inttoptr(corr_tmem_addr, 6, ctm.Float32) + leader_pair_overlap_sync_mbars = _mapa_shared_cluster(pair_overlap_sync_mbars, ctm.Int32(0)) + leader_pair_overlap_sync_mbars_dsmem = leader_pair_overlap_sync_mbars + carried_dq1024 = ctm.Int32(0) + carried_old_scale = ctm.Float32(1.0) + running_part_sum = ctm.Float32(0.0) + stream_li_idx = ctm.Int32(0) + stream_li_remaining = stream_li_total + if cutlass.const_expr(folded_query_heads == 0): + # Leading tiles that need no per-row mask run pipelined; the final tile always goes + # through the per-tile loop below. + pipelined_tiles = stream_li_total - ctm.Int32(1) + if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): + pipelined_tiles = ctm.min( + pipelined_tiles, + _load_smem_i32( + sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) + ) + // ctm.Int32(KV_TILE), + ) + if pipelined_tiles >= ctm.Int32(2): + ( + running_row_max, + running_row_anchor, + running_part_sum, + ) = _run_pipelined_leading_tiles( + sAtomicRunningRowMax, + sPipelinedRowMax, + sP, + sPSF, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + producer_valid_k, + softmax_scale_log2, + pv_psf_rescale, + pipelined_tiles, + use_mixed_imlp, + ) + stream_li_idx = pipelined_tiles + stream_li_remaining = stream_li_total - pipelined_tiles + while stream_li_remaining > ctm.Int32(1): + producer_needs_mask = ctm.Boolean(False) + if cutlass.const_expr(folded_query_heads > 0): + # The MTP window may straddle a 256-token tile boundary, so + # the penultimate tile can also need per-row causal masking. + producer_needs_mask = folded_first_valid_k < ( + (stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE) + ) + if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): + # Rows of one tile differ by at most query_len tokens, so only + # tiles past the earliest row bound need per-row masking. The + # bound is re-read from smem each tile: kept in a register it + # was spilled at the producers' register ceiling and reloaded + # from local memory here instead. + producer_needs_mask = _load_smem_i32( + sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) + ) < ((stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE)) + ( + running_row_max, + running_row_anchor, + carried_dq1024, + carried_old_scale, + running_part_sum, + ) = _runtime_t336_producer_tile_v23( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + producer_valid_k, + softmax_scale_log2, + pv_psf_rescale, + running_row_max, + running_row_anchor, + carried_dq1024, + carried_old_scale, + running_part_sum, + stream_li_idx, + stream_li_idx > ctm.Int32(0), + producer_needs_mask, + False, + use_mixed_imlp, + folded_query_heads > 0 or (packed_query_heads > 0 and query_len_per_seq > 1), + ) + stream_li_idx = stream_li_idx + ctm.Int32(1) + stream_li_remaining = stream_li_remaining - ctm.Int32(1) + final_tile_k_end = (stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE) + tail_needs_mask = valid_k_for_l < final_tile_k_end + if cutlass.const_expr(folded_query_heads > 0): + tail_needs_mask = folded_first_valid_k < final_tile_k_end + if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): + tail_needs_mask = ( + _load_smem_i32( + sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) + ) + < final_tile_k_end + ) + ( + running_row_max, + running_row_anchor, + carried_dq1024, + carried_old_scale, + running_part_sum, + ) = _runtime_t336_producer_tile_v23( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + producer_valid_k, + softmax_scale_log2, + pv_psf_rescale, + running_row_max, + running_row_anchor, + carried_dq1024, + carried_old_scale, + running_part_sum, + stream_li_idx, + stream_li_idx > ctm.Int32(0), + tail_needs_mask, + True, + use_mixed_imlp, + ) + # Publish each warp's denominator slice: part = score half * 2 + column band. + part_local_warp_idx = warp_idx & ctm.Int32(3) + part_idx = (warp_idx >> ctm.Int32(2)) * ctm.Int32(2) + (part_local_warp_idx >> ctm.Int32(1)) + part_row = (part_local_warp_idx & ctm.Int32(1)) * ctm.Int32(32) + (tidx & ctm.Int32(31)) + sScoreSlotRowMax[part_idx * ctm.Int32(SMEM_P4_BMM1_M) + part_row] = running_part_sum + prims.barrier( + barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, + number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, + ) + if cutlass.const_expr(launch_dependents == LAUNCH_DEPENDENTS_EPILOGUE): + _griddepcontrol_launch_dependents() + if is_correction_warp: + final_row_state_local_row = corr_row_band * ctm.Int32(32) + (tidx & ctm.Int32(31)) + final_anchor_row_sum = ( + sScoreSlotRowMax[final_row_state_local_row] + + sScoreSlotRowMax[ctm.Int32(2 * SMEM_P4_BMM1_M) + final_row_state_local_row] + ) + ( + sScoreSlotRowMax[ctm.Int32(SMEM_P4_BMM1_M) + final_row_state_local_row] + + sScoreSlotRowMax[ctm.Int32(3 * SMEM_P4_BMM1_M) + final_row_state_local_row] + ) + # Packed rows may be empty while their tile is not (Helix tokens + # whose KV lives on another rank); judge emptiness per row. + row_valid_k = valid_k_for_l + if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): + corr_packed_row = cta_rank * ctm.Int32(SMEM_P4_BMM1_M) + final_row_state_local_row + corr_query = ctm.min( + (packed_tile_row0 + corr_packed_row) // ctm.Int32(packed_query_heads), + ctm.Int32(query_len_per_seq - 1), + ) + corr_bound = _packed_query_bound( + mHelixKvBounds_l, + page_batch, + corr_query, + packed_request_valid_k, + query_len_per_seq, + use_helix_kv_bounds, + ) + row_valid_k = ctm.min( + ctm.max(corr_bound - split_base_k, ctm.Int32(0)), valid_k_for_l + ) + final_stat_scale = ctm.Float32(1.0) + final_row_sum = final_anchor_row_sum + final_row_max = running_row_max + if cutlass.const_expr(write_softmax_stats): + final_stat_scale = softmax_scale_log2 / ctm.Float32(LOG2_E) + final_row_sum = final_anchor_row_sum * cute.exp2( + (running_row_anchor - running_row_max) * softmax_scale_log2, + fastmath=True, + ) + if row_valid_k == ctm.Int32(0) or final_anchor_row_sum == ctm.Float32(0.0): + final_stat_scale = ctm.Float32(1.0) + final_row_max = ctm.Float32(-ctm.Float32.inf) + final_row_sum = ctm.Float32(0.0) + if cutlass.const_expr(store_split): + # Convert the lazy-anchor numerator and denominator to the + # true row maximum expected by the shared reduction kernel. + # Empty splits avoid evaluating -inf - -inf. + if row_valid_k != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): + final_stat_scale = cute.exp2( + (running_row_anchor - running_row_max) * softmax_scale_log2, + fastmath=True, + ) + final_row_sum = final_anchor_row_sum * final_stat_scale + else: + final_row_sum = ctm.Float32(0.0) + output_normalizer = ctm.Float32(0.0) + if row_valid_k != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): + output_normalizer = cute.arch.rcp_approx(final_anchor_row_sum) * pv_output_scale + last_pv_li_idx = stream_li_total - ctm.Int32(1) + last_pv_slot = last_pv_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + last_pv_phase = ( + last_pv_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pv_full_mbar.subview(last_pv_slot), last_pv_phase, time_limit=10000000 + ): + pass + if cutlass.const_expr(store_split and packed_query_heads > 0): + if packed_single_task: + _store_final_o_from_tmem( + mC_mnl, + mAccum_mnl, + mRowMax_ml, + mRowSum_ml, + tmem_holding_buf, + warp_idx, + tidx, + bidx, + bidy, + store_task_idx, + m, + n, + final_row_max=final_row_max, + final_row_sum=final_row_sum, + final_stat_scale=final_stat_scale, + output_normalizer=output_normalizer, + valid_m=valid_m, + c_batch_stride=c_batch_stride, + c_batch_idx=c_batch_idx, + write_softmax_stats=write_softmax_stats, + store_split=store_split, + direct_single_task=True, + producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, + ) + else: + _store_final_o_from_tmem( + mC_mnl, + mAccum_mnl, + mRowMax_ml, + mRowSum_ml, + tmem_holding_buf, + warp_idx, + tidx, + bidx, + bidy, + store_task_idx, + m, + n, + final_row_max=final_row_max, + final_row_sum=final_row_sum, + final_stat_scale=final_stat_scale, + output_normalizer=output_normalizer, + valid_m=valid_m, + c_batch_stride=c_batch_stride, + c_batch_idx=c_batch_idx, + write_softmax_stats=write_softmax_stats, + store_split=store_split, + producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, + ) + else: + _store_final_o_from_tmem( + mC_mnl, + mAccum_mnl, + mRowMax_ml, + mRowSum_ml, + tmem_holding_buf, + warp_idx, + tidx, + bidx, + bidy, + store_task_idx, + m, + n, + final_row_max=final_row_max, + final_row_sum=final_row_sum, + final_stat_scale=final_stat_scale, + output_normalizer=output_normalizer, + valid_m=valid_m, + c_batch_stride=c_batch_stride, + c_batch_idx=c_batch_idx, + write_softmax_stats=write_softmax_stats, + store_split=store_split, + producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, + ) + prims.barrier(barrier_id=O_STORE_BAR_ID, number_of_threads=O_STORE_BAR_THREADS) + _dealloc_tmem_cluster( + tmem_dealloc_mbar, + tmem_holding_buf, + warp_idx, + cta_rank, + num_tmem_alloc_cols=num_tmem_alloc_cols, + dealloc_warp_id=SMEM_P4_CORRECTION_WARP_ID_BEGIN, + ) + + +@cute.kernel +def kernel( + mPageTable_pl: cute.Tensor, + mPageIndptr_s: cute.Tensor, + mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, + mPlan_i: cute.Tensor, + mQGlobalScale: cute.Tensor, + mKvGlobalScale: cute.Tensor, + tma_q_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_k_page_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_q_tail_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_k_tail_page_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_v_tile_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_qsf_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_qsf_tail_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_ksf_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_ksf_stage_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_ksf_tail_stage_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_vsf_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_vsf_stage_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_v_pair_desc: ctm.GridConstant[cuda_tma.TensorMap], + tma_sfb_desc: ctm.GridConstant[cuda_tma.TensorMap], + mC_mnl: cute.Tensor, + mAccum_mnl: cute.Tensor, + mRowMax_ml: cute.Tensor, + mRowSum_ml: cute.Tensor, + problem_size: tuple, + runtime_m: ctm.Int32, + num_requests: ctm.Int32, + num_splits: ctm.Int32, + valid_m: ctm.Int32, + c_batch_stride: ctm.Int32, + v_page_offset: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_output_scale: ctm.Float32, + page_size: ctm.Constexpr, + query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, + use_smem_page_plan: ctm.Constexpr, + use_mixed_imlp: ctm.Constexpr, + use_consecutive_page_pair: ctm.Constexpr, + use_ksf_gather4: ctm.Constexpr, + write_softmax_stats: ctm.Constexpr, + store_split: ctm.Constexpr = False, + folded_query_heads: ctm.Constexpr = 0, + packed_query_heads: ctm.Constexpr = 0, + launch_dependents: ctm.Constexpr = False, + pdl_wait: ctm.Constexpr = False, +) -> None: + if cutlass.const_expr(launch_dependents == LAUNCH_DEPENDENTS_ENTRY): + # The split reduction waits for this grid with griddepcontrol.wait; letting + # it launch once every CTA here has started hides its launch latency. + _griddepcontrol_launch_dependents() + # Packed split mode: grid.z holds the task budget; slots past the plan's + # used task count have no work and exit before touching smem/tmem (both + # CTAs of a cluster share bidz, so the cluster exits together). + task_active = ctm.Boolean(True) + if cutlass.const_expr(packed_query_heads > 0): + if num_splits > ctm.Int32(1): + _, _, task_idx = cute.arch.block_idx() + task_active = task_idx < mPlan_i[1] + if task_active: + _run_mla_decode_body( + mPageTable_pl, + mPageIndptr_s, + mValidK_l, + mHelixKvBounds_l, + mPlan_i, + mQGlobalScale, + mKvGlobalScale, + tma_q_desc.get_ptr(), + tma_k_page_desc.get_ptr(), + tma_q_tail_desc.get_ptr(), + tma_k_tail_page_desc.get_ptr(), + tma_v_tile_desc.get_ptr(), + tma_qsf_desc.get_ptr(), + tma_qsf_tail_desc.get_ptr(), + tma_ksf_desc.get_ptr(), + tma_ksf_stage_desc.get_ptr(), + tma_ksf_tail_stage_desc.get_ptr(), + tma_vsf_desc.get_ptr(), + tma_vsf_stage_desc.get_ptr(), + tma_v_pair_desc.get_ptr(), + tma_sfb_desc.get_ptr(), + mC_mnl, + mAccum_mnl, + mRowMax_ml, + mRowSum_ml, + problem_size, + runtime_m, + num_requests, + num_splits, + valid_m, + c_batch_stride, + v_page_offset, + softmax_scale_log2, + pv_output_scale, + page_size, + query_len_per_seq, + use_helix_kv_bounds, + use_smem_page_plan, + use_mixed_imlp, + use_consecutive_page_pair, + use_ksf_gather4, + write_softmax_stats, + store_split, + folded_query_heads, + packed_query_heads, + pdl_wait, + # The body releases the dependents itself in the late modes. + launch_dependents, + ) + + +def _cutlass_output_dtype(output_dtype: torch.dtype): + _validate_output_dtype(output_dtype) + if output_dtype == torch.bfloat16: + return cutlass.BFloat16 + return cutlass.Float16 + + +def _make_fused_ptrs( + q_data_ptr: int, + k_data_ptr: int, + q_sf_data_ptr: int, + k_sf_data_ptr: int, + b_data_ptr: int, + sfa_data_ptr: int, + sfb_data_ptr: int, + page_table_data_ptr: int, + valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, + plan_data_ptr: int, + c_data_ptr: int, + accum_data_ptr: int, + row_max_data_ptr: int, + row_sum_data_ptr: int, + page_indptr_data_ptr: int, + q_global_scale_data_ptr: int, + kv_global_scale_data_ptr: int, + output_dtype: torch.dtype = torch.float16, + row_state_align: int = 32, + split_partial_dtype: torch.dtype = torch.float32, +) -> tuple[cute.Pointer, ...]: + accum_dtype = cutlass.Float16 if split_partial_dtype == torch.float16 else cutlass.Float32 + return ( + make_ptr(cutlass.Float4E2M1FN, q_data_ptr, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Float4E2M1FN, k_data_ptr, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, q_sf_data_ptr, cute.AddressSpace.gmem, assumed_align=32), + make_ptr(cutlass.Uint8, k_sf_data_ptr, cute.AddressSpace.gmem, assumed_align=32), + make_ptr(cutlass.Float4E2M1FN, b_data_ptr, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, sfa_data_ptr, cute.AddressSpace.gmem, assumed_align=32), + make_ptr(cutlass.Uint8, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=32), + make_ptr(cutlass.Int32, page_table_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(cutlass.Int32, valid_k_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(cutlass.Int32, helix_kv_bounds_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(cutlass.Int32, plan_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr( + _cutlass_output_dtype(output_dtype), + c_data_ptr, + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr(accum_dtype, accum_data_ptr, cute.AddressSpace.gmem, assumed_align=32), + make_ptr( + cutlass.Float32, row_max_data_ptr, cute.AddressSpace.gmem, assumed_align=row_state_align + ), + make_ptr( + cutlass.Float32, row_sum_data_ptr, cute.AddressSpace.gmem, assumed_align=row_state_align + ), + make_ptr(cutlass.Int32, page_indptr_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr( + cutlass.Float32, + q_global_scale_data_ptr, + cute.AddressSpace.gmem, + assumed_align=4, + ), + make_ptr( + cutlass.Float32, + kv_global_scale_data_ptr, + cute.AddressSpace.gmem, + assumed_align=4, + ), + ) + + +def _kv_cache_3d_layout(kv_cache: torch.Tensor, page_size: int) -> KvCache3DLayout: + if kv_cache.dtype != torch.uint8: + raise TypeError(f"kv_cache must be torch.uint8, got {kv_cache.dtype}") + if page_size <= 0 or page_size % 2 != 0: + raise ValueError(f"page_size must be a positive even integer, got {page_size}") + if kv_cache.dim() == 3: + num_pages, actual_page_size, packed_dim = kv_cache.shape + stride_page, stride_token, stride_packed_dim = kv_cache.stride() + elif kv_cache.dim() >= 5: + num_pages = kv_cache.shape[0] + actual_page_size = kv_cache.shape[2] + packed_dim = kv_cache.shape[4] + stride_page = kv_cache.stride(0) + stride_token = kv_cache.stride(2) + stride_packed_dim = kv_cache.stride(4) + else: + raise ValueError( + "kv_cache must be shaped as [num_pages, page_size, packed_dim] or expose " + "TRT-LLM's 5D paged layout with page/token/packed-dim axes at 0/2/4, " + f"got shape={tuple(kv_cache.shape)}" + ) + if actual_page_size != page_size: + raise ValueError(f"page_size mismatch: argument={page_size}, cache={actual_page_size}") + return KvCache3DLayout( + num_pages=int(num_pages), + packed_dim=int(packed_dim), + stride_page=int(stride_page), + stride_token=int(stride_token), + stride_packed_dim=int(stride_packed_dim), + ) + + +def _validate_v_packed_cache_args( + v_packed: torch.Tensor, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + block_v: int, +) -> KvCache3DLayout: + if v_packed.dtype != torch.uint8: + raise TypeError(f"v_packed must be torch.uint8, got {v_packed.dtype}") + if v_packed.dim() != 2: + raise ValueError(f"v_packed must be 2D, got shape={tuple(v_packed.shape)}") + if v_head_dim <= 0 or v_head_dim % 2 != 0: + raise ValueError(f"v_head_dim must be positive and even, got {v_head_dim}") + if block_v <= 0 or v_head_dim % block_v != 0: + raise ValueError(f"v_head_dim={v_head_dim} must be divisible by block_v={block_v}") + layout = _kv_cache_3d_layout(kv_cache, page_size) + if layout.packed_dim < v_head_dim // 2: + raise ValueError( + "kv_cache packed dimension does not cover v_head_dim: " + f"packed_dim={layout.packed_dim}, v_head_dim={v_head_dim}" + ) + expected_rows = layout.num_pages * (v_head_dim // block_v) * block_v + expected_cols = page_size // 2 + if v_packed.shape[0] < expected_rows or v_packed.shape[1] < expected_cols: + raise ValueError( + "v_packed is too small for the decode V-cache layout: " + f"got shape={tuple(v_packed.shape)}, " + f"needs at least ({expected_rows}, {expected_cols})" + ) + out_s0, out_s1 = (int(stride) for stride in v_packed.stride()) + if out_s0 <= 0 or out_s1 <= 0: + raise ValueError(f"v_packed strides must be positive, got {(out_s0, out_s1)}") + minimum_row_stride = (expected_cols - 1) * out_s1 + 1 + if out_s0 < minimum_row_stride: + raise ValueError( + f"v_packed rows must not overlap: stride={(out_s0, out_s1)}, minimum row stride={minimum_row_stride}" + ) + if v_packed.device != kv_cache.device: + raise ValueError( + f"v_packed and kv_cache must be on the same device, got {v_packed.device} and {kv_cache.device}" + ) + if min(layout.stride_page, layout.stride_token, layout.stride_packed_dim) <= 0: + raise ValueError( + "kv_cache strides must be positive, got " + f"{(layout.stride_page, layout.stride_token, layout.stride_packed_dim)}" + ) + _validate_tensor_pointer_alignment("v_packed", v_packed, alignment_bytes=16) + _validate_tensor_pointer_alignment("kv_cache", kv_cache, alignment_bytes=16) + return layout + + +_FUSED_COMPILE_CACHE: dict[tuple[object, ...], Callable] = {} +_PREPARED_FUSED_CALL_CACHE_CAPACITY = 256 +_FUSED_EXECUTOR_CACHE_CAPACITY = 256 + + +@dataclass +class _PreparedFusedCall: + executor: Any + runtime_args: tuple[object, ...] + execution_args: list[Any] + adapted_args: list[Any] + + def run(self) -> int | None: + return self.executor.run_compiled_program(self.execution_args) + + +_PREPARED_FUSED_CALL_CACHE: OrderedDict[tuple[object, ...], _PreparedFusedCall] = OrderedDict() +_FUSED_EXECUTOR_CACHE: OrderedDict[tuple[object, ...], Any] = OrderedDict() +_PREPARED_FUSED_CALL_CACHE_LOCK = threading.Lock() + + +def _class_defines_callables(value: object, *names: str) -> bool: + value_type = type(value) + return all((callable(getattr(value_type, name, None)) for name in names)) + + +def _get_fused_executor(cache_key: tuple[object, ...]) -> Any | None: + with _PREPARED_FUSED_CALL_CACHE_LOCK: + executor = _FUSED_EXECUTOR_CACHE.get(cache_key) + if executor is not None: + _FUSED_EXECUTOR_CACHE.move_to_end(cache_key) + return executor + + +def _get_prepared_fused_call( + cache_key: tuple[object, ...], +) -> _PreparedFusedCall | None: + with _PREPARED_FUSED_CALL_CACHE_LOCK: + prepared = _PREPARED_FUSED_CALL_CACHE.get(cache_key) + if prepared is not None: + _PREPARED_FUSED_CALL_CACHE.move_to_end(cache_key) + return prepared + + +def _cache_prepared_fused_call( + cache_key: tuple[object, ...], candidate: _PreparedFusedCall +) -> _PreparedFusedCall: + with _PREPARED_FUSED_CALL_CACHE_LOCK: + prepared = _PREPARED_FUSED_CALL_CACHE.get(cache_key) + if prepared is not None: + _PREPARED_FUSED_CALL_CACHE.move_to_end(cache_key) + return prepared + _PREPARED_FUSED_CALL_CACHE[cache_key] = candidate + if len(_PREPARED_FUSED_CALL_CACHE) > _PREPARED_FUSED_CALL_CACHE_CAPACITY: + _PREPARED_FUSED_CALL_CACHE.popitem(last=False) + return candidate + + +def _cache_fused_executor(cache_key: tuple[object, ...], candidate: Any) -> Any: + with _PREPARED_FUSED_CALL_CACHE_LOCK: + executor = _FUSED_EXECUTOR_CACHE.get(cache_key) + if executor is not None: + _FUSED_EXECUTOR_CACHE.move_to_end(cache_key) + return executor + _FUSED_EXECUTOR_CACHE[cache_key] = candidate + if len(_FUSED_EXECUTOR_CACHE) > _FUSED_EXECUTOR_CACHE_CAPACITY: + _FUSED_EXECUTOR_CACHE.popitem(last=False) + return candidate + + +def _compile_fused( + q_data_ptr: int, + k_data_ptr: int, + q_sf_data_ptr: int, + k_sf_data_ptr: int, + b_data_ptr: int, + sfa_data_ptr: int, + sfb_data_ptr: int, + page_table_data_ptr: int, + valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, + plan_data_ptr: int, + c_data_ptr: int, + accum_data_ptr: int, + row_max_data_ptr: int, + row_sum_data_ptr: int, + page_indptr_data_ptr: int, + q_global_scale_data_ptr: int, + kv_global_scale_data_ptr: int, + m: int, + n: int, + kv: int, + batch_size: int, + page_size: int, + use_mixed_imlp: bool, + output_dtype: torch.dtype, + stream: cuda.CUstream, + num_cache_pages: int = 0, + query_len_per_seq: int = 1, + kv_page_stride_bytes: int = 0, + ksf_page_stride_bytes: int = 0, + vsf_page_stride_bytes: int = 0, + use_consecutive_page_pair: bool = False, + write_softmax_stats: bool = False, + use_helix_kv_bounds: bool = False, + store_split: bool = False, + folded_query_heads: int = 0, + packed_query_heads: int = 0, + row_state_align: int = 32, + launch_dependents: bool = False, + split_partial_dtype: torch.dtype = torch.float32, + pdl_wait: bool = False, +) -> Callable: + if type(kv) is not int: + raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") + is_page_plan_profile = kv == SMEM_P4_PAGE_PLAN_PROFILE_KV + is_bucketed_runtime_profile = ( + SMEM_P4_PAGE_PLAN_PROFILE_KV < kv <= SMEM_P4_RUNTIME_MAX_KV + and _select_runtime_kv_profile(kv) == kv + ) + if not (is_page_plan_profile or is_bucketed_runtime_profile): + raise ValueError( + "runtime-KV compile requires the fixed page-plan profile " + f"{SMEM_P4_PAGE_PLAN_PROFILE_KV} or a larger geometric profile " + f"up to {SMEM_P4_RUNTIME_MAX_KV}, got {kv}" + ) + use_ksf_gather4 = ( + not use_consecutive_page_pair and ksf_page_stride_bytes == page_size * TRTLLM_K_SF_GROUPS + ) + cache_key = ( + n, + kv, + page_size, + use_mixed_imlp, + output_dtype, + query_len_per_seq, + use_consecutive_page_pair, + use_ksf_gather4, + write_softmax_stats, + use_helix_kv_bounds, + store_split, + folded_query_heads, + packed_query_heads, + row_state_align, + launch_dependents, + split_partial_dtype, + pdl_wait, + ) + cached = _FUSED_COMPILE_CACHE.get(cache_key) + if cached is not None: + return cached + ptrs = _make_fused_ptrs( + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + sfa_data_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + helix_kv_bounds_data_ptr, + plan_data_ptr, + c_data_ptr, + accum_data_ptr, + row_max_data_ptr, + row_sum_data_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + output_dtype, + row_state_align, + split_partial_dtype, + ) + compiled = _compile_cutedsl( + fused_fp4_mla_decode_ctm, + *ptrs, + (n, kv), + ctm.Int32(SMEM_P4_CTA_GROUP_M), + ctm.Int32(1), + ctm.Int32(1), + ctm.Int32(1), + ctm.Int32(1), + ctm.Int32(SMEM_P4_CTA_GROUP_M), + ctm.Int32(1), + ctm.Int32(SMEM_P4_CTA_GROUP_M), + ctm.Int32(1), + ctm.Int32(1), + ctm.Int32(1), + ctm.Int32(0), + ctm.Int64(1), + ctm.Int64(1), + ctm.Int64(1), + ctm.Float32(1.0), + ctm.Float32(1.0), + stream, + page_size=page_size, + use_mixed_imlp=use_mixed_imlp, + query_len_per_seq=query_len_per_seq, + use_helix_kv_bounds=use_helix_kv_bounds, + use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, + use_consecutive_page_pair=use_consecutive_page_pair, + use_ksf_gather4=use_ksf_gather4, + write_softmax_stats=write_softmax_stats, + store_split=store_split, + folded_query_heads=folded_query_heads, + packed_query_heads=packed_query_heads, + launch_dependents=launch_dependents, + pdl_wait=pdl_wait, + options="--opt-level 2 --ptxas-options '--uumn'", + ) + _FUSED_COMPILE_CACHE[cache_key] = compiled + return compiled + + +def run_trtllm_fp4_mla_decode_page_native( + q_internal: torch.Tensor, + q_sf_internal: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_packed: torch.Tensor, + v_sf: torch.Tensor, + src_page_ids: torch.Tensor, + paged_kv_indptr_decode: torch.Tensor, + valid_k: torch.Tensor, + output: torch.Tensor, + *, + max_kv_len: int, + sm_scale: float, + q_global_scale: torch.Tensor, + kv_global_scale: torch.Tensor, + page_size: int = TRTLLM_PAGE_SIZE, + query_len_per_seq: int = 1, + v_pack_block: int = SMEM_P4_V_N_PER_CTA, + v_page_offset: int = 0, + q_batch_capacity: int | None = None, + assume_valid_k_prefix_tiles: int = 0, + assume_consecutive_page_prefix_tiles: int = 0, + partition_runtime_valid_k: bool = False, + enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, + split_kv_buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + folded_query_heads: int = 0, + packed_query_heads: int = 0, + num_requests: int = 0, + num_splits: int = 1, + plan: torch.Tensor | None = None, + softmax_stats: torch.Tensor | None = None, + launch_dependents: bool = False, + pdl_wait: bool = False, +) -> None: + if type(v_pack_block) is not int: + raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") + if v_pack_block not in (128, 256): + raise ValueError( + f"mufu16 page-native decode requires v_pack_block in (128, 256), got {v_pack_block}" + ) + # Packed mode: q_internal holds one [M, 320, request] entry per request whose + # M rows are [query token, head] pairs padded to a multiple of 128; grid.z + # covers num_requests * num_splits tasks. + packed = packed_query_heads > 0 + if packed and ( + folded_query_heads + or query_len_per_seq <= 0 + or num_requests <= 0 + or num_splits <= 0 + or not 0 < packed_query_heads <= SMEM_P4_CTA_GROUP_M + ): + raise ValueError( + "packed queries require 1-128 heads, positive query length, request " + "and split counts, and no folded queries" + ) + if not math.isfinite(sm_scale) or sm_scale <= 0.0: + raise ValueError(f"sm_scale must be finite and positive, got {sm_scale}") + _validate_output_dtype(output.dtype) + if page_size != TRTLLM_PAGE_SIZE: + raise ValueError( + f"page-native decode requires page_size={TRTLLM_PAGE_SIZE}, got {page_size}" + ) + # In folded mode the bounds tensor carries the first logical query's + # unclamped endpoint relative to this split. The normal Helix ABI is unchanged. + if folded_query_heads not in (0, 6): + raise ValueError("folded queries currently require six logical heads") + if folded_query_heads and ( + query_len_per_seq != 1 or helix_kv_bounds is None or split_kv_buffers is None + ): + raise ValueError("folded queries require split output and per-task causal bounds") + physical_k = _select_runtime_kv_profile(max_kv_len) + for name, value in ( + ("assume_valid_k_prefix_tiles", assume_valid_k_prefix_tiles), + ( + "assume_consecutive_page_prefix_tiles", + assume_consecutive_page_prefix_tiles, + ), + ): + if type(value) is not int: + raise TypeError(f"{name} must be an int, got {type(value).__name__}") + if value < 0: + raise ValueError(f"{name} must be nonnegative, got {value}") + if type(partition_runtime_valid_k) is not bool: + raise TypeError( + "partition_runtime_valid_k must be a bool, got " + f"{type(partition_runtime_valid_k).__name__}" + ) + if type(enable_mxi_imlp) is not bool: + raise TypeError(f"enable_mxi_imlp must be a bool, got {type(enable_mxi_imlp).__name__}") + # The valid-K and partition hints remain interface-compatible no-ops. A + # consecutive-page promise is used only when it covers every tile that + # ``max_kv_len`` permits; partial promises fall back to the arbitrary-page + # specialization rather than adding a runtime branch to the producer. + del assume_valid_k_prefix_tiles, partition_runtime_valid_k + required_consecutive_tiles = _ceil_div(max_kv_len, KV_TILE) + use_consecutive_page_pair = assume_consecutive_page_prefix_tiles >= required_consecutive_tiles + if q_internal.dtype != torch.uint8 or q_internal.dim() != 3: + raise TypeError( + "q_internal must be a uint8 [M, Q640/2, L] tensor, " + f"got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" + ) + physical_m, q_bytes, q_rows = q_internal.shape + l_batch = q_rows + valid_m = physical_m + if packed: + if q_rows != num_requests: + raise ValueError(f"packed Q must hold {num_requests} requests, got {q_rows}") + valid_m = packed_query_heads * query_len_per_seq + if ( + physical_m % SMEM_P4_CTA_GROUP_M != 0 + or physical_m > PACKED_MAX_M + or not valid_m <= physical_m < valid_m + SMEM_P4_CTA_GROUP_M + ): + raise ValueError( + f"packed M must pad {valid_m} rows to the next multiple of " + f"{SMEM_P4_CTA_GROUP_M} (<= {PACKED_MAX_M}), got {physical_m}" + ) + l_batch = num_requests * num_splits + if (softmax_row_max is None) != (softmax_row_sum is None): + raise ValueError("softmax_row_max and softmax_row_sum must be provided together") + write_softmax_stats = softmax_row_max is not None + row_state_row_stride = 1 + row_state_batch_stride = physical_m + row_state_align = 32 + if softmax_stats is not None: + # Interleaved [request * valid_m, 2] (max, sum) written in place. + if not packed or num_splits != 1 or write_softmax_stats: + raise ValueError("softmax_stats requires packed single-split decode without row stats") + if ( + softmax_stats.dtype != torch.float32 + or softmax_stats.shape != (num_requests * valid_m, 2) + or not softmax_stats.is_contiguous() + ): + raise ValueError( + f"softmax_stats must be contiguous float32 [{num_requests * valid_m}, 2]" + ) + _validate_tensor_pointer_alignment("softmax_stats", softmax_stats, alignment_bytes=8) + write_softmax_stats = True + row_state_row_stride = 2 + row_state_batch_stride = valid_m * 2 + row_state_align = 4 + if softmax_row_max is not None: + expected_stats_shape = (l_batch, physical_m) + for name, tensor in ( + ("softmax_row_max", softmax_row_max), + ("softmax_row_sum", softmax_row_sum), + ): + if ( + tensor.dtype != torch.float32 + or tensor.shape != expected_stats_shape + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous float32 with shape {expected_stats_shape}" + ) + _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) + if l_batch <= 0 or l_batch > CUDA_GRID_Z_MAX: + raise ValueError(f"queries must be in [1, {CUDA_GRID_Z_MAX}], got {l_batch}") + if q_batch_capacity is None: + q_batch_capacity = q_rows + if not isinstance(q_batch_capacity, int): + raise TypeError( + f"q_batch_capacity must be an int or None, got {type(q_batch_capacity).__name__}" + ) + if q_batch_capacity < q_rows or q_batch_capacity > CUDA_GRID_Z_MAX: + raise ValueError( + "q_batch_capacity must cover active queries and fit grid.z: " + f"queries={q_rows}, capacity={q_batch_capacity}, limit={CUDA_GRID_Z_MAX}" + ) + if packed: + num_sequences = num_requests + else: + if query_len_per_seq <= 0 or l_batch % query_len_per_seq != 0: + raise ValueError(f"query_len_per_seq={query_len_per_seq} must divide queries={l_batch}") + num_sequences = l_batch // query_len_per_seq + if q_bytes != QK_LOGICAL_DIM // 2: + raise ValueError(f"q_internal dimension 1 must be {QK_LOGICAL_DIM // 2}, got {q_bytes}") + expected_q_l_stride = physical_m * (QK_LOGICAL_DIM // 2) + if q_internal.stride()[:2] != (QK_LOGICAL_DIM // 2, 1) or ( + q_rows > 1 and q_internal.stride(2) != expected_q_l_stride + ): + raise ValueError( + f"q_internal must use batch-slow [M,320,L] storage, got strides {q_internal.stride()}" + ) + if not packed and physical_m not in SMEM_P4_PIPELINE_DIRECT_TARGET_M_OPTIONS: + raise ValueError( + f"page-native decode M must be in {SMEM_P4_PIPELINE_DIRECT_TARGET_M_OPTIONS}, got {physical_m}" + ) + if q_batch_capacity > q_rows: + q_storage_bytes = ( + q_internal.untyped_storage().nbytes() + - q_internal.storage_offset() * q_internal.element_size() + ) + required_q_storage_bytes = q_batch_capacity * expected_q_l_stride + if q_storage_bytes < required_q_storage_bytes: + raise ValueError( + f"q_internal storage has {q_storage_bytes}B, requires {required_q_storage_bytes}B" + ) + expected_output_shape = (physical_m, TRTLLM_V_HEAD_DIM, l_batch) + if packed: + expected_output_shape = (valid_m, TRTLLM_V_HEAD_DIM, num_requests) + if output.shape != expected_output_shape: + raise ValueError( + f"output must use {expected_output_shape} batch-slow metadata, got {tuple(output.shape)}" + ) + expected_output_l_stride = expected_output_shape[0] * TRTLLM_V_HEAD_DIM + if output.stride()[:2] != (TRTLLM_V_HEAD_DIM, 1) or ( + expected_output_shape[2] > 1 and output.stride(2) != expected_output_l_stride + ): + raise ValueError(f"output must be batch-slow contiguous, got strides {output.stride()}") + if src_page_ids.dtype != torch.int32 or src_page_ids.dim() != 1: + raise ValueError( + "src_page_ids must be a 1D int32 physical-page list, " + f"got dtype={src_page_ids.dtype} shape={tuple(src_page_ids.shape)}" + ) + src_page_id_count = src_page_ids.numel() + if src_page_id_count == 0 or src_page_ids.stride(0) != 1: + raise ValueError( + "src_page_ids must be non-empty and contiguous, " + f"got numel={src_page_id_count} stride={src_page_ids.stride()}" + ) + if src_page_id_count > INT32_MAX: + raise ValueError( + f"src_page_ids exceeds the Int32 element limit: {src_page_id_count} > {INT32_MAX}" + ) + if ( + paged_kv_indptr_decode.dtype != torch.int32 + or paged_kv_indptr_decode.shape != (num_sequences + 1,) + or paged_kv_indptr_decode.stride(0) != 1 + ): + raise ValueError( + f"paged_kv_indptr_decode must be contiguous int32 [{num_sequences + 1}], " + f"got dtype={paged_kv_indptr_decode.dtype} " + f"shape={tuple(paged_kv_indptr_decode.shape)}" + ) + if valid_k.dtype != torch.int32 or valid_k.shape != (num_sequences,) or valid_k.stride(0) != 1: + raise ValueError( + f"valid_k must be contiguous int32 [{num_sequences}], " + f"got dtype={valid_k.dtype} shape={tuple(valid_k.shape)} " + f"stride={valid_k.stride()}" + ) + expected_bounds_len = num_requests * query_len_per_seq if packed else l_batch + if helix_kv_bounds is not None and ( + helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.shape != (expected_bounds_len,) + or helix_kv_bounds.stride(0) != 1 + ): + raise ValueError( + "helix_kv_bounds must be contiguous int32 with shape " + f"[{expected_bounds_len}], got dtype={helix_kv_bounds.dtype} " + f"shape={tuple(helix_kv_bounds.shape)} stride={helix_kv_bounds.stride()}" + ) + cache_layout = _validate_v_packed_cache_args( + v_packed, + kv_cache, + v_head_dim=TRTLLM_V_HEAD_DIM, + page_size=page_size, + block_v=v_pack_block, + ) + if ( + cache_layout.packed_dim < TRTLLM_K_STORAGE_DIM // 2 + or cache_layout.stride_packed_dim != 1 + or cache_layout.stride_token != TRTLLM_K_STORAGE_DIM // 2 + or (cache_layout.stride_page < page_size * (TRTLLM_K_STORAGE_DIM // 2)) + or (cache_layout.stride_page % 16 != 0) + ): + raise ValueError( + "kv_cache must expose physical [page,128,K640/2] storage with a " + "contiguous token payload and an optional padded page stride" + ) + num_cache_pages = cache_layout.num_pages + page_table_capacity = num_sequences * (physical_k // page_size) + if page_table_capacity > INT32_MAX: + raise ValueError( + "bucketed CSR capacity exceeds the Int32 layout limit: " + f"{page_table_capacity} > {INT32_MAX}" + ) + if not packed and src_page_id_count > page_table_capacity: + # Packed tasks address the page list by its real length. + raise ValueError( + f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_id_count} > {page_table_capacity}" + ) + expected_v_rows = num_cache_pages * (TRTLLM_V_HEAD_DIM // v_pack_block) * v_pack_block + if ( + v_packed.shape[0] < expected_v_rows + or v_packed.shape[1] < page_size // 2 + or v_packed.stride() != (page_size // 2, 1) + ): + raise ValueError( + "v_packed must expose the persistent physical-page layout " + f"[>= {expected_v_rows}, >= {page_size // 2}] with strides " + f"{(page_size // 2, 1)}, got shape={tuple(v_packed.shape)} " + f"strides={v_packed.stride()}" + ) + v_rows_per_page = TRTLLM_V_HEAD_DIM + if v_packed.shape[0] % v_rows_per_page != 0: + raise ValueError( + "v_packed row count must be a whole number of physical pages, " + f"got {v_packed.shape[0]} rows for {v_rows_per_page} rows/page" + ) + if not isinstance(v_page_offset, int): + raise TypeError(f"v_page_offset must be an int, got {type(v_page_offset).__name__}") + num_v_cache_pages = v_packed.shape[0] // v_rows_per_page + if v_page_offset < 0 or v_page_offset > INT32_MAX: + raise ValueError(f"v_page_offset must be in [0, {INT32_MAX}], got {v_page_offset}") + if num_v_cache_pages > INT32_MAX: + raise ValueError( + f"v_packed exceeds the Int32 physical-page limit: {num_v_cache_pages} > {INT32_MAX}" + ) + if v_page_offset + num_cache_pages > num_v_cache_pages: + raise ValueError( + "v_packed does not cover the requested physical-page range: " + f"offset={v_page_offset}, pages={num_cache_pages}, " + f"available={num_v_cache_pages}" + ) + if ( + q_sf_internal.element_size() != 1 + or not q_sf_internal.is_contiguous() + or q_sf_internal.numel() < physical_m * QK_SF_GROUPS * q_rows + ): + raise ValueError("q_sf_internal must be a contiguous byte-sized QSF buffer") + ksf_payload_bytes = page_size * TRTLLM_K_SF_GROUPS + if q_batch_capacity > q_rows: + q_sf_storage_bytes = ( + q_sf_internal.untyped_storage().nbytes() + - q_sf_internal.storage_offset() * q_sf_internal.element_size() + ) + required_q_sf_storage_bytes = q_batch_capacity * physical_m * QK_SF_GROUPS + if q_sf_storage_bytes < required_q_sf_storage_bytes: + raise ValueError( + f"q_sf_internal storage has {q_sf_storage_bytes}B, requires {required_q_sf_storage_bytes}B" + ) + vsf_payload_bytes = TRTLLM_V_HEAD_DIM * (page_size // SF_VEC_SIZE) + for name, tensor, payload_bytes, expected_pages in ( + ("sf_cache", sf_cache, ksf_payload_bytes, num_cache_pages), + ("v_sf", v_sf, vsf_payload_bytes, num_v_cache_pages), + ): + if tensor.element_size() != 1 or tensor.dim() == 0: + raise ValueError(f"{name} must be a byte-sized paged tensor") + if tensor.shape[0] != expected_pages or tensor.stride(-1) != 1: + raise ValueError( + f"{name} must expose {expected_pages} physical pages with a " + "contiguous innermost dimension, " + f"got shape={tuple(tensor.shape)} strides={tensor.stride()}" + ) + if tensor.numel() < expected_pages * payload_bytes: + raise ValueError( + f"{name} has {tensor.numel()} logical bytes, requires at least {expected_pages * payload_bytes}" + ) + if tensor.stride(0) < payload_bytes: + raise ValueError( + f"{name} page stride={tensor.stride(0)}B is smaller than its {payload_bytes}B payload" + ) + if tensor.stride(0) % 16 != 0: + raise ValueError(f"{name} page stride={tensor.stride(0)}B must be 16B aligned") + device = q_internal.device + tensors = ( + q_sf_internal, + kv_cache, + sf_cache, + v_packed, + v_sf, + src_page_ids, + paged_kv_indptr_decode, + valid_k, + output, + q_global_scale, + kv_global_scale, + ) + if softmax_row_max is not None: + tensors += (softmax_row_max, softmax_row_sum) + if softmax_stats is not None: + tensors += (softmax_stats,) + if helix_kv_bounds is not None: + tensors += (helix_kv_bounds,) + if device.type != "cuda" or any((tensor.device != device for tensor in tensors)): + raise ValueError("all page-native decode tensors must share one CUDA device") + if q_global_scale.dtype != torch.float32 or q_global_scale.numel() != 1: + raise TypeError("q_global_scale must be a scalar FP32 CUDA tensor") + if kv_global_scale.dtype != torch.float32 or kv_global_scale.numel() != 1: + raise TypeError("kv_global_scale must be a scalar FP32 CUDA tensor") + if split_kv_buffers is not None: + if write_softmax_stats: + raise ValueError("split-KV buffers and softmax stats outputs are mutually exclusive") + expected_shapes = ( + (l_batch, physical_m, TRTLLM_V_HEAD_DIM), + (l_batch, physical_m), + (l_batch, physical_m), + ) + if len(split_kv_buffers) != 3: + raise ValueError("split_kv_buffers requires numerator, maxima, and sums") + expected_dtypes = (split_kv_buffers[0].dtype, torch.float32, torch.float32) + if expected_dtypes[0] not in SPLIT_PARTIAL_DTYPES: + raise ValueError( + f"split-KV partial must be one of {SPLIT_PARTIAL_DTYPES}, got {expected_dtypes[0]}" + ) + for tensor, shape, dtype in zip(split_kv_buffers, expected_shapes, expected_dtypes): + if ( + tensor.shape != shape + or tensor.dtype != dtype + or tensor.device != device + or not tensor.is_contiguous() + ): + raise ValueError(f"split-KV buffer must be contiguous {dtype} {shape} on {device}") + _validate_tensor_pointer_alignment("split-KV buffer", tensor, alignment_bytes=32) + stream = _current_cu_stream() + device_index = torch.cuda.current_device() + context_result, current_context = cuda.cuCtxGetCurrent() + if context_result != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError( + f"Failed to query the current CUDA context for FP4 MLA launch: {context_result}." + ) + q_data_ptr = q_internal.data_ptr() + k_data_ptr = kv_cache.data_ptr() + q_sf_data_ptr = q_sf_internal.data_ptr() + k_sf_data_ptr = sf_cache.data_ptr() + b_data_ptr = v_packed.data_ptr() + scratch_ptr = output.data_ptr() + row_max_data_ptr = scratch_ptr + row_sum_data_ptr = scratch_ptr + if softmax_row_max is not None: + row_max_data_ptr = softmax_row_max.data_ptr() + row_sum_data_ptr = softmax_row_sum.data_ptr() + elif softmax_stats is not None: + row_max_data_ptr = softmax_stats.data_ptr() + row_sum_data_ptr = softmax_stats.data_ptr() + softmax_stats.element_size() + kernel_num_requests = num_requests if packed else num_sequences + sfb_data_ptr = v_sf.data_ptr() + page_table_data_ptr = src_page_ids.data_ptr() + valid_k_data_ptr = valid_k.data_ptr() + helix_kv_bounds_data_ptr = ( + helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else valid_k_data_ptr + ) + plan_data_ptr = valid_k_data_ptr + if packed and num_splits > 1: + plan_len = 3 + num_requests + 3 * num_requests * num_splits + if ( + plan is None + or plan.dtype != torch.int32 + or plan.numel() < plan_len + or plan.stride(0) != 1 + ): + raise ValueError( + f"packed split decode needs a contiguous int32 plan with >= {plan_len} entries" + ) + plan_data_ptr = plan.data_ptr() + c_data_ptr = output.data_ptr() + split_o_data_ptr, split_max_data_ptr, split_sum_data_ptr = ( + tuple(tensor.data_ptr() for tensor in split_kv_buffers) + if split_kv_buffers is not None + else (scratch_ptr, row_max_data_ptr, row_sum_data_ptr) + ) + split_partial_dtype = ( + split_kv_buffers[0].dtype if split_kv_buffers is not None else torch.float32 + ) + page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() + q_global_scale_data_ptr = q_global_scale.data_ptr() + kv_global_scale_data_ptr = kv_global_scale.data_ptr() + kv_page_stride_bytes = cache_layout.stride_page + ksf_page_stride_bytes = int(sf_cache.stride(0)) + vsf_page_stride_bytes = int(v_sf.stride(0)) + softmax_scale_log2 = sm_scale * LOG2_E + fused = _compile_fused( + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + scratch_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + helix_kv_bounds_data_ptr, + plan_data_ptr, + c_data_ptr, + split_o_data_ptr, + split_max_data_ptr, + split_sum_data_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + physical_m, + TRTLLM_V_HEAD_DIM, + physical_k, + l_batch, + page_size, + use_mixed_imlp=enable_mxi_imlp, + output_dtype=output.dtype, + stream=stream, + num_cache_pages=num_cache_pages, + query_len_per_seq=query_len_per_seq, + kv_page_stride_bytes=kv_page_stride_bytes, + ksf_page_stride_bytes=ksf_page_stride_bytes, + vsf_page_stride_bytes=vsf_page_stride_bytes, + use_consecutive_page_pair=use_consecutive_page_pair, + write_softmax_stats=write_softmax_stats, + use_helix_kv_bounds=helix_kv_bounds is not None, + store_split=split_kv_buffers is not None, + folded_query_heads=folded_query_heads, + packed_query_heads=packed_query_heads, + row_state_align=row_state_align, + launch_dependents=launch_dependents, + split_partial_dtype=split_partial_dtype, + pdl_wait=pdl_wait, + ) + supports_prepared = _class_defines_callables( + fused, "to", "generate_execution_args", "run_compiled_program" + ) + prepared_key = ( + fused, + device_index, + int(current_context), + int(stream), + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + helix_kv_bounds_data_ptr, + plan_data_ptr, + c_data_ptr, + split_o_data_ptr, + split_max_data_ptr, + split_sum_data_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + physical_m, + l_batch, + q_batch_capacity, + kernel_num_requests, + num_splits, + valid_m, + row_state_row_stride, + row_state_batch_stride, + src_page_id_count, + packed_query_heads, + num_cache_pages, + num_v_cache_pages, + v_page_offset, + kv_page_stride_bytes, + ksf_page_stride_bytes, + vsf_page_stride_bytes, + softmax_scale_log2, + output.dtype, + enable_mxi_imlp, + use_consecutive_page_pair, + ) + if supports_prepared: + prepared = _get_prepared_fused_call(prepared_key) + if prepared is not None: + prepared.run() + return + ptrs = _make_fused_ptrs( + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + scratch_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + helix_kv_bounds_data_ptr, + plan_data_ptr, + c_data_ptr, + split_o_data_ptr, + split_max_data_ptr, + split_sum_data_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + output.dtype, + row_state_align, + split_partial_dtype, + ) + runtime_args = ( + *ptrs, + (TRTLLM_V_HEAD_DIM, physical_k), + ctm.Int32(physical_m), + ctm.Int32(l_batch), + ctm.Int32(q_batch_capacity), + ctm.Int32(kernel_num_requests), + ctm.Int32(num_splits), + ctm.Int32(valid_m), + ctm.Int32(row_state_row_stride), + ctm.Int32(row_state_batch_stride), + ctm.Int32(src_page_id_count), + ctm.Int32(num_cache_pages), + ctm.Int32(num_v_cache_pages), + ctm.Int32(v_page_offset), + ctm.Int64(kv_page_stride_bytes), + ctm.Int64(ksf_page_stride_bytes), + ctm.Int64(vsf_page_stride_bytes), + ctm.Float32(softmax_scale_log2), + ctm.Float32(1.0), + stream, + ) + if not supports_prepared: + fused(*runtime_args) + return + executor_key = (fused, device_index, int(current_context)) + executor = _get_fused_executor(executor_key) + if executor is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "FP4 MLA CUDA Graph capture requires an eager warmup for the " + "compiled kernel, device, and CUDA context." + ) + candidate = fused.to(device_index) + if not _class_defines_callables( + candidate, "generate_execution_args", "run_compiled_program" + ): + fused(*runtime_args) + return + executor = _cache_fused_executor(executor_key, candidate) + execution_args, adapted_args = executor.generate_execution_args(*runtime_args) + prepared = _cache_prepared_fused_call( + prepared_key, + _PreparedFusedCall( + executor=executor, + runtime_args=runtime_args, + execution_args=execution_args, + adapted_args=adapted_args, + ), + ) + prepared.run() + + +def packed_query_rows(num_heads: int, query_len: int) -> tuple[int, int]: + """Return (valid rows, 128-padded rows) of one request's [query, head] rows.""" + if num_heads <= 0 or query_len <= 0: + raise ValueError( + f"packed rows need positive heads/query length, got {num_heads}/{query_len}" + ) + valid_m = num_heads * query_len + physical_m = _ceil_div(valid_m, SMEM_P4_CTA_GROUP_M) * SMEM_P4_CTA_GROUP_M + if physical_m > PACKED_MAX_M: + raise ValueError(f"packed rows {valid_m} exceed the {PACKED_MAX_M}-row kernel limit") + return valid_m, physical_m + + +def run_trtllm_fp4_mla_decode_packed( + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_packed: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + src_page_ids: torch.Tensor, + paged_kv_indptr_decode: torch.Tensor, + kv_lens: torch.Tensor, + output: torch.Tensor, + *, + max_kv_len: int, + sm_scale: float, + num_heads: int, + query_len: int, + num_requests: int, + num_splits: int = 1, + plan: torch.Tensor | None = None, + q_global_scale: torch.Tensor | None = None, + v_pack_block: int = SMEM_P4_V_N_PER_CTA, + v_page_offset: int = 0, + enable_mxi_imlp: bool = True, + helix_kv_bounds: torch.Tensor | None = None, + softmax_stats: torch.Tensor | None = None, + split_kv_buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + launch_dependents: bool = False, + pdl_wait: bool = False, +) -> None: + """Packed decode: ``q_fp4``/``q_sf`` hold ``num_requests`` entries of + ``physical_m`` rows ([query token, head] pairs, 128-padded), ``output`` is + ``[num_requests * num_heads * query_len, 512]``, ``kv_lens`` is per request + and ``helix_kv_bounds`` (optional) is per query token. ``max_kv_len`` is + the per-task KV length the runtime profile must cover. ``num_splits`` sets + the task budget (``num_requests * num_splits`` grid.z slots) that the plan + kernel (``fp4_mla_packed_reduce.launch_packed_plan``) shares among requests + in proportion to their KV length; with ``num_splits > 1`` the kernel reads + that ``plan`` and writes partials into ``split_kv_buffers`` for multi-task + requests: FP32, or the normalized FP16 split output when the partial buffer + is FP16 (``SPLIT_PARTIAL_DTYPES``) at ``begins[request] + split``. + ``launch_dependents`` lets a programmatic dependent grid (the split reduction + launched with ``use_pdl``) launch once every CTA of this grid has started. + ``pdl_wait`` launches this grid as a programmatic dependent of the previous + kernel on the stream: everything but the query load (barriers, TMEM, plan, + page table, the first KV tile) runs before ``griddepcontrol.wait``, so the + caller guarantees that only the query is written by that previous kernel.""" + valid_m, physical_m = packed_query_rows(num_heads, query_len) + if num_requests <= 0 or num_splits <= 0: + raise ValueError("packed decode needs positive request and split counts") + if q_fp4.dtype != torch.uint8 or not q_fp4.is_contiguous(): + raise ValueError("q_fp4 must be a contiguous uint8 tensor") + expected_q_bytes = num_requests * physical_m * (QK_LOGICAL_DIM // 2) + if q_fp4.numel() < expected_q_bytes: + raise ValueError(f"q_fp4 must contain {expected_q_bytes} bytes, got {q_fp4.numel()}") + if q_sf.element_size() != 1 or not q_sf.is_contiguous(): + raise ValueError("q_sf must be a contiguous byte-sized tensor") + expected_q_sf_bytes = num_requests * physical_m * QK_SF_GROUPS + if q_sf.numel() < expected_q_sf_bytes: + raise ValueError(f"q_sf must contain {expected_q_sf_bytes} bytes, got {q_sf.numel()}") + _validate_tensor_pointer_alignment("q_fp4", q_fp4, alignment_bytes=16) + _validate_tensor_pointer_alignment("q_sf", q_sf, alignment_bytes=16) + if global_scale.dtype != torch.float32 or global_scale.numel() != 1: + raise TypeError("global_scale must be a scalar FP32 CUDA tensor") + if output.shape != (num_requests * valid_m, TRTLLM_V_HEAD_DIM) or not output.is_contiguous(): + raise ValueError( + f"packed output must be contiguous [{num_requests * valid_m}, {TRTLLM_V_HEAD_DIM}], " + f"got {tuple(output.shape)}" + ) + q_batch_slow = ( + q_fp4.flatten()[:expected_q_bytes] + .view(num_requests, physical_m, QK_LOGICAL_DIM // 2) + .permute(1, 2, 0) + ) + q_sf_rows = q_sf.flatten()[:expected_q_sf_bytes] + output_batch_slow = output.view(num_requests, valid_m, TRTLLM_V_HEAD_DIM).permute(1, 2, 0) + run_trtllm_fp4_mla_decode_page_native( + q_batch_slow, + q_sf_rows, + kv_cache, + sf_cache, + v_packed, + v_sf, + src_page_ids, + paged_kv_indptr_decode, + kv_lens, + output_batch_slow, + max_kv_len=max_kv_len, + sm_scale=sm_scale, + q_global_scale=global_scale if q_global_scale is None else q_global_scale, + kv_global_scale=global_scale, + page_size=TRTLLM_PAGE_SIZE, + query_len_per_seq=query_len, + v_pack_block=v_pack_block, + v_page_offset=v_page_offset, + q_batch_capacity=num_requests, + enable_mxi_imlp=enable_mxi_imlp, + helix_kv_bounds=helix_kv_bounds, + split_kv_buffers=split_kv_buffers, + packed_query_heads=num_heads, + num_requests=num_requests, + num_splits=num_splits, + plan=plan, + softmax_stats=softmax_stats, + launch_dependents=launch_dependents, + pdl_wait=pdl_wait, + ) + + +def run_trtllm_fp4_mla_decode_page_native_from_raw( + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_packed: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + src_page_ids: torch.Tensor, + paged_kv_indptr_decode: torch.Tensor, + kv_lens: torch.Tensor, + output: torch.Tensor, + *, + max_kv_len: int, + sm_scale: float, + num_heads: int, + q_global_scale: torch.Tensor | None = None, + page_size: int = TRTLLM_PAGE_SIZE, + query_len_per_seq: int = 1, + v_pack_block: int = SMEM_P4_V_N_PER_CTA, + v_page_offset: int = 0, + q_batch_capacity: int | None = None, + assume_valid_k_prefix_tiles: int = 0, + assume_consecutive_page_prefix_tiles: int = 0, + partition_runtime_valid_k: bool = False, + enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, + split_kv_buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + folded_query_heads: int = 0, +) -> None: + if type(v_pack_block) is not int: + raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") + if v_pack_block not in (128, 256): + raise ValueError( + f"mufu16 page-native decode requires v_pack_block in (128, 256), got {v_pack_block}" + ) + if num_heads != SMEM_P4_CTA_GROUP_M: + raise ValueError( + f"raw page-native decode requires {SMEM_P4_CTA_GROUP_M} physical heads, got {num_heads}" + ) + if q_fp4.dtype != torch.uint8 or not q_fp4.is_contiguous(): + raise ValueError("q_fp4 must be a contiguous uint8 tensor") + num_queries = int(output.shape[0]) + expected_q_bytes = num_queries * num_heads * (QK_LOGICAL_DIM // 2) + if q_fp4.numel() != expected_q_bytes: + raise ValueError(f"q_fp4 must contain {expected_q_bytes} bytes, got {q_fp4.numel()}") + if q_sf.element_size() != 1 or not q_sf.is_contiguous(): + raise ValueError("q_sf must be a contiguous byte-sized tensor") + expected_q_sf_bytes = num_queries * num_heads * QK_SF_GROUPS + if q_sf.numel() != expected_q_sf_bytes: + raise ValueError(f"q_sf must contain {expected_q_sf_bytes} bytes, got {q_sf.numel()}") + _validate_tensor_pointer_alignment("q_fp4", q_fp4, alignment_bytes=16) + _validate_tensor_pointer_alignment("q_sf", q_sf, alignment_bytes=16) + if global_scale.dtype != torch.float32 or global_scale.numel() != 1: + raise TypeError("global_scale must be a scalar FP32 CUDA tensor") + q_batch_slow = q_fp4.view(num_queries, num_heads, QK_LOGICAL_DIM // 2).permute(1, 2, 0) + run_trtllm_fp4_mla_decode_page_native( + q_batch_slow, + q_sf, + kv_cache, + sf_cache, + v_packed, + v_sf, + src_page_ids, + paged_kv_indptr_decode, + kv_lens, + output.permute(1, 2, 0), + max_kv_len=max_kv_len, + sm_scale=sm_scale, + q_global_scale=global_scale if q_global_scale is None else q_global_scale, + kv_global_scale=global_scale, + page_size=page_size, + query_len_per_seq=query_len_per_seq, + v_pack_block=v_pack_block, + v_page_offset=v_page_offset, + q_batch_capacity=q_batch_capacity, + assume_valid_k_prefix_tiles=assume_valid_k_prefix_tiles, + assume_consecutive_page_prefix_tiles=assume_consecutive_page_prefix_tiles, + partition_runtime_valid_k=partition_runtime_valid_k, + enable_mxi_imlp=enable_mxi_imlp, + softmax_row_max=softmax_row_max, + softmax_row_sum=softmax_row_sum, + helix_kv_bounds=helix_kv_bounds, + split_kv_buffers=split_kv_buffers, + folded_query_heads=folded_query_heads, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_packed_reduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_packed_reduce.py new file mode 100644 index 000000000000..24dddc2dd226 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp4_mla_packed_reduce.py @@ -0,0 +1,1905 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuTe DSL task plan and split reduction of K3's packed FP4 MLA decode. + +The plan and reduction of the packed decode kernel in ``fp4_mla_mufu16`` (same directory): the +kernel stores partials by plan slot (ordered plans) and optionally as normalized FP16, so its plan +and reduction are these and not the FP4 MLA backend's. + +The decode kernel writes one FP32 partial per task covering the request's +128-padded [query token, head] rows together with the raw row maximum and the +row sum relative to it. Tasks follow the balanced plan of the decode kernel +(consecutive tasks per request, more for longer KV), which this kernel rederives +from the same per-request bounds. It merges the tasks of every valid row into the +BF16/FP16 output and, for Helix, the per-row (max, sum) softmax stats. Global FP4 +scales and KV lengths are read on the device, so CUDA graph replays follow them. +""" + +import heapq +import math +from collections.abc import Callable + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import make_ptr + +from ..features import active as _features +from .fp4_mla_mufu16 import ( + KV_TILE, + LOG2_E, + PACKED_MIN_TILES_PER_TASK, + _compile_cutedsl, + _current_cu_stream, + _cutlass_output_dtype, + _packed_query_bound, + _validate_output_dtype, + _validate_tensor_pointer_alignment, +) + +_PACKED_REDUCE_COMPILE_CACHE: dict[ + tuple[int, bool, bool, torch.dtype, bool, bool, torch.dtype], Callable +] = {} +_PACKED_PLAN_COMPILE_CACHE: dict[tuple[int, bool, bool, bool, bool], Callable] = {} +_REDUCE_THREADS = 128 +# The rule plan (no simulation) runs one CTA of this many threads. +_RULE_PLAN_THREADS = 128 +# One warp per candidate task length of the wave plan. +_PLAN_THREADS = 1024 +# The makespan plan simulates every candidate in its own warp. +_LPT_PLAN_THREADS = 384 +# Requests one plan CTA stages in shared memory (4 KiB of tiles + begins). +PLAN_MAX_REQUESTS = 1024 +# The plan simulates the decode grid for this many task lengths, one warp each: +# the rule's (warp 0) and the longest request cut into 1..PLAN_SIM_CANDIDATES - 1 pieces. +PLAN_SIM_CANDIDATES = _PLAN_THREADS // 32 +# Waves of more slots (four per lane) or batches of more requests keep the rule's length. +PLAN_SIM_MAX_SLOTS = 128 +PLAN_SIM_MAX_REQUESTS = 64 +# Candidates of more tasks are skipped: the simulation leaves out the reduction, whose +# cost grows with the partials. +PLAN_SIM_MAX_TASKS = 256 +# Fixed cost of one task (Q load, pipeline fill, epilogue) in KV tiles, and the extra cost of +# a task whose request is split (FP32 partial and row stats, merged by the reduction). Fitted +# on Rubin to the decode + reduction time of forced plans. +PLAN_TASK_COST_TILES = 12 +PLAN_SPLIT_COST_TILES = 24 +_VALUE_DIM = 512 +# Tasks whose partial loads a warp issues back-to-back before any is consumed +# with a full 512-column band (16 vector loads per lane; narrower bands take +# proportionally more tasks per chunk). The reduction is latency-bound, so +# this sets its memory-level parallelism; 32 loads in flight measured slower +# (register pressure), as did issuing all loads ahead of the weight math. +_TASKS_PER_CHUNK = 4 +# Persistent reducer grid: up to this many CTAs per SM, one warp per output row. +_REDUCE_CTAS_PER_SM = 4 + + +def reduce_pdl_enabled() -> bool: + """``features.fp4_mla_reduce_pdl``: the reduction launches as a programmatic dependent of the + decode kernel (which then triggers its dependents as soon as every decode CTA has started) and + waits for the decode grid with griddepcontrol.wait after staging the plan.""" + return _features().fp4_mla_reduce_pdl + + +def f16_partials_enabled() -> bool: + """``features.fp4_mla_f16_partials``: the decode kernel stores each split's normalized output + (numerator / row sum, in output units) as FP16 instead of the FP32 max-relative numerator, + which halves the partial store and load traffic; the merge then weighs every split by its + max-scaled row sum and applies no output scale.""" + return _features().fp4_mla_f16_partials + + +@cutlass.dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [], + "griddepcontrol.wait;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# Makespan plan (``plan_packed_tasks_lpt``): the task length is chosen by +# simulating the cluster schedule of a few candidates and the tasks are +# launched longest first. +# Fixed cost of one decode cluster (launch, plan and page loads, Q load, +# pipeline fill, partial store) in 256-token KV tiles. Rubin measures ~7 us for +# a cluster alone and ~16 us when a full wave starts at once, against ~0.8 us +# per tile. +LPT_CLUSTER_COST_TILES = 12 +# Cluster slots the schedule simulation tracks (four per lane of a warp). +LPT_MAX_SLOTS = 128 +# Batches up to this size choose the task length by simulation (the plan +# kernel then takes ~8 / 11 / 20 us at 8 / 16 / 32 requests on Rubin); larger +# ones keep the balanced length and only take the longest-first order. +LPT_SIM_MAX_REQUESTS = 32 +LPT_NUM_CANDIDATES = 10 +# A length other than the balanced one must shorten the modelled makespan by +# 1/50 of it and by half a cluster cost. +LPT_MIN_GAIN_DIVISOR = 50 +_LPT_SLOT_BITS = 7 +_LPT_INF_LOAD = (1 << 23) - 1 + + +@cute.jit +def _reduce_packed_row( + mPartial_trc: cute.Tensor, + mRowMax_tr: cute.Tensor, + mRowSum_tr: cute.Tensor, + mOut_rc: cute.Tensor, + mStatsMax_r: cute.Tensor, + mStatsSum_r: cute.Tensor, + lane: cutlass.Int32, + row: cutlass.Int32, + local_row: cutlass.Int32, + band: cutlass.Int32, + task_begin: cutlass.Int32, + request_tasks: cutlass.Int32, + physical_m: cutlass.Int32, + scale: cutlass.Float32, + output_scale: cutlass.Float32, + band_columns: cutlass.Constexpr, + write_stats: cutlass.Constexpr, +) -> None: + """One warp merges the planned tasks of one (output row, column band). + + Lane ``l`` owns columns ``band * band_columns + i * 128 + 4 * l .. + 4`` for + ``i < band_columns // 128``, so every vector load and store of the warp + covers 128 contiguous columns. Tasks are consumed in chunks sized to keep 16 + vector loads in flight per lane; the per-task weights are recomputed by + every lane from the row state (L1-resident), so no shuffles or shared + memory are needed. Empty tasks (row sum 0) carry zero weight and a row + with no visible key anywhere yields O=0 and stats (-inf, 0), the neutral + Helix contribution. FP16 partials are normalized split outputs, weighted + by the max-scaled row sums. + """ + columns_per_lane: cutlass.Constexpr = 4 + partial_bits: cutlass.Constexpr = mPartial_trc.element_type.width + normalized: cutlass.Constexpr = partial_bits == 16 + column_groups: cutlass.Constexpr = band_columns // (32 * columns_per_lane) + tasks_per_chunk: cutlass.Constexpr = _TASKS_PER_CHUNK * (_VALUE_DIM // band_columns) + band_base = band * cutlass.Int32(band_columns) + lane * cutlass.Int32(4) + gPartial_arr = cutlass.make_array_view(mPartial_trc) + gOut_arr = cutlass.make_array_view(mOut_rc) + global_row_max = cutlass.Float32(-cutlass.Float32.inf) + for task_idx in cutlass.range(lane, request_tasks, 32, unroll=1): + task = task_begin + task_idx + if mRowSum_tr[task, local_row] != cutlass.Float32(0.0): + global_row_max = cute.arch.fmax(global_row_max, mRowMax_tr[task, local_row] * scale) + global_row_max = cute.arch.warp_redux_sync(global_row_max, "fmax") + + # Four independent accumulators (the DSL tracks rebinding per object). + numerator0 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + numerator1 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + numerator2 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + numerator3 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + denominator = cutlass.Float32(0.0) + for chunk_base in cutlass.range(0, request_tasks, tasks_per_chunk, unroll=1): + chunk_vecs = [] + chunk_weights = [] + for slot in cutlass.range_constexpr(tasks_per_chunk): + # Out-of-range slots re-read the chunk's last valid task with zero weight. + task_idx = cutlass.min(chunk_base + slot, request_tasks - cutlass.Int32(1)) + task = task_begin + task_idx + task_weight = cutlass.Float32(0.0) + if chunk_base + slot < request_tasks: + task_row_sum = mRowSum_tr[task, local_row] + if task_row_sum != cutlass.Float32(0.0): + task_weight = cute.exp2( + mRowMax_tr[task, local_row] * scale - global_row_max, fastmath=True + ) + denominator += task_row_sum * task_weight + if cutlass.const_expr(normalized): + task_weight = task_weight * task_row_sum + chunk_weights.append(task_weight) + partial_row = (task * physical_m + local_row) * cutlass.Int32(_VALUE_DIM) + band_base + for group in cutlass.range_constexpr(column_groups): + group_vec = ( + gPartial_arr.subview(partial_row + cutlass.Int32(group * 128)) + .data_ptr() + .load(count=columns_per_lane, alignment=partial_bits // 8 * columns_per_lane) + ) + if cutlass.const_expr(normalized): + group_vec = group_vec.to(cutlass.Float32) + chunk_vecs.append(group_vec) + for slot in cutlass.range_constexpr(tasks_per_chunk): + # Empty tasks keep zero weight; their payload is never trusted. + if chunk_weights[slot] != cutlass.Float32(0.0): + numerator0 = numerator0 + chunk_vecs[slot * column_groups] * chunk_weights[slot] + if cutlass.const_expr(column_groups > 1): + numerator1 = ( + numerator1 + chunk_vecs[slot * column_groups + 1] * chunk_weights[slot] + ) + if cutlass.const_expr(column_groups > 2): + numerator2 = ( + numerator2 + chunk_vecs[slot * column_groups + 2] * chunk_weights[slot] + ) + numerator3 = ( + numerator3 + chunk_vecs[slot * column_groups + 3] * chunk_weights[slot] + ) + + normalizer = cutlass.Float32(0.0) + if denominator != cutlass.Float32(0.0): + normalizer = output_scale / denominator + out_row = row * cutlass.Int32(_VALUE_DIM) + band_base + gOut_arr.store( + (numerator0 * normalizer).to(mOut_rc.element_type), + idx=out_row, + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + if cutlass.const_expr(column_groups > 1): + gOut_arr.store( + (numerator1 * normalizer).to(mOut_rc.element_type), + idx=out_row + cutlass.Int32(128), + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + if cutlass.const_expr(column_groups > 2): + gOut_arr.store( + (numerator2 * normalizer).to(mOut_rc.element_type), + idx=out_row + cutlass.Int32(256), + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + gOut_arr.store( + (numerator3 * normalizer).to(mOut_rc.element_type), + idx=out_row + cutlass.Int32(384), + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + if cutlass.const_expr(write_stats): + if band == cutlass.Int32(0) and lane == cutlass.Int32(0): + # Helix combine convention: natural-log-domain scaled max and the + # sum relative to it; an empty row contributes (-inf, 0). + stats_max = cutlass.Float32(-cutlass.Float32.inf) + if denominator != cutlass.Float32(0.0): + stats_max = global_row_max * cutlass.Float32(1.0 / LOG2_E) + mStatsMax_r[row] = stats_max + mStatsSum_r[row] = denominator + + +@cute.kernel +def packed_split_reduce_kernel( + mPartial_trc: cute.Tensor, + mRowMax_tr: cute.Tensor, + mRowSum_tr: cute.Tensor, + mOut_rc: cute.Tensor, + mStatsMax_r: cute.Tensor, + mStatsSum_r: cute.Tensor, + mQScale: cute.Tensor, + mKvScale: cute.Tensor, + mPlan_i: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + num_requests: cutlass.Int32, + valid_m: cutlass.Int32, + physical_m: cutlass.Int32, + band_columns: cutlass.Constexpr, + write_stats: cutlass.Constexpr, + spread_stats: cutlass.Constexpr, + use_pdl: cutlass.Constexpr = False, +) -> None: + """Persistent grid: each warp walks (row, column band) items of multi-task requests. + + Every CTA stages the plan's task prefix (plan[2 .. 2 + num_requests]) and + scans the requests. The (row, band) pairs of requests with more than one + task form one compact item space that the grid's warps stride through, one + item per warp at a time (narrower bands when few rows would leave SMs + idle); a single-task request was stored by the decode kernel itself and, + for Helix, only needs its (max, sum) stats derived from its row state. + With no multi-task request the kernel exits after the plan scan, so an + even batch pays only the launch. + """ + bands: cutlass.Constexpr = _VALUE_DIM // band_columns + bidx, _, _ = cute.arch.block_idx() + grid_x, _, _ = cute.arch.grid_dim() + tidx, _, _ = cute.arch.thread_idx() + warp_idx = tidx // cutlass.Int32(32) + lane = tidx % cutlass.Int32(32) + warps_per_cta: cutlass.Constexpr = _REDUCE_THREADS // 32 + sBegins = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, PLAN_MAX_REQUESTS + 1), + cute.make_layout(PLAN_MAX_REQUESTS + 1), + ) + for entry in cutlass.range(tidx, num_requests + cutlass.Int32(1), _REDUCE_THREADS, unroll=1): + sBegins[entry] = mPlan_i[cutlass.Int32(2) + entry] + cute.arch.sync_threads() + scale = softmax_scale_log2 / (mQScale[0] * mKvScale[0]) + if cutlass.const_expr(use_pdl): + # Partials and row state come from the decode grid this one depends on; + # the plan and the global scales were written before it. + _griddepcontrol_wait() + # Compact item c of the multi-task (row, band) pairs belongs to warp slot + # c % item_stride. + item_slot = bidx * cutlass.Int32(warps_per_cta) + warp_idx + item_stride = grid_x * cutlass.Int32(warps_per_cta) + compact_offset = cutlass.Int32(0) + # With spread_stats, 128-row stats chunk c of the single-task requests + # (counted across requests) belongs to CTA c % grid. + stats_chunk_offset = cutlass.Int32(0) + for request in cutlass.range(0, num_requests, 1, unroll=1): + task_begin = cutlass.Int32(sBegins[request]) + request_tasks = cutlass.Int32(sBegins[request + cutlass.Int32(1)]) - task_begin + if request_tasks > cutlass.Int32(1): + request_items = valid_m * cutlass.Int32(bands) + first_item = compact_offset + ( + (item_slot - compact_offset % item_stride + item_stride) % item_stride + ) + for compact_item in cutlass.range( + first_item, compact_offset + request_items, item_stride, unroll=1 + ): + local_item = compact_item - compact_offset + local_row = local_item // cutlass.Int32(bands) + band = local_item - local_row * cutlass.Int32(bands) + _reduce_packed_row( + mPartial_trc, + mRowMax_tr, + mRowSum_tr, + mOut_rc, + mStatsMax_r, + mStatsSum_r, + lane, + request * valid_m + local_row, + local_row, + band, + task_begin, + request_tasks, + physical_m, + scale, + output_scale, + band_columns, + write_stats, + ) + compact_offset += request_items + else: + if cutlass.const_expr(write_stats): + # Thread-parallel over the request's rows: chunk k of 128 rows + # is owned by CTA k % grid, or by CTA (offset + k) % grid with + # spread_stats so that consecutive requests use other CTAs. + row_chunks = (valid_m + cutlass.Int32(127)) // cutlass.Int32(128) + first_chunk = bidx + if cutlass.const_expr(spread_stats): + first_chunk = (bidx - stats_chunk_offset % grid_x + grid_x) % grid_x + stats_chunk_offset += row_chunks + for chunk in cutlass.range(first_chunk, row_chunks, grid_x, unroll=1): + local_row = chunk * cutlass.Int32(128) + tidx + if local_row < valid_m: + row = request * valid_m + local_row + row_sum = mRowSum_tr[task_begin, local_row] + stats_max = cutlass.Float32(-cutlass.Float32.inf) + if row_sum != cutlass.Float32(0.0): + stats_max = ( + mRowMax_tr[task_begin, local_row] + * scale + * cutlass.Float32(1.0 / LOG2_E) + ) + mStatsMax_r[row] = stats_max + mStatsSum_r[row] = row_sum + + +@cute.jit +def packed_split_reduce_host( + partial_ptr: cute.Pointer, + row_max_ptr: cute.Pointer, + row_sum_ptr: cute.Pointer, + out_ptr: cute.Pointer, + stats_max_ptr: cute.Pointer, + stats_sum_ptr: cute.Pointer, + q_scale_ptr: cute.Pointer, + kv_scale_ptr: cute.Pointer, + plan_ptr: cute.Pointer, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + stream: cuda.CUstream, + num_requests: cutlass.Int32, + num_splits: cutlass.Int32, + valid_m: cutlass.Int32, + physical_m: cutlass.Int32, + total_rows: cutlass.Int32, + stats_stride: cutlass.Int32, + grid_ctas: cutlass.Int32, + band_columns: cutlass.Constexpr, + write_stats: cutlass.Constexpr, + spread_stats: cutlass.Constexpr, + use_pdl: cutlass.Constexpr = False, +) -> None: + tasks = num_requests * num_splits + partial_layout = cute.make_layout( + (tasks, physical_m, _VALUE_DIM), + stride=(cutlass.Int64(physical_m) * _VALUE_DIM, _VALUE_DIM, 1), + ) + row_state_layout = cute.make_layout((tasks, physical_m), stride=(physical_m, 1)) + out_layout = cute.make_layout((total_rows, _VALUE_DIM), stride=(_VALUE_DIM, 1)) + stats_layout = cute.make_layout((total_rows,), stride=(stats_stride,)) + mPartial_trc = cute.make_tensor(partial_ptr, partial_layout) + mRowMax_tr = cute.make_tensor(row_max_ptr, row_state_layout) + mRowSum_tr = cute.make_tensor(row_sum_ptr, row_state_layout) + mOut_rc = cute.make_tensor(out_ptr, out_layout) + mStatsMax_r = cute.make_tensor(stats_max_ptr, stats_layout) + mStatsSum_r = cute.make_tensor(stats_sum_ptr, stats_layout) + mQScale = cute.make_tensor(q_scale_ptr, cute.make_layout((1,), stride=(1,))) + mKvScale = cute.make_tensor(kv_scale_ptr, cute.make_layout((1,), stride=(1,))) + mPlan_i = cute.make_tensor( + plan_ptr, cute.make_layout((cutlass.Int32(3) + num_requests + 3 * tasks,), stride=(1,)) + ) + packed_split_reduce_kernel( + mPartial_trc, + mRowMax_tr, + mRowSum_tr, + mOut_rc, + mStatsMax_r, + mStatsSum_r, + mQScale, + mKvScale, + mPlan_i, + softmax_scale_log2, + output_scale, + num_requests, + valid_m, + physical_m, + band_columns, + write_stats, + spread_stats, + use_pdl, + ).launch( + grid=(grid_ctas, 1, 1), + block=(_REDUCE_THREADS, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) + + +@cute.jit +def _lpt_remainder_rank( + sTiles: cute.Tensor, + request: cutlass.Int32, + remainder: cutlass.Int32, + num_requests: cutlass.Int32, + tiles_per_task: cutlass.Int32, +) -> cutlass.Int32: + """Position of ``request``'s last, shorter task among all such tasks, longest first.""" + rank = cutlass.Int32(0) + for other in cutlass.range(0, num_requests, 1, unroll=1): + other_tiles = sTiles[other] + other_rem = other_tiles - (other_tiles // tiles_per_task) * tiles_per_task + if other_rem > remainder or (other_rem == remainder and other < request): + rank += cutlass.Int32(1) + return rank + + +@cute.jit +def _lpt_slot_key(load: cutlass.Int32, slot: cutlass.Int32, wave_clusters: cutlass.Int32): + key = cutlass.Int32(_LPT_INF_LOAD << _LPT_SLOT_BITS) | slot + if slot < wave_clusters: + key = (load << cutlass.Int32(_LPT_SLOT_BITS)) | slot + return key + + +@cute.jit +def _lpt_simulate_makespan( + sTiles: cute.Tensor, + sRem: cute.Tensor, + warp: cutlass.Int32, + lane: cutlass.Int32, + num_requests: cutlass.Int32, + task_capacity: cutlass.Int32, + tiles_per_task: cutlass.Int32, + m_tiles: cutlass.Int32, + wave_clusters: cutlass.Int32, +) -> cutlass.Int32: + """One warp: makespan (in KV tiles) of the longest-first schedule of ``tiles_per_task``. + + Every task is ``m_tiles`` clusters that each cost the task's tiles plus + ``LPT_CLUSTER_COST_TILES``; clusters start in launch order on the earliest + free of ``wave_clusters`` slots (lane ``l`` tracks slots ``l + 32 i``). + The full-length tasks come first and need no simulation (equal clusters + spread round-robin); the last, shorter task of each request follows, + longest first; ``sRem[:, warp]`` holds them. Returns the int32 maximum + when the tasks exceed the capacity. + """ + cost: cutlass.Constexpr = LPT_CLUSTER_COST_TILES + full_tasks = cutlass.Int32(0) + used_tasks = cutlass.Int32(0) + remainder_tasks = cutlass.Int32(0) + for request in cutlass.range(lane, num_requests, 32, unroll=1): + request_tiles = sTiles[request] + request_full = request_tiles // tiles_per_task + remainder = request_tiles - request_full * tiles_per_task + full_tasks += request_full + used_tasks += request_full + if remainder > cutlass.Int32(0): + used_tasks += cutlass.Int32(1) + remainder_tasks += cutlass.Int32(1) + rank = _lpt_remainder_rank(sTiles, request, remainder, num_requests, tiles_per_task) + sRem[rank, warp] = remainder + full_tasks = cute.arch.warp_redux_sync(full_tasks, "add") + used_tasks = cute.arch.warp_redux_sync(used_tasks, "add") + remainder_tasks = cute.arch.warp_redux_sync(remainder_tasks, "add") + cute.arch.sync_warp() + full_clusters = full_tasks * m_tiles + full_cost = tiles_per_task + cutlass.Int32(cost) + base = (full_clusters // wave_clusters) * full_cost + extra = full_clusters - (full_clusters // wave_clusters) * wave_clusters + load0 = base + load1 = base + load2 = base + load3 = base + if lane < extra: + load0 += full_cost + if lane + cutlass.Int32(32) < extra: + load1 += full_cost + if lane + cutlass.Int32(64) < extra: + load2 += full_cost + if lane + cutlass.Int32(96) < extra: + load3 += full_cost + for task in cutlass.range(0, remainder_tasks, 1, unroll=1): + task_cost = sRem[task, warp] + cutlass.Int32(cost) + for _ in cutlass.range(0, m_tiles, 1, unroll=1): + key = cutlass.min( + cutlass.min( + _lpt_slot_key(load0, lane, wave_clusters), + _lpt_slot_key(load1, lane + cutlass.Int32(32), wave_clusters), + ), + cutlass.min( + _lpt_slot_key(load2, lane + cutlass.Int32(64), wave_clusters), + _lpt_slot_key(load3, lane + cutlass.Int32(96), wave_clusters), + ), + ) + key = cute.arch.warp_redux_sync(key, "min") + slot = key & cutlass.Int32((1 << _LPT_SLOT_BITS) - 1) + if slot == lane: + load0 += task_cost + if slot == lane + cutlass.Int32(32): + load1 += task_cost + if slot == lane + cutlass.Int32(64): + load2 += task_cost + if slot == lane + cutlass.Int32(96): + load3 += task_cost + makespan = cutlass.Int32(0) + if lane < wave_clusters: + makespan = cutlass.max(makespan, load0) + if lane + cutlass.Int32(32) < wave_clusters: + makespan = cutlass.max(makespan, load1) + if lane + cutlass.Int32(64) < wave_clusters: + makespan = cutlass.max(makespan, load2) + if lane + cutlass.Int32(96) < wave_clusters: + makespan = cutlass.max(makespan, load3) + makespan = cute.arch.warp_redux_sync(makespan, "max") + if used_tasks > task_capacity: + makespan = cutlass.Int32((1 << 31) - 1) + # The next candidate reuses sRem. + cute.arch.sync_warp() + return makespan + + +@cute.jit +def _lpt_clamp_tiles( + tiles: cutlass.Int32, longest: cutlass.Int32, max_tiles_per_task: cutlass.Int32 +) -> cutlass.Int32: + return cutlass.min( + cutlass.max(tiles, cutlass.Int32(PACKED_MIN_TILES_PER_TASK)), + cutlass.max(cutlass.min(longest, max_tiles_per_task), cutlass.Int32(1)), + ) + + +@cute.jit +def _lpt_write_candidates( + sCand: cute.Tensor, + balanced_tiles: cutlass.Int32, + total_tiles: cutlass.Int32, + longest: cutlass.Int32, + num_requests: cutlass.Int32, + m_tiles: cutlass.Int32, + wave_clusters: cutlass.Int32, + max_tiles_per_task: cutlass.Int32, +) -> None: + """Task lengths the makespan plan simulates (see ``_lpt_candidates``).""" + cost: cutlass.Constexpr = LPT_CLUSTER_COST_TILES + sCand[0] = balanced_tiles + # The per-slot load of a perfectly balanced schedule, every stream cut once + # more at a slot boundary: ceil((W + cost * (slots + streams)) / slots) - cost. + streams = num_requests * m_tiles + ideal = ( + total_tiles * m_tiles + cutlass.Int32(cost) * (wave_clusters + streams) + wave_clusters - 1 + ) // wave_clusters - cutlass.Int32(cost) + ideal = cutlass.max(ideal, cutlass.Int32(1)) + for k in cutlass.range_constexpr(4): + sCand[1 + 2 * k] = _lpt_clamp_tiles( + ideal // cutlass.Int32(k + 1), longest, max_tiles_per_task + ) + sCand[2 + 2 * k] = _lpt_clamp_tiles( + (ideal * cutlass.Int32(26)) // cutlass.Int32(25 * (k + 1)), longest, max_tiles_per_task + ) + # Even batch closed form on the longest request: q full tasks per request + # on q * streams slots, the remainders in k rounds on the rest. + uniform_tiles = longest + uniform_makespan = cutlass.Int32((1 << 31) - 1) + for q in cutlass.range(1, 16, 1, unroll=1): + full_slots = q * streams + if full_slots < wave_clusters: + free_slots = wave_clusters - full_slots + rounds = (streams + free_slots - cutlass.Int32(1)) // free_slots + denom = cutlass.Int32(1) + rounds * q + tiles = ( + rounds * (longest + cutlass.Int32(cost)) - cutlass.Int32(cost) + denom - 1 + ) // denom + tiles = cutlass.max(cutlass.min(tiles, longest), cutlass.Int32(1)) + remainder = longest - q * tiles + makespan = tiles + cutlass.Int32(cost) + if remainder > cutlass.Int32(0): + makespan = cutlass.max(makespan, rounds * (remainder + cutlass.Int32(cost))) + if makespan < uniform_makespan: + uniform_makespan = makespan + uniform_tiles = tiles + sCand[2 * 4 + 1] = _lpt_clamp_tiles(uniform_tiles, longest, max_tiles_per_task) + + +@cute.jit +def _write_task_entry( + mPlan_i: cute.Tensor, + tables: cutlass.Int32, + task_capacity: cutlass.Int32, + task: cutlass.Int32, + request: cutlass.Int32, + split_idx: cutlass.Int32, + request_tasks: cutlass.Int32, +) -> None: + mPlan_i[tables + task] = request + mPlan_i[tables + task_capacity + task] = split_idx + mPlan_i[tables + task_capacity + task_capacity + task] = request_tasks + + +@cute.jit +def _lpt_plan_tables( + sTiles: cute.Tensor, + sBegins: cute.Tensor, + sFull: cute.Tensor, + sScalars: cute.Tensor, + sCand: cute.Tensor, + sMake: cute.Tensor, + sRemAll: cute.Tensor, + mPlan_i: cute.Tensor, + tidx: cutlass.Int32, + num_requests: cutlass.Int32, + task_capacity: cutlass.Int32, + m_tiles: cutlass.Int32, + wave_clusters: cutlass.Int32, + max_tiles_per_task: cutlass.Int32, + plan_threads: cutlass.Constexpr, +) -> None: + """Makespan plan: pick the task length, then lay the tasks out longest first. + + ``sScalars`` holds the balanced task length, the total and the longest + request tiles. For batches of up to ``LPT_SIM_MAX_REQUESTS`` requests with + spare task slots, each warp simulates candidates (``_lpt_write_candidates``) + and the shortest makespan wins (ties keep the longer tasks) if it beats the + balanced length by the margin of ``plan_packed_tasks_lpt``. A request + keeps ``ceil(tiles / length)`` consecutive partial slots (``begins``, as + the reduction reads them) while its tasks are launched in a different + order: every full-length task first, then each request's last, shorter + task, longest first. The decode kernel stores each task's partial at + ``begins[request] + split``. + """ + warp = tidx // cutlass.Int32(32) + lane = tidx - warp * cutlass.Int32(32) + warps: cutlass.Constexpr = plan_threads // 32 + slots = cutlass.min(wave_clusters, cutlass.Int32(LPT_MAX_SLOTS)) + balanced = sScalars[0] + simulate = cutlass.Boolean(task_capacity > num_requests) & cutlass.Boolean( + num_requests <= cutlass.Int32(LPT_SIM_MAX_REQUESTS) + ) + if simulate: + if tidx == cutlass.Int32(0): + _lpt_write_candidates( + sCand, + balanced, + sScalars[1], + sScalars[2], + num_requests, + m_tiles, + slots, + max_tiles_per_task, + ) + cute.arch.sync_threads() + for cand in cutlass.range(warp, LPT_NUM_CANDIDATES, warps, unroll=1): + makespan = _lpt_simulate_makespan( + sTiles, + sRemAll, + warp, + lane, + num_requests, + task_capacity, + sCand[cand], + m_tiles, + slots, + ) + if lane == cutlass.Int32(0): + sMake[cand] = makespan + cute.arch.sync_threads() + tiles_per_task = balanced + if tidx == cutlass.Int32(0): + if simulate: + best = sMake[0] + # Another length must beat the balanced one by the margin: the + # model leaves out the reduction of the extra partials and the + # contention of a fuller wave. + limit = best - cutlass.max( + best // cutlass.Int32(LPT_MIN_GAIN_DIVISOR), + cutlass.Int32(LPT_CLUSTER_COST_TILES // 2), + ) + for cand in cutlass.range(1, LPT_NUM_CANDIDATES, 1, unroll=1): + cand_makespan = sMake[cand] + cand_tiles = sCand[cand] + if cand_makespan <= limit and ( + cand_makespan < best or (cand_makespan == best and cand_tiles > tiles_per_task) + ): + best = cand_makespan + tiles_per_task = cand_tiles + task_begin = cutlass.Int32(0) + full_begin = cutlass.Int32(0) + for request_idx in cutlass.range(0, num_requests, 1, unroll=1): + request_tiles = sTiles[request_idx] + sBegins[request_idx] = task_begin + sFull[request_idx] = full_begin + mPlan_i[cutlass.Int32(2) + request_idx] = task_begin + full_begin = full_begin + request_tiles // tiles_per_task + task_begin = task_begin + (request_tiles + tiles_per_task - cutlass.Int32(1)) // ( + tiles_per_task + ) + sBegins[num_requests] = task_begin + sFull[num_requests] = full_begin + mPlan_i[cutlass.Int32(2) + num_requests] = task_begin + mPlan_i[0] = tiles_per_task + mPlan_i[1] = task_begin + sScalars[0] = tiles_per_task + cute.arch.sync_threads() + tiles_per_task = sScalars[0] + tables = cutlass.Int32(3) + num_requests + full_total = sFull[num_requests] + used_tasks = sBegins[num_requests] + for request_idx in cutlass.range(tidx, num_requests, plan_threads, unroll=1): + task_tiles = sTiles[request_idx] + task_full = task_tiles // tiles_per_task + task_rem = task_tiles - task_full * tiles_per_task + task_count = sBegins[request_idx + cutlass.Int32(1)] - sBegins[request_idx] + task_full_begin = sFull[request_idx] + for split_idx in cutlass.range(0, task_full, 1, unroll=1): + _write_task_entry( + mPlan_i, + tables, + task_capacity, + task_full_begin + split_idx, + request_idx, + split_idx, + task_count, + ) + if task_rem > cutlass.Int32(0): + rem_rank = _lpt_remainder_rank( + sTiles, request_idx, task_rem, num_requests, tiles_per_task + ) + _write_task_entry( + mPlan_i, + tables, + task_capacity, + full_total + rem_rank, + request_idx, + task_full, + task_count, + ) + for task_idx in cutlass.range(used_tasks + tidx, task_capacity, plan_threads, unroll=1): + _write_task_entry( + mPlan_i, + tables, + task_capacity, + task_idx, + num_requests - cutlass.Int32(1), + cutlass.Int32(1 << 15), + cutlass.Int32(1 << 15), + ) + + +@cute.jit +def _plan_simulated_makespan( + sTiles: cute.Tensor, + num_requests: cutlass.Int32, + tiles_per_task: cutlass.Int32, + wave_tasks: cutlass.Int32, + lane: cutlass.Int32, +) -> cutlass.Int32: + """Makespan (KV tiles) of the decode grid for one task length, simulated by one warp. + + The grid runs the tasks in the plan's table order (``plan_task_order``): every + request's tasks but its last, all ``tiles_per_task`` long, then the last tasks, + longest first. Each task takes the slot of ``wave_tasks`` that frees first and + holds it for its KV tiles plus PLAN_TASK_COST_TILES, plus PLAN_SPLIT_COST_TILES + when its request has several tasks. Slot j lives in lane j % 32, register + j // 32; request r's last task in lane r % 32, bit r // 32. + """ + # Every task but a request's last belongs to a split request. + task_cost = tiles_per_task + cutlass.Int32(PLAN_TASK_COST_TILES + PLAN_SPLIT_COST_TILES) + full_tasks = cutlass.Int32(0) + for request_idx in cutlass.range(lane, num_requests, 32, unroll=1): + full_tasks = ( + full_tasks + + (sTiles[request_idx] + tiles_per_task - cutlass.Int32(1)) // tiles_per_task + - cutlass.Int32(1) + ) + full_tasks = cute.arch.warp_redux_sync(full_tasks, "add") + # Equal tasks fill the slots round-robin. + full_rounds = full_tasks // wave_tasks + full_extra = full_tasks - full_rounds * wave_tasks + idle = cutlass.Int32(1 << 30) + slot0 = idle + slot1 = idle + slot2 = idle + slot3 = idle + if lane < wave_tasks: + slot0 = full_rounds * task_cost + if lane < full_extra: + slot0 = slot0 + task_cost + if lane + cutlass.Int32(32) < wave_tasks: + slot1 = full_rounds * task_cost + if lane + cutlass.Int32(32) < full_extra: + slot1 = slot1 + task_cost + if lane + cutlass.Int32(64) < wave_tasks: + slot2 = full_rounds * task_cost + if lane + cutlass.Int32(64) < full_extra: + slot2 = slot2 + task_cost + if lane + cutlass.Int32(96) < wave_tasks: + slot3 = full_rounds * task_cost + if lane + cutlass.Int32(96) < full_extra: + slot3 = slot3 + task_cost + makespan = full_rounds * task_cost + if full_extra > cutlass.Int32(0): + makespan = makespan + task_cost + taken = cutlass.Int32(0) + for _ in cutlass.range(0, num_requests, 1, unroll=1): + # The longest last task left (ties: lower request), as tail * 1024 + (1023 - request). + local_key = cutlass.Int32(-1) + for request_idx in cutlass.range(lane, num_requests, 32, unroll=1): + taken_bit = (taken >> (request_idx // cutlass.Int32(32))) & cutlass.Int32(1) + if taken_bit == cutlass.Int32(0): + request_tiles = sTiles[request_idx] + request_tasks = ( + request_tiles + tiles_per_task - cutlass.Int32(1) + ) // tiles_per_task + tail = request_tiles - (request_tasks - cutlass.Int32(1)) * tiles_per_task + local_key = cutlass.max( + local_key, + tail * cutlass.Int32(PLAN_MAX_REQUESTS) + + cutlass.Int32(PLAN_MAX_REQUESTS - 1) + - request_idx, + ) + key = cute.arch.warp_redux_sync(local_key, "max") + request = cutlass.Int32(PLAN_MAX_REQUESTS - 1) - key % cutlass.Int32(PLAN_MAX_REQUESTS) + if lane == request % cutlass.Int32(32): + taken = taken | (cutlass.Int32(1) << (request // cutlass.Int32(32))) + tail_cost = key // cutlass.Int32(PLAN_MAX_REQUESTS) + cutlass.Int32(PLAN_TASK_COST_TILES) + if sTiles[request] > tiles_per_task: + tail_cost = tail_cost + cutlass.Int32(PLAN_SPLIT_COST_TILES) + local_free = cutlass.min(cutlass.min(slot0, slot1), cutlass.min(slot2, slot3)) + start = cute.arch.warp_redux_sync(local_free, "min") + owner_key = cutlass.Int32(32) + if local_free == start: + owner_key = lane + owner = cute.arch.warp_redux_sync(owner_key, "min") + if lane == owner: + if slot0 == start: + slot0 = start + tail_cost + elif slot1 == start: + slot1 = start + tail_cost + elif slot2 == start: + slot2 = start + tail_cost + else: + slot3 = start + tail_cost + makespan = cutlass.max(makespan, start + tail_cost) + return makespan + + +def _plan_threads(lpt_order: bool, wave_plan: bool) -> int: + """Threads of the plan CTA: one warp per candidate of the simulating plans.""" + if lpt_order: + return _LPT_PLAN_THREADS + return _PLAN_THREADS if wave_plan else _RULE_PLAN_THREADS + + +@cute.kernel +def packed_plan_kernel( + mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, + mPromptLens_l: cute.Tensor, + mCorrectedKvLens_l: cute.Tensor, + mGenerationLens_l: cute.Tensor, + mPlan_i: cute.Tensor, + num_requests: cutlass.Int32, + task_capacity: cutlass.Int32, + wave_tasks: cutlass.Int32, + max_tiles_per_task: cutlass.Int32, + generation_len: cutlass.Int32, + m_tiles: cutlass.Int32, + wave_clusters: cutlass.Int32, + query_len_per_seq: cutlass.Constexpr, + use_helix_kv_bounds: cutlass.Constexpr, + write_generation_lengths: cutlass.Constexpr, + lpt_order: cutlass.Constexpr, + wave_plan: cutlass.Constexpr, +) -> None: + """One CTA: the balanced task plan both decode and reduction kernels read. + + With ``write_generation_lengths`` the CTA first derives the per-step + generation lengths every FP4 MLA kernel of the step reads (``kv_lens - + prompt_lens + generation_len`` and ``generation_len`` per request) and + plans on those, so the first FP4 MLA kernel of a step pays one launch for + both; otherwise ``mValidK_l`` already holds the corrected lengths. + + ``plan = [tiles_per_task, used_tasks, begins[0..num_requests], + request[task_capacity], split[task_capacity], count[task_capacity]]``. + Request ``r`` owns ``ceil(tiles_r / tiles_per_task)`` tasks, whose partials + sit at ``begins[r] + split``; the tables list them in grid order (every + request's tasks but its last, then the last ones longest first, see + ``plan_task_order``). Slots past ``used_tasks`` map to the last request + with an out-of-range split and are never scheduled (the decode kernel exits them). + + Task length: a request is only split when it is at least 1.5x longer than + the batch's work spread over one wave, ``ceil(total / wave_tasks)``, i.e. + when its tail would outlast the machine-bound makespan (an even batch that + fills the machine is never split). The length then targets one wave of + tasks for an even, under-filled batch and two waves when the machine is + full (3 * wave_tasks <= 4 * requests) or the batch is strongly skewed + (longest > 2x average), where the extra slack lets short tasks backfill + behind the long request's pieces; a Helix rank always targets one wave + (it holds 1/N of every request's KV, so two-wave pieces are short and + the per-task cost dominates); the longest request is cut into + ``round(longest / target)`` equal pieces. If that needs more tasks than + the capacity holds, the smaller of the guaranteed-fit share + ``ceil(total / (capacity - requests))`` and the per-request share + ``ceil(longest / splits)`` that fits is used. Everything is clamped to + ``max_tiles_per_task``, the KV the compiled runtime profile covers (the + host sizes it so the clamp never needs more than the capacity). + + That length is only a candidate: one warp each, the plan + simulates the grid (``_plan_simulated_makespan``) for it and for the + longest request cut into 1..PLAN_SIM_CANDIDATES - 1 equal pieces, and + takes the length that finishes first (ties keep the longer tasks). + Without ``wave_plan`` the plan keeps + the rule's length and lists every request's tasks in request order. + The host mirror is ``plan_packed_tasks``. + + With ``lpt_order`` (the makespan plan K3's decode runs, takes precedence over + ``wave_plan``) the rule's length is one candidate of the makespan plan + (``_lpt_plan_tables``, host mirror ``plan_packed_tasks_lpt``), which also + needs the decode kernel's ``m_tiles`` clusters per task and the + ``wave_clusters`` the GPU holds at once. + """ + tidx, _, _ = cute.arch.thread_idx() + plan_threads: cutlass.Constexpr = _plan_threads(lpt_order, wave_plan) + sTiles = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, PLAN_MAX_REQUESTS), cute.make_layout(PLAN_MAX_REQUESTS) + ) + sBegins = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, PLAN_MAX_REQUESTS + 1), + cute.make_layout(PLAN_MAX_REQUESTS + 1), + ) + # The rule's task length, the total and the longest request tiles. + sScalars = cute.make_tensor(cute.arch.alloc_smem(cutlass.Int32, 4), cute.make_layout(4)) + for request_idx in cutlass.range(tidx, num_requests, plan_threads, unroll=1): + request_valid_k = cutlass.Int32(mValidK_l[request_idx]) + if cutlass.const_expr(write_generation_lengths): + request_valid_k = ( + request_valid_k - cutlass.Int32(mPromptLens_l[request_idx]) + generation_len + ) + mCorrectedKvLens_l[request_idx] = request_valid_k + mGenerationLens_l[request_idx] = generation_len + # KV tiles the request spans: its last query token's bound, at least one. + bound = _packed_query_bound( + mHelixKvBounds_l, + request_idx, + cutlass.Int32(query_len_per_seq - 1), + request_valid_k, + query_len_per_seq, + use_helix_kv_bounds, + ) + sTiles[request_idx] = cutlass.max( + (bound + cutlass.Int32(KV_TILE - 1)) // cutlass.Int32(KV_TILE), cutlass.Int32(1) + ) + cute.arch.sync_threads() + if tidx < cutlass.Int32(32): + total_tiles = cutlass.Int32(0) + longest = cutlass.Int32(0) + for request_idx in cutlass.range(tidx, num_requests, 32, unroll=1): + request_tiles = sTiles[request_idx] + total_tiles = total_tiles + request_tiles + longest = cutlass.max(longest, request_tiles) + total_tiles = cute.arch.warp_redux_sync(total_tiles, "add") + longest = cute.arch.warp_redux_sync(longest, "max") + tiles_per_task = longest + if task_capacity > num_requests: + share = cutlass.max( + (total_tiles + wave_tasks - cutlass.Int32(1)) // wave_tasks, + cutlass.Int32(PACKED_MIN_TILES_PER_TASK), + ) + if longest > share + share // cutlass.Int32(2): + target_tasks = wave_tasks + if cutlass.const_expr(not use_helix_kv_bounds): + full_machine = wave_tasks * cutlass.Int32(3) <= num_requests * cutlass.Int32(4) + strong_skew = longest * num_requests > total_tiles + total_tiles + if full_machine or strong_skew: + target_tasks = wave_tasks + wave_tasks + target = cutlass.max( + (total_tiles + target_tasks - cutlass.Int32(1)) // target_tasks, + cutlass.Int32(PACKED_MIN_TILES_PER_TASK), + ) + pieces = (longest + longest + target) // (target + target) + target = cutlass.max( + (longest + pieces - cutlass.Int32(1)) // pieces, + cutlass.Int32(PACKED_MIN_TILES_PER_TASK), + ) + used_target = cutlass.Int32(0) + for request_idx in cutlass.range(tidx, num_requests, 32, unroll=1): + used_target = ( + used_target + (sTiles[request_idx] + target - cutlass.Int32(1)) // target + ) + used_target = cute.arch.warp_redux_sync(used_target, "add") + if used_target <= task_capacity: + tiles_per_task = target + else: + spare = task_capacity - num_requests + fallback = cutlass.max( + (total_tiles + spare - cutlass.Int32(1)) // spare, + cutlass.Int32(PACKED_MIN_TILES_PER_TASK), + ) + tiles_per_task = fallback + splits = task_capacity // num_requests + even = cutlass.max( + (longest + splits - cutlass.Int32(1)) // splits, + cutlass.Int32(PACKED_MIN_TILES_PER_TASK), + ) + if even < fallback: + used_even = cutlass.Int32(0) + for request_idx in cutlass.range(tidx, num_requests, 32, unroll=1): + used_even = ( + used_even + (sTiles[request_idx] + even - cutlass.Int32(1)) // even + ) + used_even = cute.arch.warp_redux_sync(used_even, "add") + if used_even <= task_capacity: + tiles_per_task = even + tiles_per_task = cutlass.min(tiles_per_task, longest) + tiles_per_task = cutlass.min(tiles_per_task, max_tiles_per_task) + if tidx == cutlass.Int32(0): + sScalars[0] = tiles_per_task + sScalars[1] = total_tiles + sScalars[2] = longest + cute.arch.sync_threads() + if cutlass.const_expr(lpt_order): + sFull = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, PLAN_MAX_REQUESTS + 1), + cute.make_layout(PLAN_MAX_REQUESTS + 1), + ) + sCand = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, LPT_NUM_CANDIDATES), + cute.make_layout(LPT_NUM_CANDIDATES), + ) + sMake = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, LPT_NUM_CANDIDATES), + cute.make_layout(LPT_NUM_CANDIDATES), + ) + sRemAll = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, LPT_SIM_MAX_REQUESTS * (plan_threads // 32)), + cute.make_layout((LPT_SIM_MAX_REQUESTS, plan_threads // 32)), + ) + _lpt_plan_tables( + sTiles, + sBegins, + sFull, + sScalars, + sCand, + sMake, + sRemAll, + mPlan_i, + tidx, + num_requests, + task_capacity, + m_tiles, + wave_clusters, + max_tiles_per_task, + plan_threads, + ) + else: + if cutlass.const_expr(wave_plan): + # Per candidate warp: simulated makespan, then task length. + sCandidates = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, 2 * PLAN_SIM_CANDIDATES), + cute.make_layout(2 * PLAN_SIM_CANDIDATES), + ) + warp_idx = tidx // cutlass.Int32(32) + lane = tidx % cutlass.Int32(32) + fits = (wave_tasks <= cutlass.Int32(PLAN_SIM_MAX_SLOTS)) & ( + num_requests <= cutlass.Int32(PLAN_SIM_MAX_REQUESTS) + ) + simulate = (task_capacity > num_requests) & fits + if simulate: + longest = sScalars[2] + candidate = sScalars[0] + if warp_idx > cutlass.Int32(0): + candidate = cutlass.min( + cutlass.max( + (longest + warp_idx - cutlass.Int32(1)) // warp_idx, + cutlass.Int32(PACKED_MIN_TILES_PER_TASK), + ), + cutlass.min(longest, max_tiles_per_task), + ) + candidate_tasks = cutlass.Int32(0) + for request_idx in cutlass.range(lane, num_requests, 32, unroll=1): + candidate_tasks = ( + candidate_tasks + + (sTiles[request_idx] + candidate - cutlass.Int32(1)) // candidate + ) + candidate_tasks = cute.arch.warp_redux_sync(candidate_tasks, "add") + candidate_makespan = cutlass.Int32(-1) + if candidate_tasks <= cutlass.min(task_capacity, cutlass.Int32(PLAN_SIM_MAX_TASKS)): + candidate_makespan = _plan_simulated_makespan( + sTiles, num_requests, candidate, wave_tasks, lane + ) + if lane == cutlass.Int32(0): + sCandidates[warp_idx] = candidate_makespan + sCandidates[PLAN_SIM_CANDIDATES + warp_idx] = candidate + cute.arch.sync_threads() + if simulate & (tidx == cutlass.Int32(0)): + best_makespan = cutlass.Int32(-1) + best_tiles = sScalars[0] + for candidate_idx in cutlass.range(0, PLAN_SIM_CANDIDATES, 1, unroll=1): + candidate_makespan = sCandidates[candidate_idx] + candidate = sCandidates[PLAN_SIM_CANDIDATES + candidate_idx] + if candidate_makespan >= cutlass.Int32(0): + better = ( + (best_makespan < cutlass.Int32(0)) + | (candidate_makespan < best_makespan) + | ((candidate_makespan == best_makespan) & (candidate > best_tiles)) + ) + if better: + best_makespan = candidate_makespan + best_tiles = candidate + sScalars[0] = best_tiles + cute.arch.sync_threads() + if tidx == cutlass.Int32(0): + tiles_per_task = sScalars[0] + task_begin = cutlass.Int32(0) + for request_idx in cutlass.range(0, num_requests, 1, unroll=1): + sBegins[request_idx] = task_begin + mPlan_i[cutlass.Int32(2) + request_idx] = task_begin + task_begin = ( + task_begin + + (sTiles[request_idx] + tiles_per_task - cutlass.Int32(1)) // tiles_per_task + ) + sBegins[num_requests] = task_begin + mPlan_i[cutlass.Int32(2) + num_requests] = task_begin + mPlan_i[0] = tiles_per_task + mPlan_i[1] = task_begin + cute.arch.sync_threads() + task_tiles = sScalars[0] + tables = cutlass.Int32(3) + num_requests + used_tasks = sBegins[num_requests] + # Grid order with wave_plan (plan_task_order): non-last tasks in request order from + # begins[r] - r, then last tasks longest first. Partials stay at begins[r] + split. + full_tasks = used_tasks - num_requests + for request_idx in cutlass.range(tidx, num_requests, plan_threads, unroll=1): + begin = sBegins[request_idx] + request_tasks = sBegins[request_idx + cutlass.Int32(1)] - begin + first_task = begin + last_task = begin + request_tasks - cutlass.Int32(1) + if cutlass.const_expr(wave_plan): + request_tiles = sTiles[request_idx] + tail = request_tiles - (request_tasks - cutlass.Int32(1)) * task_tiles + rank = cutlass.Int32(0) + for other_idx in cutlass.range(0, num_requests, 1, unroll=1): + other_tasks = sBegins[other_idx + cutlass.Int32(1)] - sBegins[other_idx] + other_tail = sTiles[other_idx] - (other_tasks - cutlass.Int32(1)) * task_tiles + if (other_tail > tail) | ((other_tail == tail) & (other_idx < request_idx)): + rank = rank + cutlass.Int32(1) + first_task = begin - request_idx + last_task = full_tasks + rank + for split_idx in cutlass.range(0, request_tasks, 1, unroll=1): + task_idx = last_task + if split_idx + cutlass.Int32(1) < request_tasks: + task_idx = first_task + split_idx + mPlan_i[tables + task_idx] = request_idx + mPlan_i[tables + task_capacity + task_idx] = split_idx + mPlan_i[tables + task_capacity + task_capacity + task_idx] = request_tasks + for task_idx in cutlass.range(used_tasks + tidx, task_capacity, plan_threads, unroll=1): + mPlan_i[tables + task_idx] = num_requests - cutlass.Int32(1) + mPlan_i[tables + task_capacity + task_idx] = cutlass.Int32(1 << 15) + mPlan_i[tables + task_capacity + task_capacity + task_idx] = cutlass.Int32(1 << 15) + + +@cute.jit +def packed_plan_host( + kv_lens_ptr: cute.Pointer, + helix_bounds_ptr: cute.Pointer, + prompt_lens_ptr: cute.Pointer, + corrected_kv_lens_ptr: cute.Pointer, + generation_lens_ptr: cute.Pointer, + plan_ptr: cute.Pointer, + stream: cuda.CUstream, + num_requests: cutlass.Int32, + task_capacity: cutlass.Int32, + wave_tasks: cutlass.Int32, + max_tiles_per_task: cutlass.Int32, + generation_len: cutlass.Int32, + m_tiles: cutlass.Int32, + wave_clusters: cutlass.Int32, + query_len_per_seq: cutlass.Constexpr, + use_helix_kv_bounds: cutlass.Constexpr, + write_generation_lengths: cutlass.Constexpr, + lpt_order: cutlass.Constexpr, + wave_plan: cutlass.Constexpr, +) -> None: + request_layout = cute.make_layout((num_requests,), stride=(1,)) + mValidK_l = cute.make_tensor(kv_lens_ptr, request_layout) + mHelixKvBounds_l = cute.make_tensor( + helix_bounds_ptr, cute.make_layout((num_requests * query_len_per_seq,), stride=(1,)) + ) + mPromptLens_l = cute.make_tensor(prompt_lens_ptr, request_layout) + mCorrectedKvLens_l = cute.make_tensor(corrected_kv_lens_ptr, request_layout) + mGenerationLens_l = cute.make_tensor(generation_lens_ptr, request_layout) + mPlan_i = cute.make_tensor( + plan_ptr, + cute.make_layout( + (cutlass.Int32(3) + num_requests + cutlass.Int32(3) * task_capacity,), + stride=(1,), + ), + ) + packed_plan_kernel( + mValidK_l, + mHelixKvBounds_l, + mPromptLens_l, + mCorrectedKvLens_l, + mGenerationLens_l, + mPlan_i, + num_requests, + task_capacity, + wave_tasks, + max_tiles_per_task, + generation_len, + m_tiles, + wave_clusters, + query_len_per_seq, + use_helix_kv_bounds, + write_generation_lengths, + lpt_order, + wave_plan, + ).launch( + grid=(1, 1, 1), + block=(_plan_threads(lpt_order, wave_plan), 1, 1), + stream=stream, + ) + + +def packed_plan_length(num_requests: int, task_capacity: int) -> int: + """Int32 entries of the plan buffer for ``num_requests`` and ``task_capacity``.""" + return 3 + num_requests + 3 * task_capacity + + +def launch_packed_plan( + kv_lens: torch.Tensor, + plan: torch.Tensor, + *, + num_requests: int, + num_splits: int, + wave_tasks: int, + max_tiles_per_task: int, + query_len: int, + helix_kv_bounds: torch.Tensor | None = None, + generation_lengths: tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None = None, + lpt_m_tiles: int = 0, + lpt_wave_clusters: int = 0, + wave_plan: bool = False, +) -> None: + """Write the balanced task plan for ``num_requests * num_splits`` slots into ``plan``. + + ``kv_lens`` (int32 per request) and the optional per-token int32 + ``helix_kv_bounds`` are the same bounds the decode kernel reads; see + ``plan_packed_tasks`` for the host mirror of the plan. With + ``generation_lengths = (prompt_lens, corrected_kv_lens, generation_lens, + generation_len)`` the same launch first writes the per-request generation + lengths ``kv_lens - prompt_lens + generation_len`` and ``generation_len`` + (``kv_lens`` then being the raw lengths) and plans on the corrected ones. + Positive ``lpt_m_tiles`` (decode clusters per task) and + ``lpt_wave_clusters`` (clusters the GPU runs at once) select the makespan + plan (``plan_packed_tasks_lpt``); otherwise ``wave_plan`` selects the wave + plan, and without it the plan keeps the rule's task length and request order (no + wave simulation, no tail reordering). The makespan and wave plans list the tasks + out of request order, so they need a decode kernel that stores each partial at + ``begins[request] + split`` (``SUPPORTS_ORDERED_PLAN``). + """ + task_capacity = num_requests * num_splits + lpt_order = lpt_m_tiles > 0 + if lpt_order and lpt_wave_clusters <= 0: + raise ValueError("the makespan plan needs the wave's cluster count") + if num_requests <= 0 or num_splits <= 0 or wave_tasks <= 0 or max_tiles_per_task <= 0: + raise ValueError("plan needs positive request, split, wave and tile counts") + if num_requests > PLAN_MAX_REQUESTS: + raise ValueError(f"plan supports at most {PLAN_MAX_REQUESTS} requests, got {num_requests}") + if query_len <= 0: + raise ValueError("query_len must be positive") + if kv_lens.dtype != torch.int32 or kv_lens.numel() < num_requests or kv_lens.stride(0) != 1: + raise ValueError(f"kv_lens must be contiguous int32 with >= {num_requests} entries") + if helix_kv_bounds is not None and ( + helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.numel() < num_requests * query_len + or helix_kv_bounds.stride(0) != 1 + ): + raise ValueError( + f"helix_kv_bounds must be contiguous int32 with >= {num_requests * query_len} entries" + ) + plan_len = packed_plan_length(num_requests, task_capacity) + if plan.dtype != torch.int32 or plan.numel() < plan_len or plan.stride(0) != 1: + raise ValueError(f"plan must be contiguous int32 with >= {plan_len} entries") + use_helix = helix_kv_bounds is not None + bounds_ptr = helix_kv_bounds.data_ptr() if use_helix else kv_lens.data_ptr() + write_lengths = generation_lengths is not None + length_ptrs = (kv_lens.data_ptr(),) * 3 + generation_len = 0 + if write_lengths: + prompt_lens, corrected_kv_lens, generation_lens, generation_len = generation_lengths + for name, tensor in ( + ("prompt_lens", prompt_lens), + ("corrected_kv_lens", corrected_kv_lens), + ("generation_lens", generation_lens), + ): + if ( + tensor.dtype != torch.int32 + or tensor.numel() < num_requests + or tensor.stride(0) != 1 + ): + raise ValueError(f"{name} must be contiguous int32 with >= {num_requests} entries") + if generation_len <= 0: + raise ValueError("generation_len must be positive") + length_ptrs = ( + prompt_lens.data_ptr(), + corrected_kv_lens.data_ptr(), + generation_lens.data_ptr(), + ) + gmem = cute.AddressSpace.gmem + ptrs = ( + make_ptr(cutlass.Int32, kv_lens.data_ptr(), gmem, assumed_align=4), + make_ptr(cutlass.Int32, bounds_ptr, gmem, assumed_align=4), + *(make_ptr(cutlass.Int32, ptr, gmem, assumed_align=4) for ptr in length_ptrs), + make_ptr(cutlass.Int32, plan.data_ptr(), gmem, assumed_align=4), + ) + stream = _current_cu_stream() + # The makespan plan replaces the wave plan. + wave_plan = wave_plan and not lpt_order + cache_key = (query_len, use_helix, write_lengths, lpt_order, wave_plan) + plan_fn = _PACKED_PLAN_COMPILE_CACHE.get(cache_key) + if plan_fn is None: + plan_fn = _compile_cutedsl( + packed_plan_host, + *ptrs, + stream, + *(cutlass.Int32(1) for _ in range(7)), + query_len, + use_helix, + write_lengths, + lpt_order, + wave_plan, + ) + _PACKED_PLAN_COMPILE_CACHE[cache_key] = plan_fn + plan_fn( + *ptrs, + stream, + cutlass.Int32(num_requests), + cutlass.Int32(task_capacity), + cutlass.Int32(wave_tasks), + cutlass.Int32(max_tiles_per_task), + cutlass.Int32(generation_len), + cutlass.Int32(max(lpt_m_tiles, 1)), + cutlass.Int32(max(lpt_wave_clusters, 1)), + ) + + +def _make_reduce_ptrs( + partial: torch.Tensor, + maxima: torch.Tensor, + sums: torch.Tensor, + output: torch.Tensor, + stats: torch.Tensor | None, + q_scale: torch.Tensor, + kv_scale: torch.Tensor, + plan: torch.Tensor, +) -> tuple[cute.Pointer, ...]: + stats_max_ptr = stats.data_ptr() if stats is not None else maxima.data_ptr() + stats_sum_ptr = ( + stats.data_ptr() + stats.element_size() if stats is not None else sums.data_ptr() + ) + gmem = cute.AddressSpace.gmem + partial_dtype = cutlass.Float16 if partial.dtype == torch.float16 else cutlass.Float32 + return ( + make_ptr(partial_dtype, partial.data_ptr(), gmem, assumed_align=32), + make_ptr(cutlass.Float32, maxima.data_ptr(), gmem, assumed_align=32), + make_ptr(cutlass.Float32, sums.data_ptr(), gmem, assumed_align=32), + make_ptr(_cutlass_output_dtype(output.dtype), output.data_ptr(), gmem, assumed_align=32), + make_ptr(cutlass.Float32, stats_max_ptr, gmem, assumed_align=4), + make_ptr(cutlass.Float32, stats_sum_ptr, gmem, assumed_align=4), + make_ptr(cutlass.Float32, q_scale.data_ptr(), gmem, assumed_align=4), + make_ptr(cutlass.Float32, kv_scale.data_ptr(), gmem, assumed_align=4), + make_ptr(cutlass.Int32, plan.data_ptr(), gmem, assumed_align=4), + ) + + +def plan_packed_tasks( + request_tiles: list[int], + task_capacity: int, + wave_tasks: int, + max_tiles_per_task: int = 1 << 30, + min_tiles_per_task: int = PACKED_MIN_TILES_PER_TASK, + one_wave: bool = False, + wave_plan: bool = False, +) -> tuple[int, list[int]]: + """Host mirror of the kernels' balanced task plan: (tiles per task, task begins). + + ``request_tiles`` are the per-request KV tile counts (>= 1), ``task_capacity`` + the grid.z slots (``num_requests * num_splits``) and ``wave_tasks`` the tasks + that fill one wave of clusters; ``one_wave`` (Helix) never takes the two-wave + target. With ``wave_plan`` the length competes + with the longest request cut into 1..PLAN_SIM_CANDIDATES - 1 pieces by + ``plan_simulated_makespan``. Returns the KV tiles each task may own and the + ``len(request_tiles) + 1`` task prefix, whose last element is the number of + tasks actually used (<= ``task_capacity``). See ``_packed_tiles_per_task``. + """ + num_requests = len(request_tiles) + if num_requests <= 0 or task_capacity <= 0 or wave_tasks <= 0 or min(request_tiles) <= 0: + raise ValueError("plan needs requests, task slots, wave tasks and positive tile counts") + total = sum(request_tiles) + longest = max(request_tiles) + tiles_per_task = longest + + def used(tiles: int) -> int: + return sum(math.ceil(t / tiles) for t in request_tiles) + + if task_capacity > num_requests: + share = max(math.ceil(total / wave_tasks), min_tiles_per_task) + if longest > share + share // 2: + full_machine = 3 * wave_tasks <= 4 * num_requests + strong_skew = longest * num_requests > 2 * total + two_waves = not one_wave and (full_machine or strong_skew) + target_tasks = 2 * wave_tasks if two_waves else wave_tasks + target = max(math.ceil(total / target_tasks), min_tiles_per_task) + pieces = (2 * longest + target) // (2 * target) + target = max(math.ceil(longest / pieces), min_tiles_per_task) + if used(target) <= task_capacity: + tiles_per_task = target + else: + fallback = max( + math.ceil(total / (task_capacity - num_requests)), min_tiles_per_task + ) + tiles_per_task = fallback + even = max(math.ceil(longest / (task_capacity // num_requests)), min_tiles_per_task) + if even < fallback and used(even) <= task_capacity: + tiles_per_task = even + tiles_per_task = min(tiles_per_task, longest) + tiles_per_task = min(tiles_per_task, max_tiles_per_task) + if ( + wave_plan + and task_capacity > num_requests + and wave_tasks <= PLAN_SIM_MAX_SLOTS + and num_requests <= PLAN_SIM_MAX_REQUESTS + ): + candidates = [tiles_per_task] + [ + min(max(math.ceil(longest / pieces), min_tiles_per_task), longest, max_tiles_per_task) + for pieces in range(1, PLAN_SIM_CANDIDATES) + ] + best = None + for candidate in candidates: + if used(candidate) > min(task_capacity, PLAN_SIM_MAX_TASKS): + continue + makespan = plan_simulated_makespan(request_tiles, candidate, wave_tasks) + if best is None or (makespan, -candidate) < best: + best = (makespan, -candidate) + if best is not None: + tiles_per_task = -best[1] + begins = [0] + for tiles in request_tiles: + begins.append(begins[-1] + math.ceil(tiles / tiles_per_task)) + return tiles_per_task, begins + + +def _lpt_candidates( + balanced: int, + request_tiles: list[int], + m_tiles: int, + slots: int, + max_tiles_per_task: int, + min_tiles_per_task: int = PACKED_MIN_TILES_PER_TASK, +) -> list[int]: + """Host mirror of ``_lpt_write_candidates``. + + The balanced length; the per-slot load of a perfect schedule (``ideal``) + and its 1/2, 1/3, 1/4 (pieces that backfill each other), each also 4% + longer; and the closed-form optimum of an even batch of the longest + request (``q`` full tasks per request plus remainders packed in rounds). + """ + cost = LPT_CLUSTER_COST_TILES + longest = max(request_tiles) + upper = max(min(longest, max_tiles_per_task), 1) + + def clamp(tiles: int) -> int: + return min(max(tiles, min_tiles_per_task), upper) + + streams = len(request_tiles) * m_tiles + ideal = (sum(request_tiles) * m_tiles + cost * (slots + streams) + slots - 1) // slots - cost + ideal = max(ideal, 1) + candidates = [balanced] + for k in range(1, 5): + candidates += [clamp(ideal // k), clamp(ideal * 26 // (25 * k))] + uniform_tiles, uniform_makespan = longest, (1 << 31) - 1 + for q in range(1, 16): + full_slots = q * streams + if full_slots >= slots: + continue + rounds = math.ceil(streams / (slots - full_slots)) + denom = 1 + rounds * q + tiles = (rounds * (longest + cost) - cost + denom - 1) // denom + tiles = max(min(tiles, longest), 1) + remainder = longest - q * tiles + makespan = tiles + cost + if remainder > 0: + makespan = max(makespan, rounds * (remainder + cost)) + if makespan < uniform_makespan: + uniform_tiles, uniform_makespan = tiles, makespan + candidates.append(clamp(uniform_tiles)) + return candidates + + +def lpt_makespan( + request_tiles: list[int], tiles_per_task: int, m_tiles: int, slots: int, task_capacity: int +) -> int: + """Host mirror of ``_lpt_simulate_makespan``: longest-first schedule length in tiles.""" + cost = LPT_CLUSTER_COST_TILES + full = sum(t // tiles_per_task for t in request_tiles) + remainders = sorted( + (t % tiles_per_task for t in request_tiles if t % tiles_per_task), reverse=True + ) + if full + len(remainders) > task_capacity: + return (1 << 31) - 1 + full_cost = tiles_per_task + cost + base = (full * m_tiles // slots) * full_cost + extra = full * m_tiles % slots + loads = [base + (full_cost if slot < extra else 0) for slot in range(slots)] + for remainder in remainders: + for _ in range(m_tiles): + slot = min(range(slots), key=lambda s: (loads[s], s)) + loads[slot] += remainder + cost + return max(loads) + + +def plan_packed_tasks_lpt( + request_tiles: list[int], + task_capacity: int, + wave_tasks: int, + m_tiles: int, + wave_clusters: int, + max_tiles_per_task: int = 1 << 30, + one_wave: bool = False, +) -> tuple[int, list[int], list[tuple[int, int]]]: + """Host mirror of the makespan plan (``lpt_order`` of ``packed_plan_kernel``). + + Returns the task length, the per-request partial-slot prefix (as + ``plan_packed_tasks``) and the launch order, ``order[z] = (request, + split)`` for the used tasks. Batches of at most ``LPT_SIM_MAX_REQUESTS`` + requests with spare slots take the candidate length + (``_lpt_candidates``) whose longest-first schedule (``lpt_makespan``) is + shortest, ties to the longer length, if it beats the balanced length by + the ``LPT_MIN_GAIN_DIVISOR`` / half-cluster margin; others keep the + balanced length. + """ + balanced, _ = plan_packed_tasks( + request_tiles, + task_capacity, + wave_tasks, + max_tiles_per_task, + one_wave=one_wave, + wave_plan=False, + ) + num_requests = len(request_tiles) + tiles_per_task = balanced + slots = min(wave_clusters, LPT_MAX_SLOTS) + if task_capacity > num_requests and num_requests <= LPT_SIM_MAX_REQUESTS: + candidates = _lpt_candidates(balanced, request_tiles, m_tiles, slots, max_tiles_per_task) + makespans = [ + lpt_makespan(request_tiles, tiles, m_tiles, slots, task_capacity) + for tiles in candidates + ] + best = makespans[0] + limit = best - max(best // LPT_MIN_GAIN_DIVISOR, LPT_CLUSTER_COST_TILES // 2) + for tiles, makespan in zip(candidates[1:], makespans[1:]): + if makespan <= limit and ( + makespan < best or (makespan == best and tiles > tiles_per_task) + ): + best, tiles_per_task = makespan, tiles + begins = [0] + for tiles in request_tiles: + begins.append(begins[-1] + math.ceil(tiles / tiles_per_task)) + full_order = [ + (request, split) + for request, tiles in enumerate(request_tiles) + for split in range(tiles // tiles_per_task) + ] + last_tasks = sorted( + (-(tiles % tiles_per_task), request) + for request, tiles in enumerate(request_tiles) + if tiles % tiles_per_task + ) + order = full_order + [ + (request, request_tiles[request] // tiles_per_task) for _, request in last_tasks + ] + return tiles_per_task, begins, order + + +def plan_task_order( + request_tiles: list[int], tiles_per_task: int, wave_plan: bool = False +) -> list[tuple[int, int]]: + """Grid order of the plan's tasks as (request, split). + + Every request's tasks but its last in request order, then the last tasks + longest first (ties: lower request); without ``wave_plan`` all in request order. + """ + tasks = [math.ceil(tiles / tiles_per_task) for tiles in request_tiles] + if not wave_plan: + return [(r, split) for r, n in enumerate(tasks) for split in range(n)] + order = [(r, split) for r, n in enumerate(tasks) for split in range(n - 1)] + tails = [ + (tiles - (n - 1) * tiles_per_task, r) + for r, (tiles, n) in enumerate(zip(request_tiles, tasks)) + ] + order += [(r, tasks[r] - 1) for _, r in sorted(tails, key=lambda tail: (-tail[0], tail[1]))] + return order + + +def plan_simulated_makespan(request_tiles: list[int], tiles_per_task: int, wave_tasks: int) -> int: + """Host mirror of ``_plan_simulated_makespan``: grid makespan in KV tiles for one length.""" + slots = [0] * wave_tasks + makespan = 0 + for request, split in plan_task_order(request_tiles, tiles_per_task, wave_plan=True): + start = heapq.heappop(slots) + tiles = min(tiles_per_task, request_tiles[request] - split * tiles_per_task) + end = start + tiles + PLAN_TASK_COST_TILES + if request_tiles[request] > tiles_per_task: + end += PLAN_SPLIT_COST_TILES + makespan = max(makespan, end) + heapq.heappush(slots, end) + return makespan + + +def request_kv_tiles( + kv_lens: list[int], query_len: int, helix_kv_bounds: list[int] | None = None +) -> list[int]: + """Per-request KV tile counts exactly as the kernels derive them.""" + tiles = [] + for request, kv_len in enumerate(kv_lens): + bound = ( + helix_kv_bounds[request * query_len + query_len - 1] + if helix_kv_bounds is not None + else kv_len + ) + tiles.append(max(math.ceil(max(bound, 0) / KV_TILE), 1)) + return tiles + + +def select_reduce_columns(num_rows: int, sm_count: int) -> int: + """Column band per warp: narrower bands when few rows would leave SMs idle.""" + if num_rows >= 2 * sm_count: + return 512 + if num_rows >= sm_count: + return 256 + return 128 + + +def select_reduce_ctas(num_items: int, sm_count: int) -> int: + """Persistent CTAs: one warp per (row, band) item, up to four CTAs per SM.""" + warps_per_cta = _REDUCE_THREADS // 32 + return max( + 1, min((num_items + warps_per_cta - 1) // warps_per_cta, _REDUCE_CTAS_PER_SM * sm_count) + ) + + +def reduce_packed_splits( + partial: torch.Tensor, + maxima: torch.Tensor, + sums: torch.Tensor, + output: torch.Tensor, + q_scale: torch.Tensor, + kv_scale: torch.Tensor, + *, + num_requests: int, + num_splits: int, + valid_m: int, + physical_m: int, + sm_scale: float, + output_scale: float, + plan: torch.Tensor, + stats: torch.Tensor | None = None, + sm_count: int = 128, + use_pdl: bool = False, +) -> None: + """Reduce ``[splits * requests, physical_m, 512]`` partials into ``output``. + + FP32 partials are max-relative numerators (``output_scale`` converts the + merge to output units); FP16 partials are normalized split outputs already + in output units (pass ``output_scale=1``). + + ``output`` is a contiguous ``[requests * valid_m, 512]`` BF16/FP16 tensor and + ``stats`` the optional contiguous ``[requests * valid_m, 2]`` FP32 Helix + (max, sum) buffer. ``q_scale``/``kv_scale`` are the device FP4 global scales + and ``plan`` the task plan written by ``launch_packed_plan`` for the same + ``splits * requests`` slots the decode kernel ran. ``use_pdl`` launches it as a + programmatic dependent of the preceding (decode) kernel. + """ + tasks = num_splits * num_requests + total_rows = num_requests * valid_m + if num_requests <= 0 or num_splits <= 0 or valid_m <= 0 or physical_m < valid_m: + raise ValueError("packed reduce needs positive requests/splits and physical_m >= valid_m") + if partial.dtype not in (torch.float32, torch.float16) or partial.shape != ( + tasks, + physical_m, + _VALUE_DIM, + ): + raise ValueError( + f"partial must be float32/float16 {(tasks, physical_m, _VALUE_DIM)}, " + f"got {partial.dtype} {tuple(partial.shape)}" + ) + if maxima.shape != (tasks, physical_m) or sums.shape != (tasks, physical_m): + raise ValueError("maxima and sums must be [tasks, physical_m]") + if maxima.dtype != torch.float32 or sums.dtype != torch.float32: + raise TypeError("maxima and sums must be float32") + if not (partial.is_contiguous() and maxima.is_contiguous() and sums.is_contiguous()): + raise ValueError("partial, maxima, and sums must be contiguous") + _validate_output_dtype(output.dtype) + if output.shape != (total_rows, _VALUE_DIM) or not output.is_contiguous(): + raise ValueError( + f"output must be contiguous [{total_rows}, {_VALUE_DIM}], got {tuple(output.shape)}" + ) + if stats is not None and ( + stats.dtype != torch.float32 or stats.shape != (total_rows, 2) or not stats.is_contiguous() + ): + raise ValueError(f"stats must be contiguous float32 [{total_rows}, 2]") + for name, tensor in (("q_scale", q_scale), ("kv_scale", kv_scale)): + if tensor.dtype != torch.float32 or tensor.numel() != 1: + raise TypeError(f"{name} must be a scalar float32 device tensor") + plan_len = packed_plan_length(num_requests, tasks) + if plan.dtype != torch.int32 or plan.numel() < plan_len or plan.stride(0) != 1: + raise ValueError(f"plan must be contiguous int32 with >= {plan_len} entries") + if not math.isfinite(sm_scale) or sm_scale <= 0.0 or not math.isfinite(output_scale): + raise ValueError("sm_scale must be finite and positive and output_scale finite") + for name, tensor in ( + ("partial", partial), + ("maxima", maxima), + ("sums", sums), + ("output", output), + ): + _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) + if stats is not None: + _validate_tensor_pointer_alignment("stats", stats, alignment_bytes=8) + deferred = DeferredPackedReduce.active() + if deferred is not None and deferred.args is None: + deferred.args = dict( + partial=partial, + maxima=maxima, + sums=sums, + output=output, + q_scale=q_scale, + kv_scale=kv_scale, + num_requests=num_requests, + num_splits=num_splits, + valid_m=valid_m, + physical_m=physical_m, + sm_scale=sm_scale, + output_scale=output_scale, + plan=plan, + stats=stats, + sm_count=sm_count, + ) + return + + band_columns = select_reduce_columns(total_rows, sm_count) + grid_ctas = select_reduce_ctas(total_rows * (_VALUE_DIM // band_columns), sm_count) + write_stats = stats is not None + # The Helix stats of single-task requests are spread over the whole reducer grid (otherwise + # every request's row chunks land on the same first CTAs). + spread_stats = write_stats + stream = _current_cu_stream() + ptrs = _make_reduce_ptrs(partial, maxima, sums, output, stats, q_scale, kv_scale, plan) + cache_key = ( + band_columns, + write_stats, + spread_stats, + output.dtype, + use_pdl, + partial.dtype, + ) + reduce_fn = _PACKED_REDUCE_COMPILE_CACHE.get(cache_key) + if reduce_fn is None: + reduce_fn = _compile_cutedsl( + packed_split_reduce_host, + *ptrs, + cutlass.Float32(1.0), + cutlass.Float32(1.0), + stream, + *(cutlass.Int32(1) for _ in range(7)), + band_columns, + write_stats, + spread_stats, + use_pdl, + ) + _PACKED_REDUCE_COMPILE_CACHE[cache_key] = reduce_fn + reduce_fn( + *ptrs, + cutlass.Float32(sm_scale * LOG2_E), + cutlass.Float32(output_scale), + stream, + cutlass.Int32(num_requests), + cutlass.Int32(num_splits), + cutlass.Int32(valid_m), + cutlass.Int32(physical_m), + cutlass.Int32(total_rows), + cutlass.Int32(2 if stats is not None else 1), + cutlass.Int32(grid_ctas), + ) + + +class DeferredPackedReduce: + """Context in which the first :func:`reduce_packed_splits` call records its arguments in + :attr:`args` instead of launching, for a consumer that folds the split partials itself (the + Kimi-K3 Helix tail, ``mla.py``). The partials, row state and plan stay + valid until the next decode on the stream overwrites them; :meth:`run` launches the recorded + reduction for a consumer that cannot fold them after all. + + Args: + release_early: With ``features.fp4_mla_reduce_pdl``, the decode releases its programmatic + dependents (the consumer) at its entry; False: only its completion does. + """ + + _stack: list["DeferredPackedReduce"] = [] + + def __init__(self, release_early: bool = True) -> None: + self.args: dict | None = None + self.release_early = release_early + + @classmethod + def active(cls) -> "DeferredPackedReduce | None": + """The innermost open context, or None.""" + return cls._stack[-1] if cls._stack else None + + def __enter__(self) -> "DeferredPackedReduce": + DeferredPackedReduce._stack.append(self) + return self + + def __exit__(self, *exc_info) -> None: + DeferredPackedReduce._stack.pop() + + def run(self) -> None: + """Launch the recorded reduction (no-op without a recorded call).""" + if self.args is None: + return + args = dict(self.args) + partial, maxima, sums, output = ( + args.pop(k) for k in ("partial", "maxima", "sums", "output") + ) + q_scale, kv_scale = args.pop("q_scale"), args.pop("kv_scale") + reduce_packed_splits(partial, maxima, sums, output, q_scale, kv_scale, **args) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp8_static_quant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp8_static_quant.py new file mode 100644 index 000000000000..cad89a910e9d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/fp8_static_quant.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Triton BF16 -> FP8 E4M3 per-tensor static quantization for the K3 FP8 context FMHA inputs. + +:func:`static_quantize_e4m3` writes the bytes of +``torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor``: every element becomes +``cvt.rn.satfinite.e4m3(float(x) * (1 / scale))`` with an IEEE FP32 reciprocal, and each program +converts a large contiguous block of BF16 pairs. :func:`static_quantize_e4m3_rows` writes the same +bytes for a row-strided view (e.g. the RoPE columns of ``latent_cache``) without a contiguous copy. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _e4m3x2_from_bf16x2(pairs, inv_scale): + """Convert int32-packed BF16 pairs to int16-packed E4M3 pairs, first element in the low byte.""" + lo = (pairs << 16).to(tl.float32, bitcast=True) * inv_scale + hi = (pairs & -65536).to(tl.float32, bitcast=True) * inv_scale + return tl.inline_asm_elementwise( + "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", + "=h,f,f", + [lo, hi], + dtype=tl.int16, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _static_quant_e4m3_kernel(x_ptr, y_ptr, scale_ptr, num_pairs, BLOCK: tl.constexpr): + offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < num_pairs + inv_scale = tl.math.div_rn(1.0, tl.load(scale_ptr).to(tl.float32)) + pairs = tl.load(x_ptr + offsets, mask=mask, other=0) + tl.store(y_ptr + offsets, _e4m3x2_from_bf16x2(pairs, inv_scale), mask=mask) + + +@triton.jit +def _static_quant_e4m3_rows_kernel( + x_ptr, + y_ptr, + scale_ptr, + num_rows, + x_row_stride, + ROW_PAIRS: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + rows = tl.program_id(0).to(tl.int64) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + cols = tl.arange(0, ROW_PAIRS) + mask = (rows < num_rows)[:, None] + inv_scale = tl.math.div_rn(1.0, tl.load(scale_ptr).to(tl.float32)) + pairs = tl.load(x_ptr + rows[:, None] * x_row_stride + cols[None, :], mask=mask, other=0) + tl.store( + y_ptr + rows[:, None] * ROW_PAIRS + cols[None, :], + _e4m3x2_from_bf16x2(pairs, inv_scale), + mask=mask, + ) + + +def _check_scale(scale: torch.Tensor, device: torch.device) -> None: + if scale.dtype != torch.float32 or scale.numel() != 1 or scale.device != device: + raise ValueError( + "FP8 static quantization requires a one-element FP32 scale on the input device." + ) + + +def static_quantize_e4m3(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Quantize a contiguous BF16 tensor with an even element count to FP8 E4M3 (same shape).""" + if x.dtype != torch.bfloat16 or not x.is_contiguous() or x.numel() % 2 != 0: + raise ValueError("FP8 static quantization requires a contiguous BF16 tensor of even size.") + _check_scale(scale, x.device) + y = torch.empty(x.shape, dtype=torch.float8_e4m3fn, device=x.device) + num_pairs = x.numel() // 2 + if num_pairs == 0: + return y + block = 8192 + _static_quant_e4m3_kernel[(triton.cdiv(num_pairs, block),)]( + x.view(-1).view(torch.int32), + y.view(-1).view(torch.int16), + scale, + num_pairs, + BLOCK=block, + num_warps=8, + ) + return y + + +def supports_rows(x: torch.Tensor) -> bool: + """Whether :func:`static_quantize_e4m3_rows` accepts ``x``: a 2-D BF16 view with unit column + stride, even row stride and storage offset, and a power-of-two width of at least 2.""" + if x.dtype != torch.bfloat16 or x.dim() != 2: + return False + width = x.shape[1] + return ( + x.stride(1) == 1 + and x.stride(0) % 2 == 0 + and x.storage_offset() % 2 == 0 + and width >= 2 + and width & (width - 1) == 0 + ) + + +def static_quantize_e4m3_rows(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Quantize a row-strided BF16 [rows, width] view (see :func:`supports_rows`) into a + contiguous FP8 E4M3 [rows, width] tensor.""" + if not supports_rows(x): + raise ValueError( + "FP8 row quantization requires unit column stride, even row stride and offset, " + f"and a power-of-two width; got shape {tuple(x.shape)}, stride {x.stride()}." + ) + _check_scale(scale, x.device) + rows, width = x.shape + y = torch.empty((rows, width), dtype=torch.float8_e4m3fn, device=x.device) + if rows == 0: + return y + pairs = x.view(torch.int32) + block_rows = 64 + _static_quant_e4m3_rows_kernel[(triton.cdiv(rows, block_rows),)]( + pairs, + y.view(torch.int16), + scale, + rows, + pairs.stride(0), + ROW_PAIRS=width // 2, + BLOCK_ROWS=block_rows, + num_warps=4, + ) + return y diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/helix_a2a_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/helix_a2a_fused.py new file mode 100644 index 000000000000..cf1a7b316dcb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/helix_a2a_fused.py @@ -0,0 +1,879 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused Helix context-parallel all-to-all and combine (SM100 / SM103). + +Helix runs the MLA decode attention of every head over a ``1 / P`` slice of the KV cache +(``P`` = CP size). Each rank then needs, for its ``H`` heads, the partial outputs and softmax +statistics of all ``P`` ranks and combines them with a log-sum-exp weighted sum. One launch of +:class:`HelixA2AFusedKernel` does the exchange and the combine: + +* Sender CTAs stage chunks of ``(token, destination rank)`` records of this rank in shared memory + -- the BF16 partial output of the destination's ``H`` heads (``H * D`` values) followed by + their FP32 ``(max, sum)`` pairs -- and send each record with one bulk copy (TMA) straight into + slot ``(token, source rank)`` of the destination's receive buffer (MNNVL fabric memory). + Rows this rank owns no KV for (``zero_kv_mask``) are replaced on the way by the no-op + contribution (zeros, ``(-inf, 0)``). The bulk copies keep the remote stores out of the SM's + load/store pipeline, which the receivers on the same SMs poll through. +* Receiver CTAs own tokens; one thread owns one 16-byte column (8 output values) of one head. + Lane ``l`` of a 16-lane head group polls the statistics of source rank ``l``; the group + computes the combine weights with the shuffles of the reference kernel; every lane then polls + its column's ``P`` vectors, accumulates them with FMAs in source order, rounds to BF16, stores + the output and re-arms the words it read. + +Receive buffers are Lamport buffers: every word starts as the sentinel ``0xffffffff``, a NaN +pattern the sender never stores (it rewrites such words to ``0x7fffffff``, also NaN for BF16 +pairs and FP32 alike, and NaN inputs combine to the same canonical NaN). A word that is not the +sentinel has therefore arrived; no flag, fence or counter sits between a sender's stores and +the receiver's loads. The buffer has two halves used by alternate launches. A launch reads its +half and re-arms what it read before it exits. The next stores into that half come two launches +later, from a rank that has finished the launch in between, which needed this rank's records of +that launch, which this rank sends only after this launch has finished. + +The half is the parity of a launch counter: every CTA takes a ticket from a packed 64-bit word +(high half: launch number, low half: CTAs of the launch that have arrived; the last arrival +advances the launch number), so the grid can vary between launches and the kernel replays +inside CUDA graphs. Launches that share a workspace must be ordered (one stream). A CTA takes +its ticket after ``griddepcontrol.wait``, when the previous grid has completed, and only then +lets the next grid launch, so the tickets of consecutive launches never interleave (a dependent +grid can start before every CTA of a grid larger than the GPU has started). + +The output is bitwise equal to ``helixAllToAllKernel`` followed by +``helix_postprocess_kernel_native_v2``: same ``expf`` and reciprocal instruction sequences, +same max / sum butterflies, FMA accumulation in source order, round-to-nearest BF16. That holds +for ``P <= 8`` and ``P = 16`` (other sizes pad the reference's source loop with zero terms). +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +#: The empty-word sentinel of the receive buffers (and its replacement in sent data). +SENTINEL = -1 +#: FP32 bits of -inf (the max of a no-op contribution). +_NEG_INF_BITS = -8388608 # 0xff800000 +#: Largest supported CP size (one 16-lane shuffle group holds the statistics of all sources). +MAX_CP = 16 +#: Records per sender chunk (one bulk copy each) and shared-memory stages per sender CTA. +_CHUNK_RECORDS = 8 +_STAGES = 3 +#: Threads per CTA (receivers: whole tokens of ``H * D / 8`` threads each). +_TARGET_THREADS = 256 + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _lv(value, *, loc=None, ip=None): + return Int64(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return Float32(value).ir_value(loc=loc, ip=ip) + + +def _unpack(res, ty, cls, n, *, loc=None, ip=None): + return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n)) + + +@dsl_user_op +def _atom_add_u64(addr, value, *, loc=None, ip=None): + return Int64( + _asm( + T.i64(), + [_lv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "atom.relaxed.gpu.global.add.u64 $0, [$1], $2;", + "=l,l,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sts_b64(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "st.shared.b64 [$0], $1;", + "r,l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_b64(addr, *, loc=None, ip=None): + return Int64( + _asm( + T.i64(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.b64 $0, [$1];", "=l,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _ld_nc_v4(addr, *, loc=None, ip=None): + """16 bytes of a tensor the kernel does not write.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _ld_vol_v4(addr, *, loc=None, ip=None): + """16 bytes of the receive buffer, read past L1 (peers write it over NVLink).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.volatile.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _ld_vol_v2(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_lv(addr, loc=loc, ip=ip)], + "ld.volatile.global.v2.b32 {$0, $1}, [$2];", + "=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _ld_nc_u8(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.u8 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_sentinel_v4(addr, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)], + "{\n.reg .b32 t;\nmov.b32 t, -1;\nst.global.v4.b32 [$0], {t, t, t, t};\n}", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_v4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _all_valid2(w0, w1, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(w0, loc=loc, ip=ip), _iv(w1, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.b32 p, $1, -1;\nsetp.ne.and.b32 p, $2, -1, p;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _prep4(w0, w1, w2, w3, zero_kv, is_stats, *, loc=None, ip=None): + """Words of one 16-byte vector as sent: the no-op contribution on zero-KV rows (zeros, or + ``(-inf, 0, -inf, 0)`` for a statistics vector), and no word equal to the sentinel.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)] + + [_iv(zero_kv, loc=loc, ip=ip), _iv(is_stats, loc=loc, ip=ip)], + "{\n.reg .pred pz, ps, p;\n.reg .b32 m, x;\n" + "setp.ne.u32 pz, $8, 0;\nsetp.ne.u32 ps, $9, 0;\n" + "selp.b32 m, 0xff800000, 0, ps;\n" + "selp.b32 x, m, $4, pz;\nsetp.eq.b32 p, x, -1;\nselp.b32 $0, 0x7fffffff, x, p;\n" + "selp.b32 x, 0, $5, pz;\nsetp.eq.b32 p, x, -1;\nselp.b32 $1, 0x7fffffff, x, p;\n" + "selp.b32 x, m, $6, pz;\nsetp.eq.b32 p, x, -1;\nselp.b32 $2, 0x7fffffff, x, p;\n" + "selp.b32 x, 0, $7, pz;\nsetp.eq.b32 p, x, -1;\nselp.b32 $3, 0x7fffffff, x, p;\n}", + "=r,=r,=r,=r,r,r,r,r,r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _all_valid4(w0, w1, w2, w3, *, loc=None, ip=None): + """1 if none of the four words is the sentinel.""" + return Int32( + _asm( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)], + "{\n.reg .pred p;\n" + "setp.ne.b32 p, $1, -1;\nsetp.ne.and.b32 p, $2, -1, p;\n" + "setp.ne.and.b32 p, $3, -1, p;\nsetp.ne.and.b32 p, $4, -1, p;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_sentinel_v2(addr, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)], + "{\n.reg .b32 t;\nmov.b32 t, -1;\nst.global.v2.b32 [$0], {t, t};\n}", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _as_f32(bits, *, loc=None, ip=None): + return Float32( + _asm(T.f32(), [_iv(bits, loc=loc, ip=ip)], "mov.b32 $0, $1;", "=f,r", loc=loc, ip=ip) + ) + + +@dsl_user_op +def _fmax(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "max.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fadd(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fmul(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """IEEE reciprocal (what ``1.F / x`` compiles to without fast math).""" + return Float32( + _asm(T.f32(), [_fv(a, loc=loc, ip=ip)], "rcp.rn.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def _corrected_exp(mx, gmax, total, *, loc=None, ip=None): + """``total * expf(mx - gmax)`` with CUDA's (non fast-math) ``expf`` instruction sequence.""" + return Float32( + _asm( + T.f32(), + [_fv(mx, loc=loc, ip=ip), _fv(gmax, loc=loc, ip=ip), _fv(total, loc=loc, ip=ip)], + "{\n.reg .f32 d, t, j, f, e, s;\n.reg .b32 b;\n" + "sub.rn.f32 d, $1, $2;\n" + "fma.rn.sat.f32 t, d, 0f3BBB989D, 0f3F000000;\n" + "fma.rm.f32 t, t, 0f437C0000, 0f4B400001;\n" + "add.rn.f32 j, t, 0fCB40007F;\n" + "neg.f32 j, j;\n" + "fma.rn.f32 f, d, 0f3FB8AA3B, j;\n" + "fma.rn.f32 f, d, 0f32A57060, f;\n" + "ex2.approx.ftz.f32 e, f;\n" + "mov.b32 b, t;\nshl.b32 b, b, 23;\nmov.b32 s, b;\n" + "mul.rn.f32 e, s, e;\n" + "mul.rn.f32 $0, e, $3;\n}", + "=f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_bfly(value, offset: int, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(value, loc=loc, ip=ip)], + f"shfl.sync.bfly.b32 $0, $1, {offset}, 31, -1;", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_idx(value, lane, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(value, loc=loc, ip=ip), _iv(lane, loc=loc, ip=ip)], + "shfl.sync.idx.b32 $0, $1, $2, 31, -1;", + "=f,f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bf16x2_fma(word, scale, acc_lo, acc_hi, *, loc=None, ip=None): + """``(fma(lo, scale, acc_lo), fma(hi, scale, acc_hi))`` of the BF16 pair in ``word``.""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [ + _iv(word, loc=loc, ip=ip), + _fv(scale, loc=loc, ip=ip), + _fv(acc_lo, loc=loc, ip=ip), + _fv(acc_hi, loc=loc, ip=ip), + ], + "{\n.reg .b32 l, h;\n.reg .f32 fl, fh;\n" + "shl.b32 l, $2, 16;\nand.b32 h, $2, 0xffff0000;\n" + "mov.b32 fl, l;\nmov.b32 fh, h;\n" + "fma.rn.f32 $0, fl, $3, $4;\nfma.rn.f32 $1, fh, $3, $5;\n}", + "=f,=f,r,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _pack_bf16x2(lo, hi, *, loc=None, ip=None): + """Round two FP32 values to BF16 (nearest even), ``lo`` in the low half.""" + return Int32( + _asm( + T.i32(), + [_fv(lo, loc=loc, ip=ip), _fv(hi, loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sts_v4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_proxy_async_shared(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.shared::cta;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _bulk_s2g(dst, src, nbytes: int, *, loc=None, ip=None): + """Bulk copy (TMA) of ``nbytes`` from shared ``src`` to global ``dst`` (local or peer).""" + _asm( + None, + [_lv(dst, loc=loc, ip=ip), _iv(src, loc=loc, ip=ip)], + f"cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], {nbytes};", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_commit(*, loc=None, ip=None): + _asm(None, [], "cp.async.bulk.commit_group;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _bulk_wait_read(n: int, *, loc=None, ip=None): + _asm(None, [], f"cp.async.bulk.wait_group.read {n};", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.wait;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.launch_dependents;", "", loc=loc, ip=ip) + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +class HelixA2AFusedKernel: + """Exchange + combine of one Helix MLA layer (see the module docstring). + + Args: + cp_size: ``P`` (at most 8, or 16). + num_heads: ``H``, heads per rank after the exchange (even). + head_dim: ``D``, values per head (128 or 256). + has_mask: Apply a per-token ``zero_kv_mask`` (uint8) on the sender side. + use_pdl: Launch with programmatic dependent launch. + """ + + def __init__(self, cp_size: int, num_heads: int, head_dim: int, has_mask: bool, use_pdl: bool): + if not supports(cp_size, num_heads, head_dim): + raise ValueError(f"cp_size={cp_size}, num_heads={num_heads}, head_dim={head_dim}") + self.P = cp_size + self.has_mask = has_mask + self.use_pdl = use_pdl + # 16-byte vectors of a record: partial output, then statistics. + self.DV = num_heads * head_dim // 8 + self.SV = num_heads // 2 + self.RV = self.DV + self.SV + # Columns (16-byte vectors) per head; a 16-lane group of a warp is one head. + self.CPH = head_dim // 8 + self.tok_threads = _ceil_div(self.DV, 32) * 32 + self.tpb = max(1, _TARGET_THREADS // self.tok_threads) + self.threads = self.tpb * self.tok_threads + + @cute.jit + def __call__( + self, + o_addr: Int64, + s_addr: Int64, + mask_addr: Int64, + out_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + rank_stride: Int64, + half_bytes: Int64, + num_tokens: Int32, + rank: Int32, + n_send: Int32, + n_recv: Int32, + stream: cuda.CUstream, + ): + @cute.struct + class SharedStorage: + half: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 2], 16] + stage: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, _STAGES * _CHUNK_RECORDS * self.RV * 4], 128 + ] + + self.shared_storage = SharedStorage + self.kernel( + o_addr, + s_addr, + mask_addr, + out_addr, + ws_addr, + state_addr, + rank_stride, + half_bytes, + num_tokens, + rank, + n_send, + n_recv, + ).launch( + grid=(n_send + n_recv, 1, 1), + block=(self.threads, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + o_addr: Int64, + s_addr: Int64, + mask_addr: Int64, + out_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + rank_stride: Int64, + half_bytes: Int64, + num_tokens: Int32, + rank: Int32, + n_send: Int32, + n_recv: Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + storage = utils.SmemAllocator().allocate(self.shared_storage) + half_slot = storage.half.data_ptr().toint() + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + # The ticket is taken after the wait for the previous grid: the previous launch of this + # kernel has then completed, so all of its CTAs have their tickets (a dependent grid may + # start before every CTA of a grid larger than the GPU has started). The senders' first + # loads overlap the atomic; only thread 0 needs its result before the first bulk copy. + if tidx == 0: + n_ctas = Int64(n_send + n_recv) + ticket = _atom_add_u64(state_addr, Int64(1)) + if (ticket & Int64(0xFFFFFFFF)) == n_ctas - Int64(1): + # Last arrival: clear the arrivals and advance the launch number in one add. + _atom_add_u64(state_addr, Int64(1 << 32) - n_ctas) + _sts_b64(half_slot, ((ticket >> 32) & Int64(1)) * half_bytes) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + if bidx < n_send: + # Only thread 0, which took the ticket, issues the bulk copies and needs the half. + self.send( + tidx, bidx, o_addr, s_addr, mask_addr, ws_addr, half_slot, rank_stride, + num_tokens, rank, n_send, storage.stage.data_ptr().toint(), + ) # fmt: skip + else: + cute.arch.barrier() + self.receive( + tidx, + bidx - n_send, + out_addr, + ws_addr + Int64(rank) * rank_stride + _lds_b64(half_slot), + num_tokens, + n_recv, + ) + + @cute.jit + def send( + self, + tidx: Int32, + bidx: Int32, + o_addr: Int64, + s_addr: Int64, + mask_addr: Int64, + ws_addr: Int64, + half_slot: Int32, + rank_stride: Int64, + num_tokens: Int32, + rank: Int32, + n_send: Int32, + stage_base: Int32, + ): + """Stage chunks of records in shared memory; thread 0 sends each record with one bulk + copy (record ``t * P + d`` of this rank goes to slot ``t * P + rank`` of rank ``d``).""" + P, RV, DV, SV = self.P, self.RV, self.DV, self.SV + R, NS = _CHUNK_RECORDS, _STAGES + VPC = R * RV + UL = _ceil_div(VPC, self.threads) + n_rec = num_tokens * P + n_chunks = (n_rec + (R - 1)) // R + it = Int32(0) + for c in cutlass.range(bidx, n_chunks, n_send, unroll=1): + sbase = stage_base + (it % NS) * (VPC * 16) + if tidx == 0: + # The stage's previous bulk copies (NS chunks ago) have read shared memory. + _bulk_wait_read(NS - 1) + cute.arch.barrier() + rec0 = c * R + # All loads of the chunk (data and zero-KV flags) are in flight together. + words = cute.make_rmem_tensor((UL * 4,), Int32) + zero_kv = cute.make_rmem_tensor((UL,), Int32) + for u in cutlass.range_constexpr(UL): + j = tidx + u * self.threads + if j < VPC: + rl = j // RV + k = j - rl * RV + rec = rec0 + rl + if rec < n_rec: + src = o_addr + Int64(rec * DV + k) * 16 + if k >= DV: + src = s_addr + Int64(rec * SV + k - DV) * 16 + w = _ld_nc_v4(src) + for e in cutlass.range_constexpr(4): + words[u * 4 + e] = w[e] + if cutlass.const_expr(self.has_mask): + zero_kv[u] = _ld_nc_u8(mask_addr + Int64(rec // P)) + for u in cutlass.range_constexpr(UL): + j = tidx + u * self.threads + if j < VPC: + rl = j // RV + k = j - rl * RV + if rec0 + rl < n_rec: + zk = Int32(0) + if cutlass.const_expr(self.has_mask): + zk = zero_kv[u] + w = _prep4( + words[u * 4], + words[u * 4 + 1], + words[u * 4 + 2], + words[u * 4 + 3], + zk, + Int32(k >= DV), + ) + _sts_v4(sbase + j * 16, *w) + _fence_proxy_async_shared() + cute.arch.barrier() + if tidx == 0: + dst_base = ws_addr + _lds_b64(half_slot) + for r in cutlass.range_constexpr(R): + rec = rec0 + r + if rec < n_rec: + t = rec // P + dst = rec - t * P + _bulk_s2g( + dst_base + Int64(dst) * rank_stride + Int64((t * P + rank) * RV) * 16, + sbase + r * (RV * 16), + RV * 16, + ) + _bulk_commit() + it = it + 1 + if tidx == 0: + _bulk_wait_read(0) + + @cute.jit + def receive( + self, + tidx: Int32, + rcta: Int32, + out_addr: Int64, + recv_base: Int64, + num_tokens: Int32, + n_recv: Int32, + ): + """Combine the records of this CTA's tokens (see the module docstring).""" + P, RV, DV, CPH = self.P, self.RV, self.DV, self.CPH + slot_t = tidx // self.tok_threads + c = tidx - slot_t * self.tok_threads + lane = tidx % 32 + group = lane - lane % 16 + h = c // CPH + col = c - h * CPH + src_l = col % 16 + valid = Int32(c < DV) + has_stats = valid * Int32(src_l < P) + # With 32 columns per head both 16-lane groups of a warp read the statistics; the first + # re-arms them (after the shuffles, so after both have read them). + owns_stats = has_stats * Int32(col < 16) + for t in cutlass.range(rcta * self.tpb + slot_t, num_tokens, n_recv * self.tpb, unroll=1): + rec0 = recv_base + Int64(t * (P * RV)) * 16 + st_addr = rec0 + Int64(src_l * RV + DV) * 16 + Int64(h * 8) + d_addr = rec0 + Int64(c) * 16 + stats = cute.make_rmem_tensor((2,), Int32) + stats[0] = Int32(_NEG_INF_BITS) + stats[1] = Int32(0) + vals = cute.make_rmem_tensor((P * 4,), Int32) + if valid != 0: + # Poll the statistics word and the P vectors of this column (all loads of a + # round in flight together) until no word holds the sentinel. + ready = Int32(0) + while ready == 0: + ready = Int32(1) + if has_stats != 0: + sv = _ld_vol_v2(st_addr) + stats[0] = sv[0] + stats[1] = sv[1] + ready = _all_valid2(sv[0], sv[1]) + for s in cutlass.range_constexpr(P): + w = _ld_vol_v4(d_addr + Int64(s * RV * 16)) + for e in cutlass.range_constexpr(4): + vals[s * 4 + e] = w[e] + ready = ready & _all_valid4(*w) + + # Combine weights of the group's head, in the reference kernel's order: lane l holds + # source l; max and sum butterflies over the 16 lanes (the reference's other 16 + # lanes hold -inf / +0 and change nothing). + mx = _as_f32(stats[0]) + gmax = mx + for off in (1, 2, 4, 8): + gmax = _fmax(gmax, _shfl_bfly(gmax, off)) + corrected = _corrected_exp(mx, gmax, _as_f32(stats[1])) + total = corrected + for off in (1, 2, 4, 8): + total = _fadd(total, _shfl_bfly(total, off)) + weight = _fmul(corrected, _rcp(total)) + + acc = [Float32(0.0)] * 8 + for s in cutlass.range_constexpr(P): + ws = _shfl_idx(weight, group + s) + for q in cutlass.range_constexpr(4): + acc[2 * q], acc[2 * q + 1] = _bf16x2_fma( + vals[s * 4 + q], ws, acc[2 * q], acc[2 * q + 1] + ) + if valid != 0: + o = [_pack_bf16x2(acc[2 * q], acc[2 * q + 1]) for q in range(4)] + _st_v4(out_addr + Int64(t * DV + c) * 16, *o) + for s in cutlass.range_constexpr(P): + _st_sentinel_v4(d_addr + Int64(s * RV * 16)) + if owns_stats != 0: + _st_sentinel_v2(st_addr) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def supports(cp_size: int, num_heads: int, head_dim: int) -> bool: + """Whether :func:`helix_a2a_fused` handles this shape bitwise like the reference path.""" + return ( + (1 <= cp_size <= 8 or cp_size == MAX_CP) and num_heads % 2 == 0 and head_dim in (128, 256) + ) + + +def record_bytes(num_heads: int, head_dim: int) -> int: + """Bytes of one ``(token, source rank)`` record in a receive buffer.""" + return num_heads * head_dim * 2 + num_heads * 8 + + +@functools.cache +def _compiled(cp_size: int, num_heads: int, head_dim: int, has_mask: bool, use_pdl: bool): + kernel = HelixA2AFusedKernel(cp_size, num_heads, head_dim, has_mask, use_pdl) + addr = Int64(1 << 20) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled = cute.compile( + kernel, + addr, + addr, + addr, + addr, + addr, + addr, + Int64(1 << 29), + Int64(1 << 24), + Int32(1), + Int32(0), + Int32(1), + Int32(1), + stream, + ) + return kernel, compiled + + +def helix_a2a_fused( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + zero_kv_mask: torch.Tensor | None, + out: torch.Tensor, + *, + cp_rank: int, + cp_size: int, + num_heads: int, + head_dim: int, + ws_addr: int, + rank_stride: int, + half_bytes: int, + state: torch.Tensor, + max_send_ctas: int, + max_recv_ctas: int, + use_pdl: bool = True, +) -> None: + """Launch the exchange + combine of ``num_tokens = partial_o.shape[0]`` tokens into ``out``. + + ``partial_o``: BF16 ``[T, P * H * D]`` (destination-rank major), ``softmax_stats``: FP32 + ``[T, P * H * 2]``, ``zero_kv_mask``: bool ``[T]`` or None, ``out``: BF16 ``[T, H * D]``. + ``ws_addr`` is rank 0's receive buffer in this process's address space, rank ``r``'s at + ``ws_addr + r * rank_stride``; each holds two halves of ``half_bytes``, every word armed + with the sentinel. ``state`` is the int64 launch counter shared by all launches on the + workspace. ``T`` records of ``record_bytes(H, D)`` bytes per source must fit a half. + """ + num_tokens = partial_o.shape[0] + if num_tokens == 0: + return + kernel, compiled = _compiled(cp_size, num_heads, head_dim, zero_kv_mask is not None, use_pdl) + # Up to max_send_ctas senders once there are that many chunks; fewer, larger senders below + # (half the cap) send small messages faster. + n_chunks = _ceil_div(num_tokens * cp_size, _CHUNK_RECORDS) + n_send = ( + max_send_ctas if n_chunks >= max_send_ctas else min(n_chunks, max(1, max_send_ctas // 2)) + ) + n_recv = min(max_recv_ctas, _ceil_div(num_tokens, kernel.tpb)) + mask_addr = zero_kv_mask.data_ptr() if zero_kv_mask is not None else 0 + compiled( + Int64(partial_o.data_ptr()), + Int64(softmax_stats.data_ptr()), + Int64(mask_addr), + Int64(out.data_ptr()), + Int64(ws_addr), + Int64(state.data_ptr()), + Int64(rank_stride), + Int64(half_bytes), + Int32(num_tokens), + Int32(cp_rank), + Int32(n_send), + Int32(n_recv), + cuda.CUstream(torch.cuda.current_stream(partial_o.device).cuda_stream), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_attn_res.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_attn_res.py new file mode 100644 index 000000000000..a35b1960653b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_attn_res.py @@ -0,0 +1,1706 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 decode attention-residual mix + RMSNorm (+ residual add) with a PDL prologue. + +Runs on SM100, SM103 and SM107: the kernels use no tensor-core instructions, at most five named +barriers per CTA and portable cluster sizes. + +For ``M`` tokens of hidden size 7168, ``N - 1`` snapshot rows ``block[n]`` and the layer row +``x``, one launch computes, per token, + +* ``rsigma_n = rsqrt(mean(v_n^2) + eps)``, ``logit_n = rsigma_n * ``, +* ``p = softmax(logit)``, ``mixed = bf16(sum_n p_n v_n)``, +* ``out = bf16(bf16(mixed * rsqrt(mean(mixed^2) + out_eps)) * out_w)``, + +where ``x`` is the given row, ``bf16(prefix + addend)`` (``trtllm::attn_res_add_rmsnorm_fwd``'s +fused add), or ``bf16(residual + bf16(routed + shared))`` (the MoE residual update of the +previous layer, ``kimi_k3_moe_residual._fused_moe_residual_add``); the updated prefix ``x`` is +stored as well in the last two cases, bit-identical to those adds. + +The kernel is launched with programmatic dependent launch and does everything that does not +depend on the previous kernel before ``griddepcontrol.wait``: it loads the weights and the +snapshot rows (and the prefix sum, which is older than the previous kernel) and reduces the +snapshots' statistics over the cluster. After the wait it only loads the dependent row, reduces +its two statistics, finishes the softmax and the mix, reduces the output sum of squares and +stores. Cluster partials travel by ``st.async`` into every peer's shared memory and are +counted by a per-round mbarrier (one one-way message per round). + +:class:`K3AttnResFastKernel` runs that schedule with a cluster of ``c`` CTAs per token (see +:func:`fast_cluster_size`), one 16-byte vector per thread. It sums the statistics in another order +than the C++ kernels, so the normalized output can differ from theirs by one BF16 rounding step in +a small fraction of the elements; the snapshots' share of the softmax and of the mixture is also +folded before the wait. + +Snapshot rows and the prefix sum may be read before the wait because they are older than the +previous kernel: the bank rows are written by plain copy kernels (no early trigger) at least one +layer earlier, and the prefix sum before the previous kernel's own predecessor, which waits for +its predecessor before it triggers its dependents. + +The fast kernel can take the addend from an MNNVL Lamport workspace (``lamport_addend`` with +``lamport_finish``): the all-reduce whose input every rank pushed, which the kernel then finishes +itself (rank-order sum of the slots, rotation state, re-arm of the previous buffer; +:mod:`k3_lamport_ar`). With +``early_poll`` it finishes that all-reduce and computes the mix before its ``griddepcontrol.wait`` +and only stores after it (see :func:`run_k3_attn_res`). +""" + +from __future__ import annotations + +import functools +import os + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from ..features import active as _features +from . import k3_fp8w_lowm_mxfp8 as mx_ops +from . import k3_lamport_ar as _lam +from . import k3_launch_attrs + +_F = _features() + +#: Hidden size the kernel is specialized for. +HIDDEN_SIZE = 7168 +#: Candidate counts (snapshots + the layer row) the kernel is instantiated for. +MIN_CANDIDATES, MAX_CANDIDATES = 2, 12 +#: Largest token count of the decode path (one CTA or cluster per token). +MAX_TOKENS = 1024 +#: Inputs of the layer row: the row itself; ``bf16(prefix + addend)`` (the C++ fused add); +#: ``bf16(residual + bf16(routed + shared))`` (the MoE residual update). +ADD_NONE, ADD_PREFIX, ADD_MOE = 0, 1, 2 + +#: Programmatic dependent launch unless TRTLLM_ENABLE_PDL=0 (the stack-wide switch). +_USE_PDL = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" +#: TRTLLM_K3_ATTN_RES_EARLY_TRIGGER=1: release the dependent grid right after ``griddepcontrol.wait`` +#: instead of after the last store, so the next kernel (a decode projection) can become resident and +#: stream its weights while this one finishes. +#: +#: Why it is safe in the Kimi-K3 decode graph. Since the trigger follows this kernel's own wait, every +#: grid upstream of it has completed when a dependent starts; only this grid may still run. The grids +#: that can start before it completes are its programmatic dependents (under stream capture that +#: includes cross-stream event joins) and, through dependents that trigger before their own wait, +#: theirs: the low-M / mid-M FP8 GEMMs (qkvg, q_a / kv_a, latent down, shared gate-up; weight warp), +#: the router GEMM, ``k3_kda_fab`` (warp 0) and behind them the KDA MTP decode and the MoE routing / +#: expert kernels. Before their wait these read only weights, weight scales, routing metadata and (KDA +#: decode) the layer's state pool, conv windows and replay rows, and write no global memory; this +#: kernel's outputs (normed row, updated prefix, MXFP8 copy) are read only after a wait. The next +#: attn_res launch, which reads the prefix this kernel stores (and snapshot rows) before its wait, is +#: separated from this one on every path by an MNNVL all-reduce (the o_proj all-reduce between A_pre +#: and A_mid; the latent and shared-expert all-reduces of the MoE block between A_mid and the next +#: layer's A_pre). Those kernels call ``cudaGridDependencySynchronize`` before they trigger, so the +#: next attn_res starts only after this grid completed. +_EARLY_TRIGGER = _F.attn_res_early_trigger + +_THREADS = 256 +_LOG2E = "0f3FB8AA3B" +_NEG_FLT_MAX = "0fFF7FFFFF" +_HIDDEN_F32 = "0f45E00000" # 7168.0 + +# --------------------------------------------------------------------------- +# Inline PTX helpers (fixed instructions: the C++ kernels are built with --use_fast_math) +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm_se(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _lv(value, *, loc=None, ip=None): + return Int64(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return Float32(value).ir_value(loc=loc, ip=ip) + + +def _f32_op(text, *args, loc=None, ip=None): + constraints = "=f," + ",".join("f" for _ in args) + return Float32( + _asm(T.f32(), [_fv(a, loc=loc, ip=ip) for a in args], text, constraints, loc=loc, ip=ip) + ) + + +@dsl_user_op +def _fma(a, b, c, *, loc=None, ip=None): + return _f32_op("fma.rn.f32 $0, $1, $2, $3;", a, b, c, loc=loc, ip=ip) + + +@dsl_user_op +def _add(a, b, *, loc=None, ip=None): + return _f32_op("add.rn.f32 $0, $1, $2;", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def _add_ftz(a, b, *, loc=None, ip=None): + return _f32_op("add.ftz.f32 $0, $1, $2;", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def _mul_ftz(a, b, *, loc=None, ip=None): + return _f32_op("mul.ftz.f32 $0, $1, $2;", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def _zero_f32(*, loc=None, ip=None): + return Float32(_asm(T.f32(), [], "mov.b32 $0, 0f00000000;", "=f", loc=loc, ip=ip)) + + +@dsl_user_op +def _rsigma(sum_sq, eps, *, loc=None, ip=None): + """``rsqrtf(sum_sq / 7168 + eps)`` as nvcc --use_fast_math emits it.""" + return _f32_op( + "{\n.reg .f32 t;\n" + f"div.approx.ftz.f32 t, $1, {_HIDDEN_F32};\n" + "add.ftz.f32 t, t, $2;\nrsqrt.approx.ftz.f32 $0, t;\n}", + sum_sq, + eps, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _max_ftz(a, b, *, loc=None, ip=None): + return _f32_op("max.ftz.f32 $0, $1, $2;", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def _neg_flt_max(*, loc=None, ip=None): + return Float32(_asm(T.f32(), [], f"mov.b32 $0, {_NEG_FLT_MAX};", "=f", loc=loc, ip=ip)) + + +@dsl_user_op +def _rcp_ftz(a, *, loc=None, ip=None): + return _f32_op("rcp.approx.ftz.f32 $0, $1;", a, loc=loc, ip=ip) + + +@dsl_user_op +def _round_bf16(a, *, loc=None, ip=None): + """``float(bf16_rn(a))``.""" + return _f32_op( + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.f32.bf16 $0, t;\n}", a, loc=loc, ip=ip + ) + + +@dsl_user_op +def _shfl_idx(value, lane: int, *, loc=None, ip=None): + return Float32( + _asm_se( + T.f32(), + [_fv(value, loc=loc, ip=ip)], + f"shfl.sync.idx.b32 $0, $1, {lane}, 0x1f, 0xffffffff;", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sts_f32(addr, a, *, loc=None, ip=None): + _asm_se( + None, + [_iv(addr, loc=loc, ip=ip), _fv(a, loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_f32(addr, *, loc=None, ip=None): + return Float32( + _asm_se( + T.f32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.f32 $0, [$1];", "=f,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _bar_sync(barrier_id: int, num_threads: int = _THREADS, *, loc=None, ip=None): + _asm_se(None, [], f"bar.sync {barrier_id}, {num_threads};", "", loc=loc, ip=ip) + + +@dsl_user_op +def _mbar_init(addr, count: int, *, loc=None, ip=None): + _asm_se( + None, + [_iv(addr, loc=loc, ip=ip)], + f"mbarrier.init.shared::cta.b64 [$0], {count};", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_mbar_init(*, loc=None, ip=None): + _asm_se(None, [], "fence.mbarrier_init.release.cluster;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _mbar_arrive_expect_tx(addr, nbytes: int, *, loc=None, ip=None): + _asm_se( + None, + [_iv(addr, loc=loc, ip=ip)], + f"mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], {nbytes};", + "r", + loc=loc, + ip=ip, + ) + + +_WAIT_LABEL = [0] + + +@dsl_user_op +def _mbar_wait_cluster(addr, *, loc=None, ip=None): + """Wait for phase 0 of a local mbarrier completed by peers' ``st.async`` (cluster acquire).""" + _WAIT_LABEL[0] += 1 + lab = f"K3AR_WAIT{_WAIT_LABEL[0]}" + _asm_se( + None, + [_iv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\n" + f"{lab}:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [$0], 0;\n" + f"@!p bra {lab};\n}}", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_async_b32(local_addr, rank, local_mbar, a, *, loc=None, ip=None): + _asm_se( + None, + [ + _iv(local_addr, loc=loc, ip=ip), + _iv(rank, loc=loc, ip=ip), + _iv(local_mbar, loc=loc, ip=ip), + _fv(a, loc=loc, ip=ip), + ], + "{\n.reg .b32 ra, rm;\nmapa.shared::cluster.u32 ra, $0, $1;\n" + "mapa.shared::cluster.u32 rm, $2, $1;\n" + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [ra], $3, [rm];\n}", + "r,r,r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cluster_arrive_relaxed(*, loc=None, ip=None): + _asm_se(None, [], "barrier.cluster.arrive.relaxed.aligned;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _cluster_wait(*, loc=None, ip=None): + _asm_se(None, [], "barrier.cluster.wait.acquire.aligned;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None): + _asm_se(None, [], "griddepcontrol.wait;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None): + _asm_se(None, [], "griddepcontrol.launch_dependents;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _ldg_v4(addr, nc: bool, *, loc=None, ip=None): + """16-byte global load of four packed bf16x2 words.""" + space = "global.nc" if nc else "global" + res = _asm_se( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + f"ld.{space}.v4.b32 {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return [Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _ldg_v4_if(valid, addr, nc: bool, *, loc=None, ip=None): + """:func:`_ldg_v4` predicated on ``valid != 0`` (zeros otherwise).""" + space = "global.nc" if nc else "global" + res = _asm_se( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(valid, loc=loc, ip=ip), _lv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $4, 0;\n" + "mov.b32 $0, 0;\nmov.b32 $1, 0;\nmov.b32 $2, 0;\nmov.b32 $3, 0;\n" + f"@p ld.{space}.v4.b32 {{$0, $1, $2, $3}}, [$5];\n}}", + "=&r,=&r,=&r,=&r,r,l", + loc=loc, + ip=ip, + ) + return [Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _stg_v4_if(valid, addr, words, *, loc=None, ip=None): + _asm_se( + None, + [_iv(valid, loc=loc, ip=ip), _lv(addr, loc=loc, ip=ip)] + + [_iv(w, loc=loc, ip=ip) for w in words], + "{\n.reg .pred p;\nsetp.ne.u32 p, $0, 0;\n@p st.global.v4.b32 [$1], {$2, $3, $4, $5};\n}", + "r,l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _unpack_bf16x2(word, *, loc=None, ip=None): + """``(float(lo), float(hi))`` of a packed bf16x2 word.""" + res = _asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [_iv(word, loc=loc, ip=ip)], + "{\n.reg .b32 t;\nshl.b32 t, $2, 16;\nmov.b32 $0, t;\n" + "and.b32 t, $2, 0xffff0000;\nmov.b32 $1, t;\n}", + "=f,=f,r", + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), res, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _pack_bf16x2(lo, hi, *, loc=None, ip=None): + """bf16x2 word ``{bf16_rn(hi), bf16_rn(lo)}``.""" + return Int32( + _asm( + T.i32(), + [_fv(hi, loc=loc, ip=ip), _fv(lo, loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_xor(value, offset: int, *, loc=None, ip=None): + return Float32( + _asm_se( + T.f32(), + [_fv(value, loc=loc, ip=ip)], + f"shfl.sync.bfly.b32 $0, $1, {offset}, 0x1f, 0xffffffff;", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_xor_b32(value, offset: int, *, loc=None, ip=None): + return Int32( + _asm_se( + T.i32(), + [_iv(value, loc=loc, ip=ip)], + f"shfl.sync.bfly.b32 $0, $1, {offset}, 0x1f, 0xffffffff;", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _stg_v2_if(valid, addr, w0, w1, *, loc=None, ip=None): + _asm_se( + None, + [_iv(valid, loc=loc, ip=ip), _lv(addr, loc=loc, ip=ip)] + + [_iv(w, loc=loc, ip=ip) for w in (w0, w1)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $0, 0;\n@p st.global.v2.b32 [$1], {$2, $3};\n}", + "r,l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u8_if(valid, addr, value, *, loc=None, ip=None): + _asm_se( + None, + [_iv(valid, loc=loc, ip=ip), _lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $0, 0;\n@p st.global.u8 [$1], $2;\n}", + "r,l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _selp_b32(cond, a, b, *, loc=None, ip=None): + """32-bit ``a if cond != 0 else b``.""" + return Int32( + _asm_se( + T.i32(), + [_iv(cond, loc=loc, ip=ip), _iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\nselp.b32 $0, $2, $3, p;\n}", + "=r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _selp(cond, a, b, *, loc=None, ip=None): + """``a if cond != 0 else b``.""" + return Float32( + _asm( + T.f32(), + [_iv(cond, loc=loc, ip=ip), _fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\nselp.f32 $0, $2, $3, p;\n}", + "=f,r,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ex2_scaled(x, *, loc=None, ip=None): + """``exp(x)`` as ``ex2.approx.ftz(x * log2(e))``.""" + return _f32_op( + f"{{\n.reg .f32 t;\nmul.ftz.f32 t, $1, {_LOG2E};\nex2.approx.ftz.f32 $0, t;\n}}", + x, + loc=loc, + ip=ip, + ) + + +def _warp_transpose_sum(values, lane): + """Sums every value over the warp in ``len(values) - 1 + 5 - log2(len)`` shuffles. + + ``len(values)`` must be a power of two <= 32. Afterwards lane ``l`` holds the warp total of + ``values[l >> (5 - log2(len(values)))]``. + """ + offset = 16 + while len(values) > 1: + half = len(values) // 2 + upper = Int32((lane & offset) != 0) + reduced = [] + for i in range(half): + mine = _selp(upper, values[i + half], values[i]) + send = _selp(upper, values[i], values[i + half]) + reduced.append(_add(mine, _shfl_xor(send, offset))) + values = reduced + offset //= 2 + total = values[0] + while offset >= 1: + total = _add(total, _shfl_xor(total, offset)) + offset //= 2 + return total + + +def _warp_sum(value): + for offset in (16, 8, 4, 2, 1): + value = _add(value, _shfl_xor(value, offset)) + return value + + +def _warp_max(value): + for offset in (16, 8, 4, 2, 1): + value = _max_ftz(value, _shfl_xor(value, offset)) + return value + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +class K3AttnResFastKernel: + """Vectorized variant: a cluster of ``c`` CTAs per token, one 16-byte vector per thread. + + Each CTA owns ``7168 / c`` contiguous elements and each thread eight of them (one 16-byte + load per row), so a token's rows are read with ``8 x`` fewer instructions than by the + per-element layout of the C++ kernels, and the statistics are summed in a different order: + the results are not bit-identical to the C++ kernels (the updated prefix sum, an elementwise + add, is). Before ``griddepcontrol.wait`` the snapshots' statistics, logits and their part of + the softmax and of the mixture (``sum_n exp(logit_n - m) v_n``) are complete; after it the + kernel reduces two numbers of the layer row (one cluster round), rescales the mixture and + reduces its sum of squares (a second round). + """ + + def __init__( + self, + num_candidates: int, + add_mode: int, + use_pdl: bool, + cluster: int = 4, + mxfp8: bool = False, + lamport: bool = False, + finish_world: int = 0, + finish_chunk: int = 0, + add_rn: bool = False, + finish_two_shot: bool = False, + entry_trigger: bool = False, + mx_atom: bool = False, + early_poll: bool = False, + finish_ag: bool = False, + ): + if not MIN_CANDIDATES <= num_candidates <= MAX_CANDIDATES: + raise ValueError(f"num_candidates={num_candidates} not in [2, 12]") + if early_poll and not finish_world: + raise ValueError("early_poll finishes a pushed all-reduce: it needs finish_world") + if finish_ag and (not finish_world or finish_two_shot): + raise ValueError("finish_ag finishes a one-shot call with all-gather rows") + # The finished call carries all-gather rows (k3_lamport_ar "ag"): the addend is + # bf16(y + bf16(slot sum)), the BF16 add of a column-parallel projection's output y + # (all-gathered) and the all-reduced partials (the K3 column-parallel MoE tail). + self.finish_ag = finish_ag + # c240: the MXFP8 scales in the mid-M GEMM's tcgen05 atom layout (one 128-row block) + # instead of row-major [M, 224] (see ``mx_atom`` of run_k3_attn_res). + self.mx_atom = mx_atom and mxfp8 + if add_mode not in (ADD_NONE, ADD_PREFIX, ADD_MOE) or cluster not in (1, 2, 4, 8): + raise ValueError(f"add_mode={add_mode} cluster={cluster}") + if lamport and add_mode != ADD_PREFIX: + raise ValueError("the Lamport addend is an ADD_PREFIX variant") + if (lamport or finish_world) and (not lamport or finish_world < 2 or finish_chunk < 1): + raise ValueError( + f"the finish needs the Lamport addend, finish_world={finish_world} >= 2 " + f"and finish_chunk={finish_chunk} >= 1" + ) + if finish_two_shot and not finish_world: + raise ValueError("finish_two_shot needs finish_world") + if add_rn and add_mode != ADD_PREFIX: + raise ValueError("add_rn selects the rounding of the ADD_PREFIX add") + self.mxfp8 = mxfp8 + self.lamport = lamport + # Finish: the addend is the pushed all-reduce over ``finish_world`` ranks, which this + # kernel completes with the one-shot or (``finish_two_shot``) the two-shot transport. + self.finish_world = finish_world + self.finish_chunk = finish_chunk + self.finish_two_shot = finish_two_shot + # bf16(prefix + addend) with a round-to-nearest add (the MoE residual update's + # arithmetic) instead of the C++ kernels' fast-math add. + self.add_rn = add_rn + # Release the dependent grid at kernel entry (see ``entry_trigger`` of run_k3_attn_res). + self.entry_trigger = entry_trigger and use_pdl + # Finish the pushed all-reduce before ``griddepcontrol.wait`` (see ``early_poll`` of + # run_k3_attn_res); without PDL the grid starts after its predecessor anyway. + self.early_poll = early_poll and use_pdl + self.n = num_candidates + self.k = num_candidates - 1 + self.c = cluster + self.add_mode = add_mode + self.use_pdl = use_pdl + self.vecs = HIDDEN_SIZE // 8 // cluster # 16-byte vectors per CTA, one per thread + self.threads = (self.vecs + 31) // 32 * 32 + self.warps = self.threads // 32 + self.parts = cluster * self.warps # warp partials per token + assert self.parts <= 32 + # Snapshot statistics (sq, dot) per warp, padded to a power of two for the transpose sum. + self.nvals = 2 * self.k + self.width = 1 << (self.nvals - 1).bit_length() + self.shift = 5 - (self.width.bit_length() - 1) + + @cute.jit + def __call__( + self, + layer_addr: Int64, + add_addr: Int64, + add2_addr: Int64, + block_addr: Int64, + res_addr: Int64, + rms_addr: Int64, + out_w_addr: Int64, + out_addr: Int64, + upd_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + lam_mc: Int64, + lam_rank: Int32, + m: Int32, + eps: Float32, + out_eps: Float32, + stream: cuda.CUstream, + ): + parts, nvals = self.parts, self.nvals + + @cute.struct + class SharedStorage: + snap: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, parts * nvals], 16] + last: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, parts * 2], 16] + outsq: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, parts], 16] + mbar: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 4], 16] + + self.shared_storage = SharedStorage + kernel = self.kernel( + layer_addr, + add_addr, + add2_addr, + block_addr, + res_addr, + rms_addr, + out_w_addr, + out_addr, + upd_addr, + xq_addr, + xs_addr, + lam_mc, + lam_rank, + m, + eps, + out_eps, + ) + if cutlass.const_expr(self.c > 1): + kernel.launch( + grid=(m * self.c, 1, 1), + block=(self.threads, 1, 1), + cluster=(self.c, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + else: + kernel.launch( + grid=(m, 1, 1), + block=(self.threads, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + layer_addr: Int64, + add_addr: Int64, + add2_addr: Int64, + block_addr: Int64, + res_addr: Int64, + rms_addr: Int64, + out_w_addr: Int64, + out_addr: Int64, + upd_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + lam_mc: Int64, + lam_rank: Int32, + m: Int32, + eps: Float32, + out_eps: Float32, + ): + c, k, parts, nvals = self.c, self.k, self.parts, self.nvals + warps, shift = self.warps, self.shift + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + storage = utils.SmemAllocator().allocate(self.shared_storage) + snap_smem = storage.snap.data_ptr().toint() + last_smem = storage.last.data_ptr().toint() + outsq_smem = storage.outsq.data_ptr().toint() + mbar_smem = storage.mbar.data_ptr().toint() + + if cutlass.const_expr(self.entry_trigger): + _griddepcontrol_launch_dependents() + token = bidx // c + rank = bidx - token * c + if cutlass.const_expr(self.finish_two_shot): + # The rows this rank owns in the two-shot call go to the first CTA groups. + token = _lam.two_shot_token(token, m, lam_rank, self.finish_world) + if cutlass.const_expr(c > 1): + if tidx == 0: + for r in cutlass.range_constexpr(3): + _mbar_init(mbar_smem + r * 8, 1) + _fence_mbar_init() + _cluster_arrive_relaxed() + + valid = Int32(1) + if cutlass.const_expr(self.vecs != self.threads): + valid = Int32(tidx < self.vecs) + col = Int64(rank * self.vecs + tidx) * 16 + row = Int64(token) * (HIDDEN_SIZE * 2) + + def load(addr, nc): + if cutlass.const_expr(self.vecs != self.threads): + return _ldg_v4_if(valid, addr, nc) + return _ldg_v4(addr, nc) + + def unpack(words): + out = [] + for w in words: + lo, hi = _unpack_bf16x2(w) + out.append(lo) + out.append(hi) + return out + + # ---- prologue: everything that does not depend on the previous grid ---- + res_w = unpack(load(res_addr + col, True)) + rms_w = unpack(load(rms_addr + col, True)) + out_w = load(out_w_addr + col, True) + plane = Int64(m) * (HIDDEN_SIZE * 2) + snaps = [load(block_addr + plane * n + row + col, True) for n in range(k)] + prefix = [] + if cutlass.const_expr(self.add_mode != ADD_NONE): + prefix = load(layer_addr + row + col, True) + q = [_mul_ftz(res_w[e], rms_w[e]) for e in range(8)] + + snap_f = [unpack(snaps[n]) for n in range(k)] + stats = [] + for n in cutlass.range_constexpr(k): + sq = _zero_f32() + dot = _zero_f32() + for e in cutlass.range_constexpr(8): + sq = _fma(snap_f[n][e], snap_f[n][e], sq) + dot = _fma(snap_f[n][e], q[e], dot) + stats.append(sq) + stats.append(dot) + for _ in cutlass.range_constexpr(self.width - nvals): + stats.append(_zero_f32()) + warp_stat = _warp_transpose_sum(stats, lane) + vidx = lane >> shift + sub = lane & ((1 << shift) - 1) + if cutlass.const_expr(c == 1): + if sub == 0: + if vidx < nvals: + _sts_f32(snap_smem + (warp * nvals + vidx) * 4, warp_stat) + _bar_sync(1, self.threads) + else: + _cluster_wait() + if tidx == 0: + _mbar_arrive_expect_tx(mbar_smem, parts * nvals * 4) + _mbar_arrive_expect_tx(mbar_smem + 8, parts * 2 * 4) + _mbar_arrive_expect_tx(mbar_smem + 16, parts * 4) + slot = snap_smem + ((rank * warps + warp) * nvals + vidx) * 4 + for r in cutlass.range_constexpr(c): + if vidx < nvals: + if sub == r % (1 << shift): + _st_async_b32(slot, r, mbar_smem, warp_stat) + _mbar_wait_cluster(mbar_smem) + + # Snapshot totals (lane n < k), logits and the snapshot part of the softmax. + logit = _neg_flt_max() + if lane < k: + tot_sq = _zero_f32() + tot_dot = _zero_f32() + for j in cutlass.range_constexpr(parts): + tot_sq = _add(tot_sq, _lds_f32(snap_smem + (j * nvals + 2 * lane) * 4)) + tot_dot = _add(tot_dot, _lds_f32(snap_smem + (j * nvals + 2 * lane + 1) * 4)) + logit = _mul_ftz(tot_dot, _rsigma(tot_sq, eps)) + snap_max = _warp_max(logit) + weight = _zero_f32() + if lane < k: + weight = _ex2_scaled(_add(logit, -snap_max)) + snap_sum = _warp_sum(weight) + premix = [_zero_f32() for _ in range(8)] + for n in cutlass.range_constexpr(k): + w_n = _shfl_idx(weight, n) + for e in cutlass.range_constexpr(8): + premix[e] = _fma(w_n, snap_f[n][e], premix[e]) + + # ---- wait for the previous grid (early_poll: after the mix, before the first store) ---- + if cutlass.const_expr(self.use_pdl and not self.early_poll): + _griddepcontrol_wait() + if cutlass.const_expr(_EARLY_TRIGGER and not self.entry_trigger): + _griddepcontrol_launch_dependents() + + upd_words = [] + arr_old = Int32(0) # early_poll: thread 0's arrival count (arrive_count / arrive_rotate) + if cutlass.const_expr(self.add_mode == ADD_PREFIX): + if cutlass.const_expr(self.lamport): + # The addend is the all-reduce whose input every rank pushed into the next call on a + # Lamport workspace (add_addr: its local unicast base, add2_addr: its state flags, + # lam_mc: its multicast base, lam_rank: this rank), and this kernel is the call's + # finish: it reads the rotation state (final: the producer completed before the + # wait above), counts its CTA (the last one advances the state) and re-arms the + # previous call's buffer. One-shot transport: it sums the token's slot of every + # rank in rank order like k3_lamport_ar.run_lamport_finish. Two-shot transport: on + # the token's owner it sums the scatter slots and multicasts the row, elsewhere it + # polls the broadcast row. Threads past the row work on the row's first vector + # (and do not store) and use zeros. + # early_poll: all of it before the wait (the producer released this grid after + # its own wait, so the state is final and the dirty buffer's readers are done), + # except the state advance, which follows the wait (below). + W = self.finish_world + cur, dirty, nbytes, dstages, c0, c1, c2, c3 = _lam.lamport_state(add2_addr) + _bar_sync(1, self.threads) + if cutlass.const_expr(not self.early_poll): + if tidx == 0: + _lam.arrive_finish( + add2_addr, + m * c, + cur, + nbytes, + m, + W, + HIDDEN_SIZE, + not self.finish_two_shot, + 1 if self.finish_ag else 0, + ) + _lam.clear_dirty_stages( + add_addr, + dirty, + nbytes, + dstages, + c0, + c1, + c2, + c3, + Int32(bidx * self.threads + tidx), + m * (c * self.threads), + ) + off = _selp_b32(valid, Int32(rank * self.vecs + tidx) * 16, Int32(0)) + # early_poll: every thread polls its slots from the start, while the producer + # still runs. + if cutlass.const_expr(self.finish_two_shot): + words = _lam.finish_two_shot_v4( + add_addr, + lam_mc, + cur, + nbytes, + token, + off, + lam_rank, + valid, + W, + self.finish_chunk, + HIDDEN_SIZE, + ) + addend = unpack([_selp_b32(valid, w, Int32(0)) for w in words]) + else: + acc = _lam.poll_reduce_v4( + _lam.stage_addr(add_addr, cur, Int32(0), 1, nbytes) + + Int64(token * W) * (HIDDEN_SIZE * 2) + + Int64(off), + HIDDEN_SIZE * 2, + W, + self.finish_chunk, + ) + if cutlass.const_expr(self.finish_ag): + # bf16(y + bf16(sum)) with y this vector of the token's all-gather row. + yw = _lam.poll_v4( + _lam.ag_row_addr(add_addr, cur, nbytes, m, token, W, HIDDEN_SIZE) + + Int64(off) + ) + sums = [_round_bf16(a) for a in acc] + ys = unpack(yw) + addend = [ + _selp(valid, _round_bf16(_add(ys[e], sums[e])), _zero_f32()) + for e in range(8) + ] + else: + addend = [_selp(valid, _round_bf16(a), _zero_f32()) for a in acc] + if cutlass.const_expr(self.early_poll): + # Every thread of this CTA has seen its vectors of every rank, so every + # producer CTA that pushed them has read the rotation state (the push address + # depends on it); once every CTA of this grid has counted itself, the last + # one may advance it. The count is a relaxed atomic issued here whose result + # thread 0 reads only after its stores (arrive_rotate below): its round trip + # overlaps the mix instead of delaying warp 0 into the cluster rounds. + _bar_sync(1, self.threads) + if tidx == 0: + arr_old = _lam.arrive_count(add2_addr) + else: + addend = unpack(load(add_addr + row + col, False)) + pre = unpack(prefix) + if cutlass.const_expr(self.add_rn): + x = [_round_bf16(_add(pre[e], addend[e])) for e in range(8)] + else: + x = [_round_bf16(_add_ftz(pre[e], addend[e])) for e in range(8)] + elif cutlass.const_expr(self.add_mode == ADD_MOE): + routed = unpack(load(add_addr + row + col, False)) + shared = unpack(load(add2_addr + row + col, False)) + pre = unpack(prefix) + x = [ + _round_bf16(_add(pre[e], _round_bf16(_add(routed[e], shared[e])))) for e in range(8) + ] + else: + x = unpack(load(layer_addr + row + col, False)) + if cutlass.const_expr(self.add_mode != ADD_NONE): + upd_words = [_pack_bf16x2(x[2 * i], x[2 * i + 1]) for i in range(4)] + if cutlass.const_expr(not self.early_poll): + _stg_v4_if(valid, upd_addr + row + col, upd_words) + + sq = _zero_f32() + dot = _zero_f32() + for e in cutlass.range_constexpr(8): + sq = _fma(x[e], x[e], sq) + dot = _fma(x[e], q[e], dot) + # Lanes 0-15 hold the warp's sum of squares, lanes 16-31 its dot product. + pair = _warp_transpose_sum([sq, dot], lane) + half = lane & 15 + if cutlass.const_expr(c == 1): + if half == 0: + _sts_f32(last_smem + (warp * 2 + (lane >> 4)) * 4, pair) + _bar_sync(2, self.threads) + else: + if half < c: + _st_async_b32( + last_smem + ((rank * warps + warp) * 2 + (lane >> 4)) * 4, + half, + mbar_smem + 8, + pair, + ) + _mbar_wait_cluster(mbar_smem + 8) + part_sq = _zero_f32() + part_dot = _zero_f32() + if lane < parts: + part_sq = _lds_f32(last_smem + lane * 8) + part_dot = _lds_f32(last_smem + lane * 8 + 4) + tot = _warp_transpose_sum([part_sq, part_dot], lane) + tot_sq = _shfl_idx(tot, 0) + tot_dot = _shfl_idx(tot, 16) + last_logit = _mul_ftz(tot_dot, _rsigma(tot_sq, eps)) + + # Softmax over all candidates: p_n = exp(logit_n - m_s) * a / den, p_last = b / den. + new_max = _max_ftz(snap_max, last_logit) + a = _ex2_scaled(_add(snap_max, -new_max)) + b = _ex2_scaled(_add(last_logit, -new_max)) + inv = _rcp_ftz(_fma(snap_sum, a, b)) + scale = _mul_ftz(a, inv) + p_last = _mul_ftz(b, inv) + mixed = [_round_bf16(_fma(p_last, x[e], _mul_ftz(scale, premix[e]))) for e in range(8)] + out_sq = _zero_f32() + for e in cutlass.range_constexpr(8): + out_sq = _fma(mixed[e], mixed[e], out_sq) + out_sq = _warp_sum(out_sq) + if cutlass.const_expr(c == 1): + if lane == 0: + _sts_f32(outsq_smem + warp * 4, out_sq) + _bar_sync(3, self.threads) + else: + if lane < c: + _st_async_b32(outsq_smem + (rank * warps + warp) * 4, lane, mbar_smem + 16, out_sq) + _mbar_wait_cluster(mbar_smem + 16) + part = _zero_f32() + if lane < parts: + part = _lds_f32(outsq_smem + lane * 4) + out_rs = _rsigma(_warp_sum(part), out_eps) + + w_out = unpack(out_w) + y = [_mul_ftz(_round_bf16(_mul_ftz(mixed[e], out_rs)), w_out[e]) for e in range(8)] + words = [_pack_bf16x2(y[2 * i], y[2 * i + 1]) for i in range(4)] + if cutlass.const_expr(self.early_poll): + # Every store of this grid follows the wait: its output buffers may reuse memory that + # the producer, or a grid between it and this one, still reads. The dependents are + # released after the stores (below) or at entry. + _griddepcontrol_wait() + _stg_v4_if(valid, upd_addr + row + col, upd_words) + _stg_v4_if(valid, out_addr + row + col, words) + if cutlass.const_expr(self.mxfp8): + # MXFP8 form of the stored BF16 row for the low-M FP8 GEMM (k3_fp8w_lowm_mxfp8): four + # consecutive lanes hold one 1x32 group; the GEMM's own amax / exponent / E4M3 + # conversion, so its operands are bitwise those of its in-kernel quantization. + amax = mx_ops.amax_bf16x8_words(*words) + amax = mx_ops.amax_pair(amax, _shfl_xor_b32(amax, 1)) + amax = mx_ops.amax_pair(amax, _shfl_xor_b32(amax, 2)) + e_a = mx_ops.ue8m0_of_amax(amax) + q0, q1 = mx_ops.quant_bf16x8_to_e4m3(*words, mx_ops.bf16x2_inv_scale(e_a)) + chan = Int64(rank * self.vecs + tidx) * 8 + qrow = Int64(token) * HIDDEN_SIZE + _stg_v2_if(valid, xq_addr + qrow + chan, q0, q1) + if cutlass.const_expr(self.mx_atom): + # Mid-M scale atom of group g = chan / 32 for token t < 128 (k3_fp8w_lowm_gemm + # K3MxFp8QuantKernel with one scale block): byte (g / 4) * 512 + (t % 32) * 16 + # + (t / 32) * 4 + g % 4. + grp = chan >> 5 + sf_off = ( + (grp >> 2) * 512 + + Int64(token & 31) * 16 + + Int64((token >> 5) & 3) * 4 + + (grp & 3) + ) + _stg_u8_if(valid & Int32((tidx & 3) == 0), xs_addr + sf_off, e_a) + else: + _stg_u8_if( + valid & Int32((tidx & 3) == 0), + xs_addr + Int64(token) * (HIDDEN_SIZE // 32) + (chan >> 5), + e_a, + ) + if cutlass.const_expr(self.early_poll): + # The last CTA to count advances the workspace (its count was issued after the poll). + if tidx == 0: + _lam.arrive_rotate( + add2_addr, + arr_old, + m * c, + cur, + nbytes, + m, + self.finish_world, + HIDDEN_SIZE, + not self.finish_two_shot, + 1 if self.finish_ag else 0, + ) + + # Release at the end unless released at entry or (TRTLLM_K3_ATTN_RES_EARLY_TRIGGER, not + # with early_poll) after the wait. + if cutlass.const_expr( + self.use_pdl and not self.entry_trigger and (self.early_poll or not _EARLY_TRIGGER) + ): + _bar_sync(4, self.threads) + _griddepcontrol_launch_dependents() + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def fast_cluster_size(num_tokens: int, num_candidates: int) -> int: + """CTAs per token of :class:`K3AttnResFastKernel`. + + Measured on GB300 behind a PDL producer: eight CTAs per token are fastest for a handful of + tokens with at least three snapshots (smaller per-CTA prologue), four from ~32 tokens on + (eight-CTA clusters then fill the machine) and for two or three candidates. On SM107 (212 + SMs) this choice is within 0.25 us of the fastest cluster size for up to 64 tokens. + """ + return 8 if num_tokens <= 16 and num_candidates >= 4 else 4 + + +@functools.cache +def _compiled( + num_candidates: int, + add_mode: int, + use_pdl: bool, + cluster: int | None, + mxfp8: bool = False, + lamport: bool = False, + finish_world: int = 0, + finish_chunk: int = 0, + add_rn: bool = False, + finish_two_shot: bool = False, + entry_trigger: bool = False, + mx_atom: bool = False, + early_poll: bool = False, + finish_ag: bool = False, +): + kernel = K3AttnResFastKernel( + num_candidates, + add_mode, + use_pdl, + cluster, + mxfp8, + lamport, + finish_world, + finish_chunk, + add_rn, + finish_two_shot, + entry_trigger, + mx_atom, + early_poll, + finish_ag, + ) + addr = Int64(1 << 20) + compiled = cute.compile( + kernel, + addr, + addr, + addr, + addr, + addr, + addr, + addr, + addr, + addr, + addr, + addr, + addr, + Int32(0), + Int32(1), + Float32(1e-6), + Float32(1e-6), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + # TRTLLM_K3_CUTE_CARVEOUT_FIX=1: the SMs this grid lands on keep the maximum shared-memory + # split, so the CTAs of its programmatic dependent (a GEMM with a deep shared-memory weight + # ring) can become resident beside it. Without a preference the driver gives these SMs a small + # split, and the dependent's CTAs wait for this grid's CTAs to exit. + return k3_launch_attrs.apply_shared_carveout(compiled, k3_launch_attrs.MAX_SHARED_CARVEOUT) + + +@functools.cache +def _prefer_max_carveout(compiled, device_index: int) -> None: + """Set the preferred shared-memory carve-out of the kernel in ``compiled`` to 100 % (the + CuTe launcher does not pass one), so that an SM running one of its CTAs keeps the shared + memory a CTA of a large-footprint dependent needs to land beside it.""" + del device_index # CUDA libraries are context free; the attribute is set once per kernel + attribute = cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT + library = compiled.library + status, count = cuda.cuLibraryGetKernelCount(library) + if status != cuda.CUresult.CUDA_SUCCESS or count != 1: + raise RuntimeError(f"cuLibraryGetKernelCount: {status}, {count} kernels") + status, kernels = cuda.cuLibraryEnumerateKernels(count, library) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuLibraryEnumerateKernels: {status}") + status, function = cuda.cuKernelGetFunction(kernels[0]) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuKernelGetFunction: {status}") + (status,) = cuda.cuFuncSetAttribute(function, attribute, 100) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuFuncSetAttribute(carve-out): {status}") + + +@functools.cache +def _is_supported_device(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3), (10, 7)) + + +def can_use_k3_attn_res( + layer: torch.Tensor, + addends: tuple[torch.Tensor, ...], + block: torch.Tensor, + *weights: torch.Tensor, +) -> bool: + """Whether :func:`run_k3_attn_res` accepts ``layer [M, H]``, 0-2 addends ``[M, H]`` and + ``block [N - 1, M, H]``.""" + if not layer.is_cuda or not _is_supported_device(layer.device.index or 0): + return False + if layer.dim() != 2 or layer.shape[1] != HIDDEN_SIZE or not 1 <= layer.shape[0] <= MAX_TOKENS: + return False + m = layer.shape[0] + if block.dim() != 3 or tuple(block.shape[1:]) != (m, HIDDEN_SIZE): + return False + if not MIN_CANDIDATES <= block.shape[0] + 1 <= MAX_CANDIDATES or len(addends) > 2: + return False + for t in (layer, block, *weights, *addends): + if t.dtype != torch.bfloat16 or not t.is_contiguous() or t.device != layer.device: + return False + if any(a.shape != layer.shape for a in addends): + return False + return all(w.numel() == HIDDEN_SIZE for w in weights) + + +def run_k3_attn_res( + layer: torch.Tensor, + addends: tuple[torch.Tensor, ...], + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + output: torch.Tensor, + updated: torch.Tensor | None, + *, + use_pdl: bool = _USE_PDL, + cluster: int | None = None, + mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, + lamport_addend: tuple[int, int] | None = None, + lamport_finish: tuple[int, int, bool, int, int] | None = None, + add_rn: bool = False, + entry_trigger: bool = False, + mx_atom: bool = False, + early_poll: bool = False, + finish_ag: bool = False, +) -> None: + """Launch on the current stream (inputs are not re-validated; see :func:`can_use_k3_attn_res`). + + Args: + layer: ``bf16 [M, 7168]``: the layer row without addends, else the prefix sum. + addends: ``()``; ``(addend,)`` for ``bf16(layer + addend)``; ``(routed, shared)`` for + ``bf16(layer + bf16(routed + shared))`` (all ``bf16 [M, 7168]``). + block: ``bf16 [N - 1, M, 7168]`` snapshot rows, ``2 <= N <= 12``. + res_weight, rms_weight: Score projection and its RMSNorm weight, ``bf16 [7168]``. + out_weight: Weight of the trailing RMSNorm, ``bf16 [7168]``. + rms_eps, out_eps: Epsilons of the two norms. + output: ``bf16 [M, 7168]`` normalized mixture. + updated: ``bf16 [M, 7168]`` updated prefix sum; required with addends. + use_pdl: Launch with programmatic dependent launch. + cluster: CTAs per token (default: :func:`fast_cluster_size`). + mxfp8_out: ``(xq, xs)``: also store the MXFP8 form of ``output`` for + :func:`k3_fp8w_lowm_mxfp8.run_k3_fp8w_lowm_gemm_mxfp8`: E4M3 ``xq [M, 7168]`` and one + UE8M0 exponent per 1x32 group ``xs [M, 224]`` (uint8), bitwise what that GEMM's BF16 + twin quantizes ``output`` to. + lamport_addend: ``(uc_local, flags)`` device addresses of an MNNVL Lamport workspace + (``k3_lamport_ar.LamportWorkspace``) instead of ``addends`` (which must be ``()``): + ``bf16(layer + addend)`` with the ``[M, 7168]`` all-reduce ``lamport_finish`` + describes. + lamport_finish: ``(world, chunk, two_shot, mc_base, rank)``, required with + ``lamport_addend``: the addend is the all-reduce whose input every rank pushed into + the next call on the workspace (``mc_base``: its multicast base, ``rank``: this rank + of ``world``) with + the one-shot or (``two_shot``) the two-shot transport, and this kernel finishes that + call as ``k3_lamport_ar.run_lamport_finish`` does (rank-order sums, ``chunk`` ranks per + partial sum; two-shot: the owner of each row reduces and multicasts it, the others + poll it), advances the rotation state and re-arms the previous call's buffer, so that + no finish kernel runs in front of it. The addend has the bits of the finished + all-reduce. The producer must trigger only after its own ``griddepcontrol.wait``; with + the two-shot transport every rank's CTAs must become co-resident (the grid is small: + ``M * cluster`` CTAs, the owned rows first). + add_rn: ``bf16(layer + addend)`` with a round-to-nearest FP32 add (the MoE residual + update's arithmetic) instead of the C++ kernels' fast-math add (one addend). + entry_trigger: (PDL) release the dependent grid at kernel entry + instead of after ``griddepcontrol.wait`` or the last store, and prefer the largest + shared-memory carve-out for this kernel, so that a weight-streaming projection + behind it (the KDA QKVG GEMM that stages its weights) becomes resident + while this kernel and its predecessor still run, also on the SMs of this kernel's + CTAs. The caller guarantees that the dependents, and the grids they release before + their own wait, read before their ``griddepcontrol.wait`` only memory that no grid + still running writes (weights, and for the KDA verify kernel the layer's state pool, + conv windows and replay rows of the previous step) and write no global memory + before it. The results are the same bits. + mx_atom: (c240, with ``mxfp8_out``, M <= 128) ``xs`` is ``uint8 [(7168 / 128) * 512]`` + in the scale layout of the mid-M FP8 GEMM (``k3_fp8w_lowm_gemm.run_mid_prequantized`` + with one 128-row scale block): the same exponents, so that GEMM on ``(xq, xs)`` is + bitwise its quantization kernel + GEMM on ``output``. Atom rows past M are not + written (the GEMM's columns of those rows are not stored). + early_poll: (with ``lamport_finish``, PDL) finish the pushed all-reduce before + ``griddepcontrol.wait``: as soon as it is resident the kernel reads the rotation + state, re-arms the previous buffer, polls its slots, counts its CTA (a relaxed + atomic: every producer CTA that pushed those slots has read the state) and computes + the mix and the norm; then it waits for the previous grid, stores its outputs (which + may reuse memory that grids still running read) and the last CTA advances the state; + it releases its dependents after its stores (or at entry). + This is the post-projection half of a collective "sandwich" as two PDL grids: the + producer (the push GEMM) releases this grid after its own wait, every rank's slots + are polled while the producer runs, and the kernel boundary leaves the critical path. + Requires that this grid is launched (programmatically or not) after the producer + and that the producer triggers only after its own ``griddepcontrol.wait`` (the push + variants do), so every earlier call on the workspace has completed at launch. Same + bits as without it. + finish_ag: (with a one-shot ``lamport_finish``) the call carries all-gather rows + (``k3_lamport_ar`` "ag"): the addend is ``bf16(y + bf16(slot sum))`` with ``y`` the + token's all-gather row, the BF16 add of a column-parallel projection's output and + the all-reduce of the ranks' partials (``ADD_MOE``'s ``bf16(routed + shared)``). + """ + if early_poll and lamport_finish is None: + raise ValueError("early_poll needs lamport_finish") + mx_atom = mx_atom and mxfp8_out is not None + if mx_atom and layer.shape[0] > 128: + raise ValueError(f"mx_atom covers at most 128 rows, got {layer.shape[0]}") + lamport = lamport_addend is not None + if lamport and addends: + raise ValueError("pass either addends or lamport_addend") + if (lamport_finish is not None) != lamport: + raise ValueError("lamport_addend and lamport_finish go together") + add_mode = ADD_PREFIX if lamport else len(addends) + if add_rn and add_mode != ADD_PREFIX: + raise ValueError("add_rn needs exactly one addend") + if cluster is None: + cluster = fast_cluster_size(layer.shape[0], block.shape[0] + 1) + mxfp8 = mxfp8_out is not None + finish_world, finish_chunk, finish_two_shot, lam_mc, lam_rank = ( + lamport_finish if lamport_finish is not None else (0, 0, False, 0, 0) + ) + entry_trigger = entry_trigger and use_pdl + if finish_ag and (lamport_finish is None or finish_two_shot): + raise ValueError("finish_ag needs a one-shot lamport_finish") + compiled = _compiled( + block.shape[0] + 1, + add_mode, + use_pdl, + cluster, + mxfp8, + lamport, + finish_world, + finish_chunk, + add_rn, + finish_two_shot, + entry_trigger, + mx_atom, + early_poll and use_pdl, + finish_ag, + ) + if entry_trigger: + _prefer_max_carveout(compiled, layer.device.index or 0) + if lamport: + add_addr, add2_addr = lamport_addend + else: + add_addr = addends[0].data_ptr() if add_mode >= 1 else 0 + add2_addr = addends[1].data_ptr() if add_mode == 2 else 0 + compiled( + Int64(layer.data_ptr()), + Int64(add_addr), + Int64(add2_addr), + Int64(block.data_ptr()), + Int64(res_weight.data_ptr()), + Int64(rms_weight.data_ptr()), + Int64(out_weight.data_ptr()), + Int64(output.data_ptr()), + Int64(updated.data_ptr() if add_mode else 0), + Int64(mxfp8_out[0].data_ptr() if mxfp8 else 0), + Int64(mxfp8_out[1].data_ptr() if mxfp8 else 0), + Int64(lam_mc), + Int32(lam_rank), + Int32(layer.shape[0]), + Float32(rms_eps), + Float32(out_eps), + cuda.CUstream(torch.cuda.current_stream(layer.device).cuda_stream), + ) + + +def k3_attn_res_rmsnorm( + layer: torch.Tensor, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, +) -> torch.Tensor: + """Attention-residual mix of ``block`` and ``layer`` followed by RMSNorm. + + ``trtllm::attn_res_rmsnorm_fwd`` (``layer [M, 7168]``, ``block [N - 1, M, 7168]``) up to one + BF16 rounding step in a small fraction of the output elements. + """ + output = torch.empty_like(layer) + run_k3_attn_res( + layer, (), block, res_weight, rms_weight, out_weight, rms_eps, out_eps, output, None + ) + return output + + +def k3_attn_res_add_rmsnorm( + prefix: torch.Tensor, + addend: torch.Tensor, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """``(bf16(prefix + addend), output)``: the residual add, then the mix and RMSNorm. + + ``trtllm::attn_res_add_rmsnorm_fwd``: the same updated prefix, the output up to one BF16 + rounding step in a small fraction of the elements. + """ + output = torch.empty_like(prefix) + updated = torch.empty_like(prefix) + run_k3_attn_res( + prefix, + (addend,), + block, + res_weight, + rms_weight, + out_weight, + rms_eps, + out_eps, + output, + updated, + ) + return updated, output + + +def k3_attn_res_moe_add_rmsnorm( + residual: torch.Tensor, + routed: torch.Tensor, + shared: torch.Tensor, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """``(prefix, output)`` with ``prefix = bf16(residual + bf16(routed + shared))``. + + The MoE residual update of the previous layer (bit-identical to + ``kimi_k3_moe_residual._fused_moe_residual_add``) followed by the mix and RMSNorm of + :func:`k3_attn_res_rmsnorm` on ``prefix``. + """ + output = torch.empty_like(residual) + prefix = torch.empty_like(residual) + run_k3_attn_res( + residual, + (routed, shared), + block, + res_weight, + rms_weight, + out_weight, + rms_eps, + out_eps, + output, + prefix, + ) + return prefix, output + + +def _mxfp8_outputs(rows: torch.Tensor, atom: bool = False) -> tuple[torch.Tensor, torch.Tensor]: + m = rows.shape[0] + xs_shape = ((HIDDEN_SIZE // 128) * 512,) if atom else (m, HIDDEN_SIZE // 32) + return ( + torch.empty((m, HIDDEN_SIZE), dtype=torch.float8_e4m3fn, device=rows.device), + torch.empty(xs_shape, dtype=torch.uint8, device=rows.device), + ) + + +def k3_attn_res_moe_add_rmsnorm_mx( + residual: torch.Tensor, + routed: torch.Tensor, + shared: torch.Tensor, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + entry_trigger: bool = False, + mx_atom: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """:func:`k3_attn_res_moe_add_rmsnorm` plus the MXFP8 form ``(xq, xs)`` of its output; + ``entry_trigger`` and ``mx_atom`` (the scales in the mid-M GEMM's layout): see + :func:`run_k3_attn_res`.""" + output = torch.empty_like(residual) + prefix = torch.empty_like(residual) + xq, xs = _mxfp8_outputs(residual, mx_atom) + run_k3_attn_res( + residual, + (routed, shared), + block, + res_weight, + rms_weight, + out_weight, + rms_eps, + out_eps, + output, + prefix, + mxfp8_out=(xq, xs), + entry_trigger=entry_trigger, + mx_atom=mx_atom, + ) + return prefix, output, xq, xs + + +def _run_lamport_add( + prefix, + uc_local, + flags, + block, + weights, + rms_eps, + out_eps, + finish, + add_rn, + mx, + early_poll=False, + entry_trigger=False, + finish_ag=False, + updated=None, +): + if updated is None: + updated = torch.empty_like(prefix) + output = torch.empty_like(prefix) + mxfp8_out = _mxfp8_outputs(prefix) if mx else None + run_k3_attn_res( + prefix, + (), + block, + *weights, + rms_eps, + out_eps, + output, + updated, + mxfp8_out=mxfp8_out, + lamport_addend=(uc_local, flags), + lamport_finish=finish, + add_rn=add_rn, + early_poll=early_poll, + entry_trigger=entry_trigger, + finish_ag=finish_ag, + ) + return (updated, output) + (tuple(mxfp8_out) if mx else ()) + + +_LAMPORT_ADD_DOC = """``(bf16(prefix + addend), output)``: the mix and RMSNorm of + :func:`k3_attn_res_add_rmsnorm` with the addend read from the Lamport workspace ``(uc_local, + flags)``: the pushed all-reduce over ``finish_world`` ranks that this kernel finishes, + ``finish_chunk`` ranks per partial sum, with + the one-shot or (``finish_two_shot``) the two-shot transport (``mc_base``: the workspace's + multicast base, ``rank``: this rank; ``lamport_finish`` of :func:`run_k3_attn_res`). + ``add_rn``: round-to-nearest add; ``early_poll`` (with a finish): finish the call before the + grid wait and store after it (``early_poll`` of :func:`run_k3_attn_res`, same bits); + ``entry_trigger``: release the dependent at kernel entry (see :func:`run_k3_attn_res`); + ``finish_ag`` (one-shot finish): the call carries all-gather rows, the addend is + ``bf16(y + bf16(slot sum))`` (``finish_ag`` of :func:`run_k3_attn_res`).""" + + +def k3_attn_res_lamport_add_rmsnorm( + prefix: torch.Tensor, + uc_local: int, + flags: int, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + finish_world: int, + finish_chunk: int, + finish_two_shot: bool, + mc_base: int, + rank: int, + add_rn: bool, + early_poll: bool = False, + entry_trigger: bool = False, + finish_ag: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + return _run_lamport_add( + prefix, + uc_local, + flags, + block, + (res_weight, rms_weight, out_weight), + rms_eps, + out_eps, + (finish_world, finish_chunk, finish_two_shot, mc_base, rank), + add_rn, + False, + early_poll, + entry_trigger, + finish_ag, + ) + + +k3_attn_res_lamport_add_rmsnorm.__doc__ = _LAMPORT_ADD_DOC + + +def k3_attn_res_lamport_add_rmsnorm_mx( + prefix: torch.Tensor, + uc_local: int, + flags: int, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + finish_world: int, + finish_chunk: int, + finish_two_shot: bool, + mc_base: int, + rank: int, + add_rn: bool, + early_poll: bool = False, + entry_trigger: bool = False, + finish_ag: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """:func:`k3_attn_res_lamport_add_rmsnorm` plus the MXFP8 form ``(xq, xs)`` of its output.""" + return _run_lamport_add( + prefix, + uc_local, + flags, + block, + (res_weight, rms_weight, out_weight), + rms_eps, + out_eps, + (finish_world, finish_chunk, finish_two_shot, mc_base, rank), + add_rn, + True, + early_poll, + entry_trigger, + finish_ag, + ) + + +def k3_attn_res_lamport_add_rmsnorm_into( + updated: torch.Tensor, + mx: bool, + prefix: torch.Tensor, + uc_local: int, + flags: int, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + finish_world: int, + finish_chunk: int, + finish_two_shot: bool, + mc_base: int, + rank: int, + add_rn: bool, + early_poll: bool = False, + entry_trigger: bool = False, + finish_ag: bool = False, +) -> tuple[torch.Tensor, ...]: + """:func:`k3_attn_res_lamport_add_rmsnorm` (``mx``: the ``_mx`` variant) storing the updated + prefix sum into ``updated`` (contiguous, ``prefix``'s shape and dtype, not overlapping the + inputs), e.g. the next snapshot slot of the attention-residual bank. Same bits as + :func:`k3_attn_res_lamport_add_rmsnorm`.""" + if ( + updated.shape != prefix.shape + or updated.dtype != prefix.dtype + or not updated.is_contiguous() + ): + raise ValueError("updated must be a contiguous tensor of prefix's shape and dtype") + return _run_lamport_add( + prefix, + uc_local, + flags, + block, + (res_weight, rms_weight, out_weight), + rms_eps, + out_eps, + (finish_world, finish_chunk, finish_two_shot, mc_base, rank), + add_rn, + mx, + early_poll, + entry_trigger, + finish_ag, + updated=updated, + ) + + +__all__ = [ + "ADD_MOE", + "ADD_NONE", + "ADD_PREFIX", + "HIDDEN_SIZE", + "MAX_CANDIDATES", + "MAX_TOKENS", + "MIN_CANDIDATES", + "can_use_k3_attn_res", + "k3_attn_res_add_rmsnorm", + "k3_attn_res_lamport_add_rmsnorm", + "k3_attn_res_lamport_add_rmsnorm_mx", + "k3_attn_res_lamport_add_rmsnorm_into", + "k3_attn_res_moe_add_rmsnorm", + "k3_attn_res_moe_add_rmsnorm_mx", + "k3_attn_res_rmsnorm", + "run_k3_attn_res", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha.py new file mode 100644 index 000000000000..f4c426e951ce --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha.py @@ -0,0 +1,3732 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Rubin SM107 runtime FMHA kernel adapted from DKG's attention examples. + +The K3 path adds separate K-nope/K-pe loads and in-epilogue chunk merging. +Only the runtime kernel class is vendored; DKG CLI and reference code are excluded. +The launcher is ``k3_context_fmha_op.run_k3_context_fmha`` (entry ``ops/attention/k3_context_fmha``). + +``fmha_utils.cfence()`` / ``warp_switch()`` / ``sched_res_busy_xu64()`` / +``reset_sched_res_busy_xu64()`` emit the internal CuTe DSL build's scheduling hints where that +build provides them and nothing otherwise (``k3_context_fmha_helpers.HAS_SCHEDULING_HINTS``). +""" + +import math +from functools import partial +from typing import Callable, Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as cute_testing +import cutlass.pipeline as pipeline +import cutlass.utils as utils + +# Shared tcgen05 layout/copy helpers are also used by Rubin. +import cutlass.utils.blackwell_helpers as tcgen05_utils +from cutlass import Arch, testing +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.nvgpu.common import OperandMajorMode +from cutlass.cute.typing import Float32, Int8, Int32, Int64 +from cutlass.cutlass_dsl import BaseDSL +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from . import k3_context_fmha_helpers as fmha_utils + + +def make_thread_cooperative_group(size: int): + return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) + + +class RubinFusedMultiHeadAttentionForward: + """K3 context FMHA using Rubin SM107 FP8 MMA instructions.""" + + arch_str: str = "sm_107" + arch_name: str = "Rubin SM107" + supports_split_k_epilogue_merge: bool = True + _MMA_INST_K: int = 64 + + def __init__( + self, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + mma_tiler: Tuple[int, int], + head_dim: Union[int, Tuple[int, int]], + is_persistent: bool, + mask_type: fmha_utils.MaskEnum, + enable_ex2_emulation: bool, + enable_skip_correction: bool, + use_tma_store: bool = True, + split_k_dims: Optional[Tuple[int, int]] = None, + ): + """Initializes the configuration for a Rubin Fused Multi-Head Attention (FMHA) kernel. + + This configuration includes several key aspects: + + 1. Data Type Settings: + - qk_acc_dtype: Data type for Q*K^T matrix multiplication accumulator + - pv_acc_dtype: Data type for P*V matrix multiplication accumulator + + 2. MMA Instruction Settings: + - mma_tiler: The shape of the MMA instruction unit: (M, N) for BMM1 and (M, K) for BMM2 + - head_dim: The head dimension, it can be a single integer or a tuple of two integers (D, Dv). + If it is a tuple, Dv is the head dimension of the value & output tensors. + It also determines the K dimension of the BMM1's MMA instruction unit + & N dimension of the BMM2's MMA instruction unit. + - qk_mma_tiler: MMA shape for Q*K^T computation + - pv_mma_tiler: MMA shape for P*V computation + + 3. Kernel Execution Mode: + - is_persistent: Boolean indicating whether to use persistent kernel mode + - mask_type: Specifies the type of mask to use (no mask, residual mask, or causal mask) + - window_size_left/right: Sliding window size for attention masking + - enable_ex2_emulation: Whether to enable exp2 emulation + - enable_skip_correction: Whether to skip the correction when rowmax is not updated larger than a threshold + + :param qk_acc_dtype: Data type for Q*K^T matrix multiplication accumulator + :type qk_acc_dtype: Type[cutlass.Numeric] + :param pv_acc_dtype: Data type for P*V matrix multiplication accumulator + :type pv_acc_dtype: Type[cutlass.Numeric] + :param mma_tiler: The (M, N) shape of the MMA instruction + :type mma_tiler: Tuple[int, int] + :param head_dim: The head dimension, it can be a single integer or a tuple of two integers (D, Dv). + :type head_dim: Union[int, Tuple[int, int]] + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param mask_type: Type of mask to use + :type mask_type: fmha_utils.MaskEnum + :param window_size_left: Left-side sliding window size for attention masking + :type window_size_left: int + :param window_size_right: Right-side sliding window size for attention masking + :type window_size_right: int + """ + + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + if isinstance(head_dim, tuple): + self.head_dim = head_dim[0] + self.head_dim_v = head_dim[1] + assert self.head_dim == 192 and self.head_dim_v == 128, ( + f"When Headdim is a tuple, it's for MLA. Must be (192, 128), but got {head_dim}" + ) + else: + self.head_dim = head_dim + self.head_dim_v = head_dim + if split_k_dims is not None: + if ( + len(split_k_dims) != 2 + or min(split_k_dims) <= 0 + or sum(split_k_dims) != self.head_dim + ): + raise ValueError( + "split_k_dims must contain two positive dimensions that " + f"sum to head_dim={self.head_dim}, got {split_k_dims}" + ) + self.split_k_dims = split_k_dims + self.cta_tiler = ( + 2 * mma_tiler[0], # 2 O tile per CTA + mma_tiler[1], + self.head_dim_v, + ) + self.qk_mma_tiler = ( + *mma_tiler, + self.head_dim, + ) + self.pv_mma_tiler = ( + mma_tiler[0], + self.head_dim_v, + mma_tiler[1], + ) + self.cluster_shape_mn = (1, 1) + self.is_persistent = is_persistent + self.mask_type = mask_type + self.enable_skip_correction = enable_skip_correction + self.enable_ex2_emulation = enable_ex2_emulation + self.use_tma_store = use_tma_store + + self.softmax0_warp_ids = (0, 1, 2, 3) + self.softmax1_warp_ids = (4, 5, 6, 7) + self.correction_warp_ids = (8, 9, 10, 11) + self.mma_warp_id = 12 + self.load_warp_id = 13 + self.epilogue_warp_id = 14 + self.empty_warp_id = 15 + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols(self.arch_str) + + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.softmax0_warp_ids, + *self.softmax1_warp_ids, + *self.correction_warp_ids, + self.mma_warp_id, + self.load_warp_id, + self.epilogue_warp_id, + self.empty_warp_id, + ) + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_warp + * sum( + ( + len((self.mma_warp_id,)), + len(self.softmax0_warp_ids), + len(self.softmax1_warp_ids), + len(self.correction_warp_ids), + ) + ), + ) + self.sequence_s0_s1_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=self.threads_per_warp + * len((*self.softmax0_warp_ids, *self.softmax1_warp_ids)), + ) + self.sequence_s1_s0_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp + * len((*self.softmax0_warp_ids, *self.softmax1_warp_ids)), + ) + self.s0_warpgroup_barrier = pipeline.NamedBarrier( + barrier_id=5, + num_threads=self.threads_per_warp * len(self.softmax0_warp_ids), + ) + self.s1_warpgroup_barrier = pipeline.NamedBarrier( + barrier_id=6, + num_threads=self.threads_per_warp * len(self.softmax1_warp_ids), + ) + self.tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=7, + num_threads=self.threads_per_warp * len(self.correction_warp_ids), + ) + + self.tmem_s0_offset = 0 + self.tmem_s1_offset = 128 + self.tmem_o0_offset = 256 + self.tmem_o1_offset = 384 + # inplaced with s1 + self.tmem_p0_offset = 160 + # inplaced with s0 + self.tmem_p1_offset = 32 + # vec buffer for row_max & row_sum + # inplaced with s0 + self.tmem_vec0_offset = 0 + # inplaced with s1 + self.tmem_vec1_offset = 128 + # skip mma pv flag offset regarding to the vec buffer + # inplaced with s1 + self.tmem_skip_softmax0_offset = 136 + # inplaced with s0 + self.tmem_skip_softmax1_offset = 8 + + self.num_regs_softmax = 192 + self.num_regs_correction = 96 + self.num_regs_other = 32 + self.buffer_align_bytes = 1024 + self.arch = BaseDSL._get_dsl().get_arch_enum() + + if self.arch >= Arch.sm_103: + assert not self.enable_ex2_emulation, ( + f"Don't enable exp2 emulation for {self.arch}, it doesn't help performance" + ) + + num_warps_per_warpgroup = 4 + self.softmax_warpgroup_count = ( + len((*self.softmax0_warp_ids, *self.softmax1_warp_ids)) // num_warps_per_warpgroup + ) + + def _make_qk_tiled_mma(self, cta_group): + return utils.sm107.make_trivial_tiled_mma( + self.q_dtype, + self.k_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + (*self.qk_mma_tiler[:2], self.mma_inst_k), + ) + + def _make_pv_tiled_mma(self, cta_group, p_major_mode, p_source): + return utils.sm107.make_trivial_tiled_mma( + self.v_dtype, + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + (*self.pv_mma_tiler[:2], self.pv_mma_inst_k), + a_source=p_source, + ) + + def _make_split_k_smem_layout(self, tiled_mma, k_dim): + """Build a split-K layout with the full-D atom and physical swizzle.""" + full_shape = tiled_mma.partition_shape_B( + cute.dice( + (*self.qk_mma_tiler[:2], self.head_dim), + (None, 1, 1), + ) + ) + full_shape_nk = ( + cute.size(full_shape[0][0]) * full_shape[1], + cute.size(full_shape[0][1]) * full_shape[2], + ) + full_atom_kind = tcgen05_utils.get_smem_layout_atom_ab( + OperandMajorMode.K, self.k_dtype, full_shape_nk + ) + full_atom = tcgen05.make_smem_layout_atom(full_atom_kind, self.k_dtype) + split_shape = tiled_mma.partition_shape_B( + cute.dice((*self.qk_mma_tiler[:2], k_dim), (None, 1, 1)) + ) + return tcgen05.tile_to_mma_shape( + full_atom, + cute.append(split_shape, 1), + order=(1, 2, 3), + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the FMHA kernel operation. + + This method initializes and configures various attributes required for the + execution of the fused multi-head attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.q_stage = 2 + k_stage = 4 if self.q_dtype.width == 8 else 3 + v_stage = 4 if self.v_dtype.width == 8 else 3 + self.kv_stage = min(k_stage, v_stage) + # For D192, the smem usage of Q & K is larger. So, we need to reduce the stage count. + if self.head_dim == 192 and self.q_dtype.width == 16: + self.kv_stage = 2 + self.p_mma_stage = 1 + self.acc_stage = 1 + self.softmax_corr_stage = 1 + self.mma_corr_stage = 2 + self.mma_softmax_stage = 1 + self.epi_stage = 2 + + # Tunable parameters + self.rescale_threshold = 8.0 if self.enable_skip_correction else 0.0 + # FP8 P pre-scale: offset added to exp2 exponent so that P*2^offset fills + # more of E4M3's [0, 448] range, improving quantization precision. + # Derived from rescale_threshold to guarantee P*2^offset <= 448. + self.p_fp8_prescale_log2 = max(0.0, math.floor(math.log2(448) - self.rescale_threshold)) + # ln(2) * offset correction for LSE when pre-scale is active + self.p_fp8_prescale_lse_correction = self.p_fp8_prescale_log2 * math.log(2) + # Optional double buffering for correction rescale. + self.enable_correction_double_buffer = False + + is_f16_bf16_qk = self.q_dtype in {cutlass.Float16, cutlass.BFloat16} + is_f16_bf16_pv = self.v_dtype in {cutlass.Float16, cutlass.BFloat16} + if not (is_f16_bf16_qk or is_f16_bf16_pv): + self.kv_stage = 6 + self.mma_inst_k = ( + 16 if is_f16_bf16_qk else (32 if self.head_dim == 32 else self._MMA_INST_K) + ) + self.pv_mma_inst_k = 16 if is_f16_bf16_pv else self._MMA_INST_K + # Always keep the sequence barrier on and disable ex2 emulation on Rubin. + self.enable_sequence_barrier = True + self.enable_ex2_emulation = False + + @cute.jit + def __call__( + self, + q_tensor: cute.Tensor, + k_tensor: cute.Tensor, + k_pe_tensor: Optional[cute.Tensor], + v_tensor: cute.Tensor, + o_tensor: cute.Tensor, + problem_size: Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Int32, Int32], + cum_seqlen_q: Optional[cute.Tensor], + cum_seqlen_k: Optional[cute.Tensor], + lse_tensor: Optional[cute.Tensor], + sink_tensor: Optional[cute.Tensor], + merge_stats_tensor: Optional[cute.Tensor], + merge_op_tensor: Optional[cute.Tensor], + scale_softmax_log2: Float32, + scale_softmax: Float32, + scale_output: Float32, + skip_softmax_threshold_log2: Optional[Float32], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + skip_softmax_count: Optional[cute.Tensor], + total_softmax_count: Optional[cute.Tensor], + stream: cuda.CUstream, + use_pdl: bool, + ): + """Execute the Fused Multi-Head Attention operation on the provided tensors. + + This method prepares the input tensors for processing, validates their shapes and types, + configures the computation parameters, and launches the CUDA kernel. + + The method handles: + 1. Tensor layout transformations for specific memory access patterns + 2. Validation of tensor shapes and data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch with appropriate parameters + + :param q_tensor: The query tensor with shape (b, s_q, h_k, h_r, d) + :type q_tensor: cute.Tensor + :param k_tensor: The key tensor with shape (b, s_k, h_k, 1, d) + :type k_tensor: cute.Tensor + :param v_tensor: The value tensor with shape (b, s_v, h_k, 1, dv) + :type v_tensor: cute.Tensor + :param o_tensor: The output tensor with shape (b, s_q, h_k, h_r, dv) + :type o_tensor: cute.Tensor + :param problem_size: The problem size with shape + [b, s_q_max, s_lse_max, s_k_max, h_q, h_k, d, dv]. With cumulative + sequence lengths, s_q_max and s_k_max are the per-batch maxima. + :type problem_size: Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Int32, Int32] + :param cum_seqlen_q: The cumulative sequence length tensor for query + :type cum_seqlen_q: Optional[cute.Tensor] + :param cum_seqlen_k: The cumulative sequence length tensor for key + :type cum_seqlen_k: Optional[cute.Tensor] + :param scale_softmax_log2: The log2 scale factor for softmax + :type scale_softmax_log2: Float32 + :param scale_softmax: The scale factor for softmax + :type scale_softmax: Float32 + :param scale_output: The scale factor for the output + :type scale_output: Float32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + :raises TypeError: If tensor data types don't match or aren't supported + :raises RuntimeError: If tensor layouts aren't in supported formats + """ + b, s_q_max, s_lse_max, s_k_max, h_q, h_k, d, dv = problem_size + h_r = h_q // h_k + # setup static attributes before smem/grid/tma computation + self.q_dtype = q_tensor.element_type + self.k_dtype = k_tensor.element_type + self.v_dtype = v_tensor.element_type + self.o_dtype = o_tensor.element_type + if cutlass.const_expr(self.split_k_dims is None): + if cutlass.const_expr(k_pe_tensor is not None): + raise ValueError("k_pe_tensor requires split_k_dims") + else: + if cutlass.const_expr(k_pe_tensor is None): + raise ValueError("split_k_dims requires k_pe_tensor") + if cutlass.const_expr(k_pe_tensor.element_type != self.k_dtype): + raise TypeError( + f"Split K dtype mismatch: {self.k_dtype} != {k_pe_tensor.element_type}" + ) + if cutlass.const_expr((merge_stats_tensor is None) != (merge_op_tensor is None)): + raise ValueError("merge_stats_tensor and merge_op_tensor must be paired") + if cutlass.const_expr(merge_stats_tensor is not None): + if cutlass.const_expr(self.split_k_dims is None): + raise ValueError("attention merge requires split K") + if cutlass.const_expr(self.use_tma_store): + raise ValueError("attention merge requires direct global store") + if cutlass.const_expr(lse_tensor is not None or sink_tensor is not None): + raise ValueError("attention merge does not support LSE or sink") + if cutlass.const_expr( + self.q_dtype.width != self.v_dtype.width + or self.q_dtype.width not in {8, 16} + or self.o_dtype.width != 16 + or self.head_dim_v != 128 + ): + raise ValueError( + "attention merge requires matching 8-bit or 16-bit Q/K/V, " + "16-bit O, and value head dim 128" + ) + if cutlass.const_expr(merge_stats_tensor.element_type is not Float32): + raise TypeError("merge_stats_tensor must use Float32") + if cutlass.const_expr(merge_op_tensor.element_type is not Int64): + raise TypeError("merge_op_tensor must use Int64") + + # s_q, s_k, s_v are the actual tensor dimensions (total seqlen for varlen) + s_q = q_tensor.shape[1] + s_k = k_tensor.shape[1] + s_v = v_tensor.shape[1] + s_lse = s_lse_max + # Important for performance + align = 256 // self.o_dtype.width + d = cute.assume(Int32(d), align) + dv = cute.assume(Int32(dv), align) + + stride_b_q = h_r * h_k * s_q * d if cum_seqlen_q is None else 0 + stride_b_o = h_r * h_k * s_q * dv if cum_seqlen_q is None else 0 + k_storage_d = self.split_k_dims[0] if self.split_k_dims is not None else d + stride_b_k = h_k * s_k * k_storage_d if cum_seqlen_k is None else 0 + stride_b_k_pe = ( + s_k * self.split_k_dims[1] + if cum_seqlen_k is None and self.split_k_dims is not None + else 0 + ) + stride_b_v = h_k * s_v * dv if cum_seqlen_k is None else 0 + stride_b_lse = h_r * h_k * s_lse if cum_seqlen_q is None else 0 + + # (b, s_q, h_k, h_r, d) -> (s_q, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=(d * h_r * h_k, 1, ((d, d * h_r), stride_b_q)), + ) + q = cute.make_tensor(q_tensor.iterator, q_layout) + # (b, s_k, h_k, 1, d) -> (s_k, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + if cutlass.const_expr(self.split_k_dims is not None): + # Split K and V can be views into a combined B-proj output. Preserve + # their runtime token/head strides instead of assuming compact + # storage; the innermost dimension remains contiguous for TMA. + k_stride = ( + k_tensor.stride[1], + k_tensor.stride[4], + ((0, k_tensor.stride[2]), stride_b_k), + ) + else: + k_stride = ( + k_storage_d * h_k, + 1, + ((0, k_storage_d), stride_b_k), + ) + k_layout = cute.make_layout( + (s_k, k_storage_d, ((h_r, h_k), b)), + stride=k_stride, + ) + k = cute.make_tensor(k_tensor.iterator, k_layout) + k_pe = None + if cutlass.const_expr(self.split_k_dims is not None): + # K-pe is shared by all query heads without materializing expand(). + k_pe_layout = cute.make_layout( + (s_k, self.split_k_dims[1], ((h_r, h_k), b)), + stride=( + k_pe_tensor.stride[1], + k_pe_tensor.stride[4], + ((0, 0), stride_b_k_pe), + ), + ) + k_pe = cute.make_tensor(k_pe_tensor.iterator, k_pe_layout) + # (b, s_v, h_k, 1, dv) -> (dv, s_v, ((h_r, h_k), b)), 0-stride for h_r to broadcast + if cutlass.const_expr(self.split_k_dims is not None): + v_stride = ( + v_tensor.stride[4], + v_tensor.stride[1], + ((0, v_tensor.stride[2]), stride_b_v), + ) + else: + v_stride = (1, dv * h_k, ((0, dv), stride_b_v)) + v_layout = cute.make_layout( + (dv, s_v, ((h_r, h_k), b)), + stride=v_stride, + ) + v = cute.make_tensor(v_tensor.iterator, v_layout) + # (b, s_q, h_k, h_r, dv) -> (s_q, dv, ((h_r, h_k), b)) + o_layout = cute.make_layout( + (s_q, dv, ((h_r, h_k), b)), + stride=(dv * h_r * h_k, 1, ((dv, dv * h_r), stride_b_o)), + ) + o = cute.make_tensor(o_tensor.iterator, o_layout) + if cutlass.const_expr(lse_tensor is not None): + # (s, ((h_r, h_k), b)) - head stride=1 to match FlashInfer (total_q, h_q) convention + lse_layout = cute.make_layout( + (s_lse, ((h_r, h_k), b)), + stride=(h_r * h_k, ((1, h_r), stride_b_lse)), + ) + lse = cute.make_tensor(lse_tensor.iterator, lse_layout) + else: + lse = None + + if cutlass.const_expr(sink_tensor is not None): + # sink_tensor is 1D with shape (h_q,) = (h_k * h_r,) + # Create layout ((h_r, h_k), b) with stride 0 for batch so blk_coord[2] works + sink_layout = cute.make_layout( + ((h_r, h_k), b), + stride=((1, h_r), 0), + ) + sink = cute.make_tensor(sink_tensor.iterator, sink_layout) + else: + sink = None + + if cutlass.const_expr(merge_stats_tensor is not None): + stride_b_stats = h_r * h_k * s_lse * 2 if cum_seqlen_q is None else 0 + merge_stats_layout = cute.make_layout( + (s_lse, ((h_r, h_k), b), 2), + stride=( + h_r * h_k * 2, + ((2, 2 * h_r), stride_b_stats), + 1, + ), + ) + merge_stats = cute.make_tensor(merge_stats_tensor.iterator, merge_stats_layout) + merge_op = merge_op_tensor + else: + merge_stats = None + merge_op = None + + self.tile_sched_params, grid = fmha_utils.compute_grid( + cute.shape((s_q_max, d, ((h_r, h_k), b))), + self.cta_tiler, + self.is_persistent, + ) + self.q_major_mode = cutlass.tensor_utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = cutlass.tensor_utils.LayoutEnum.from_tensor(k).mma_major_mode() + k_pe_major_mode = self.k_major_mode + if cutlass.const_expr(self.split_k_dims is not None): + k_pe_major_mode = cutlass.tensor_utils.LayoutEnum.from_tensor(k_pe).mma_major_mode() + self.v_major_mode = cutlass.tensor_utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = cutlass.tensor_utils.LayoutEnum.from_tensor(o) + + if cutlass.const_expr(self.q_major_mode != OperandMajorMode.K): + raise RuntimeError("The layout of q is not supported") + if cutlass.const_expr(self.k_major_mode != OperandMajorMode.K): + raise RuntimeError("The layout of k is not supported") + if cutlass.const_expr( + self.split_k_dims is not None and k_pe_major_mode != OperandMajorMode.K + ): + raise RuntimeError("The layout of k_pe is not supported") + if cutlass.const_expr(self.v_major_mode != OperandMajorMode.MN): + raise RuntimeError("The layout of v is not supported") + + # check type consistency: Q and K must share the same dtype (qk_dtype); + # V may use a different dtype (pv_dtype) to allow qk_dtype != pv_dtype. + if cutlass.const_expr(self.q_dtype != self.k_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}") + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.ONE + # the intermediate tensor p is from tmem & k-major + p_source = tcgen05.OperandSource.TMEM + p_major_mode = cute.nvgpu.OperandMajorMode.K + qk_tiled_mma = self._make_qk_tiled_mma(cta_group) + pv_tiled_mma = self._make_pv_tiled_mma(cta_group, p_major_mode, p_source) + + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.pv_mma_tiler[:2] + + q_smem_layout_staged = tcgen05_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = tcgen05_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.k_dtype, + self.kv_stage, + ) + p_tmem_layout_staged = tcgen05_utils.make_smem_layout_a( + pv_tiled_mma, + self.pv_mma_tiler, + self.v_dtype, + self.acc_stage, + ) + v_smem_layout_staged_origin = tcgen05_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.v_dtype, + self.kv_stage, + ) + # k & v share the same smem buffer. Pad the smaller-tile operand's stage stride to + # match the larger tile. Stride is scaled to the target operand's element width. + # sK_cosize covers all kv_stage slots at the larger tile's byte footprint. + if cutlass.const_expr(self.k_dtype.width >= self.v_dtype.width): + # k tile >= v tile: give v k's stage stride in v_dtype units + k_tile_cosize = cute.cosize(cute.select(k_smem_layout_staged, mode=[0, 1, 2])) + v_stage_stride = k_tile_cosize * self.k_dtype.width // self.v_dtype.width + v_smem_layout_staged = cute.append( + cute.select(v_smem_layout_staged_origin, mode=[0, 1, 2]), + cute.make_layout(self.kv_stage, stride=v_stage_stride), + ) + sK_cosize = cute.cosize(k_smem_layout_staged) + else: + # v tile > k tile: give k v's stage stride in k_dtype units + v_tile_cosize = cute.cosize(cute.select(v_smem_layout_staged_origin, mode=[0, 1, 2])) + k_stage_stride = v_tile_cosize * self.v_dtype.width // self.k_dtype.width + k_smem_layout_staged = cute.append( + cute.select(k_smem_layout_staged, mode=[0, 1, 2]), + cute.make_layout(self.kv_stage, stride=k_stage_stride), + ) + v_smem_layout_staged = v_smem_layout_staged_origin + # cute.cosize gives (kv_stage-1)*k_stage_stride + k_tile_cosize, but the + # last stage must also accommodate a full V tile, so size for kv_stage slots. + sK_cosize = self.kv_stage * k_stage_stride + + k_nope_smem_layout_staged = None + k_pe_smem_layout_staged = None + if cutlass.const_expr(self.split_k_dims is not None): + k_nope_dim, k_pe_dim = self.split_k_dims + k_nope_smem_layout = self._make_split_k_smem_layout(qk_tiled_mma, k_nope_dim) + k_pe_smem_layout = self._make_split_k_smem_layout(qk_tiled_mma, k_pe_dim) + k_nope_tile_cosize = cute.cosize(k_nope_smem_layout) + k_pe_tile_cosize = cute.cosize(k_pe_smem_layout) + k_stage_stride = cute.cosize(cute.select(k_smem_layout_staged, mode=[0, 1, 2])) + if cutlass.const_expr(k_nope_tile_cosize + k_pe_tile_cosize != k_stage_stride): + raise ValueError("Split K shared layouts do not cover the full K tile") + k_nope_smem_layout_staged = cute.append( + cute.select(k_nope_smem_layout, mode=[0, 1, 2]), + cute.make_layout(self.kv_stage, stride=k_stage_stride), + ) + k_pe_smem_layout_staged = cute.append( + cute.select(k_pe_smem_layout, mode=[0, 1, 2]), + cute.make_layout(self.kv_stage, stride=k_stage_stride), + ) + o_smem_layout_staged = tcgen05_utils.make_smem_layout_epi( + self.o_dtype, + self.o_layout, + self.epi_tile, + self.epi_stage, + ) + + # TMA load for Q + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + tma_store_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k_pe = None + tma_tensor_k_pe = None + if cutlass.const_expr(self.split_k_dims is not None): + k_nope_dim, k_pe_dim = self.split_k_dims + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + k, + cute.select(k_nope_smem_layout_staged, mode=[0, 1, 2]), + (*self.qk_mma_tiler[:2], k_nope_dim), + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_k_pe, tma_tensor_k_pe = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + k_pe, + cute.select(k_pe_smem_layout_staged, mode=[0, 1, 2]), + (*self.qk_mma_tiler[:2], k_pe_dim), + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + else: + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + v, + v_smem_layout, + self.pv_mma_tiler, + pv_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + o_smem_layout = cute.select(o_smem_layout_staged, mode=[0, 1]) + tma_atom_o, tma_tensor_o = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_store_op, + o, + o_smem_layout, + self.epi_tile, + ) + + tma_atom_old_o = None + tma_tensor_old_o = None + self.tma_copy_old_o_bytes = 0 + if cutlass.const_expr(merge_stats is not None): + if cutlass.const_expr(self.epi_tile != (128, 128) or self.head_dim_v != 128): + raise ValueError("old-O TMA load requires a 128x128 output tile") + self.tma_copy_old_o_bytes = cute.size_in_bytes(self.o_dtype, o_smem_layout) + if cutlass.const_expr(self.tma_copy_old_o_bytes != 128 * 128 * 2): + raise ValueError("old-O TMA load must transfer exactly 32768 bytes") + tma_atom_old_o, tma_tensor_old_o = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_op, + o, + o_smem_layout, + self.epi_tile, + ) + + q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) + k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) + if cutlass.const_expr(self.split_k_dims is not None): + k_copy_size = cute.size_in_bytes( + self.k_dtype, + cute.select(k_nope_smem_layout_staged, mode=[0, 1, 2]), + ) + cute.size_in_bytes( + self.k_dtype, + cute.select(k_pe_smem_layout_staged, mode=[0, 1, 2]), + ) + v_copy_size = cute.size_in_bytes(self.v_dtype, v_smem_layout) + self.tma_copy_q_bytes = q_copy_size + self.tma_copy_k_bytes = k_copy_size + self.tma_copy_v_bytes = v_copy_size + + @cute.struct + class SharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[Int64, self.q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_stage * 2] + mma_s0_mbar_ptr: cute.struct.MemRange[Int64, self.mma_softmax_stage * 2] + mma_s1_mbar_ptr: cute.struct.MemRange[Int64, self.mma_softmax_stage * 2] + p0_mma_mbar_ptr: cute.struct.MemRange[Int64, self.p_mma_stage * 2] + p1_mma_mbar_ptr: cute.struct.MemRange[Int64, self.p_mma_stage * 2] + s0_corr_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_corr_stage * 2] + s1_corr_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_corr_stage * 2] + corr_epi_mbar_ptr: cute.struct.MemRange[Int64, self.epi_stage * 2] + mma_corr_mbar_ptr: cute.struct.MemRange[Int64, self.mma_corr_stage * 2] + s0_p1_inplace_barrier_ptr: cute.struct.MemRange[Int64, self.p_mma_stage * 2] + s1_p0_inplace_barrier_ptr: cute.struct.MemRange[Int64, self.p_mma_stage * 2] + # Independent barriers prevent O1 from overtaking a slow O0 waiter. + old_o_tma_mbar_ptr: cute.struct.MemRange[Int64, 2] + # Tmem holding buffer + tmem_holding_buf: Int32 + # Smem tensors + sO: cute.struct.Align[ + cute.struct.MemRange[self.o_dtype, cute.cosize(o_smem_layout_staged)], + self.buffer_align_bytes, + ] + sQ: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(q_smem_layout_staged)], + self.buffer_align_bytes, + ] + sK: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, sK_cosize], + self.buffer_align_bytes, + ] + # Skip softmax and PV warpgroup votes + s0_warp_wants_skip_softmax_exchange: cute.struct.MemRange[Int8, 4] + s1_warp_wants_skip_softmax_exchange: cute.struct.MemRange[Int8, 4] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_k_pe, + tma_tensor_k_pe, + tma_atom_v, + tma_tensor_v, + tma_atom_old_o, + tma_tensor_old_o, + tma_atom_o, + tma_tensor_o, + o, + cum_seqlen_q, + cum_seqlen_k, + lse, + sink, + merge_stats, + merge_op, + scale_softmax_log2, + scale_softmax, + scale_output, + skip_softmax_threshold_log2, + window_size_left, + window_size_right, + q_smem_layout_staged, + k_smem_layout_staged, + p_tmem_layout_staged, + v_smem_layout_staged, + o_smem_layout_staged, + skip_softmax_count, + total_softmax_count, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + use_pdl=use_pdl, + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_k_pe: Optional[cute.CopyAtom], + mK_pe_kdl: Optional[cute.Tensor], + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + tma_atom_old_o: Optional[cute.CopyAtom], + mOldO_qdl: Optional[cute.Tensor], + tma_atom_o: cute.CopyAtom, + mO_qdl: cute.Tensor, + mO: cute.Tensor, + cum_seqlen_q: Optional[cute.Tensor], + cum_seqlen_k: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mSink: Optional[cute.Tensor], + mMergeStats: Optional[cute.Tensor], + mMergeOp: Optional[cute.Tensor], + scale_softmax_log2: Float32, + scale_softmax: Float32, + scale_output: Float32, + skip_softmax_threshold_log2: Optional[Float32], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + p_tmem_layout_staged: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + o_smem_layout_staged: cute.ComposedLayout, + skip_softmax_count: Optional[cute.Tensor], + total_softmax_count: Optional[cute.Tensor], + tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams, + ): + """The device kernel implementation of the Fused Multi-Head Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the FMHA computation: + 1. Load warp: Loads Q, K, V data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Softmax warps: Compute softmax normalization on attention scores + 4. Correction warps: Apply adjustments to intermediate results + 5. Epilogue warp: Handles final output transformation and storage + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases, and optional attention masking. + + :param qk_tiled_mma: Tiled MMA for Q*K^T + :type qk_tiled_mma: cute.TiledMma + :param pv_tiled_mma: Tiled MMA for P*V + :type pv_tiled_mma: cute.TiledMma + :param tma_atom_q: TMA copy atom for query tensor + :type tma_atom_q: cute.CopyAtom + :param mQ_qdl: Partitioned query tensor + :type mQ_qdl: cute.Tensor + :param tma_atom_k: TMA copy atom for key tensor + :type tma_atom_k: cute.CopyAtom + :param mK_kdl: Partitioned key tensor + :type mK_kdl: cute.Tensor + :param tma_atom_v: TMA copy atom for value tensor + :type tma_atom_v: cute.CopyAtom + :param mV_dkl: Partitioned value tensor + :type mV_dkl: cute.Tensor + :param tma_atom_o: TMA copy atom for output tensor + :type tma_atom_o: cute.CopyAtom + :param mO_qdl: Partitioned output tensor + :type mO_qdl: cute.Tensor + :param mO: Non-partitioned output tensor + :type mO: cute.Tensor + :param scale_softmax_log2: The log2 scale factor for softmax + :type scale_softmax_log2: Float32 + :param scale_output: The scale factor for the output + :type scale_output: Float32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param q_smem_layout_staged: Shared memory layout for query tensor + :type q_smem_layout_staged: cute.ComposedLayout + :param k_smem_layout_staged: Shared memory layout for key tensor + :type k_smem_layout_staged: cute.ComposedLayout + :param p_tmem_layout_staged: Tensor memory layout for probability matrix + :type p_tmem_layout_staged: cute.ComposedLayout + :param v_smem_layout_staged: Shared memory layout for value tensor + :type v_smem_layout_staged: cute.ComposedLayout + :param o_smem_layout_staged: Shared memory layout for output tensor + :type o_smem_layout_staged: cute.ComposedLayout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Prefetch tma desc + # + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + if cutlass.const_expr(self.split_k_dims is not None): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k_pe) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + if cutlass.const_expr(tma_atom_old_o is not None): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_old_o) + if cutlass.const_expr(self.use_tma_store): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_o) + + # Alloc + smem = cutlass.memory.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + load_kv_producer, load_kv_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.kv_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_k_bytes, + barrier_storage=storage.load_kv_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + mma_s0_producer, mma_s0_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_softmax_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax0_warp_ids) + ), + barrier_storage=storage.mma_s0_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + mma_s1_producer, mma_s1_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_softmax_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax1_warp_ids) + ), + barrier_storage=storage.mma_s1_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + p0_mma_producer, p0_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.p_mma_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax0_warp_ids) + ), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.p0_mma_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + p1_mma_producer, p1_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.p_mma_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax1_warp_ids) + ), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.p1_mma_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + s0_corr_producer, s0_corr_consumer = pipeline.PipelineAsync.create( + num_stages=self.softmax_corr_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len((*self.softmax0_warp_ids, self.mma_warp_id)) + ), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.correction_warp_ids) + ), + barrier_storage=storage.s0_corr_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + s1_corr_producer, s1_corr_consumer = pipeline.PipelineAsync.create( + num_stages=self.softmax_corr_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len((*self.softmax1_warp_ids, self.mma_warp_id)) + ), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.correction_warp_ids) + ), + barrier_storage=storage.s1_corr_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + corr_epi_producer, corr_epi_consumer = pipeline.PipelineAsync.create( + num_stages=self.epi_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.correction_warp_ids) + ), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len([self.epilogue_warp_id]) + ), + barrier_storage=storage.corr_epi_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + mma_corr_producer, mma_corr_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_corr_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.correction_warp_ids) + ), + barrier_storage=storage.mma_corr_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + s0_p1_inplace_producer, s0_p1_inplace_consumer = pipeline.PipelineAsync.create( + num_stages=self.p_mma_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax0_warp_ids) + ), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax1_warp_ids) + ), + barrier_storage=storage.s0_p1_inplace_barrier_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + s1_p0_inplace_producer, s1_p0_inplace_consumer = pipeline.PipelineAsync.create( + num_stages=self.p_mma_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax0_warp_ids) + ), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax1_warp_ids) + ), + barrier_storage=storage.s1_p0_inplace_barrier_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + old_o_tma_mbar_ptr = storage.old_o_tma_mbar_ptr.data_ptr() + if cutlass.const_expr(tma_atom_old_o is not None): + if warp_idx == self.load_warp_id: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(old_o_tma_mbar_ptr, 1) + cute.arch.mbarrier_init(old_o_tma_mbar_ptr + 1, 1) + tmem = cutlass.memory.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + # Correction warp is the last one that accesses tmem + allocator_warp_id=self.correction_warp_ids[0], + arch=self.arch_str, + ) + pipeline_init_arrive(is_relaxed=True) + + # Generate smem tensor Q/K/V/O + # (MMA, MMA_Q, MMA_D, PIPE) + sQ = storage.sQ.get_tensor(q_smem_layout_staged.outer, swizzle=q_smem_layout_staged.inner) + # (MMA, MMA_K, MMA_D, PIPE) + sK = storage.sK.get_tensor(k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner) + sK_nope = sK + sK_pe = None + if cutlass.const_expr(self.split_k_dims is not None): + k_nope_dim, k_pe_dim = self.split_k_dims + k_nope_smem_layout = self._make_split_k_smem_layout(qk_tiled_mma, k_nope_dim) + k_pe_smem_layout = self._make_split_k_smem_layout(qk_tiled_mma, k_pe_dim) + k_stage_stride = cute.cosize(cute.select(k_smem_layout_staged, mode=[0, 1, 2])) + k_nope_smem_layout_staged = cute.append( + cute.select(k_nope_smem_layout, mode=[0, 1, 2]), + cute.make_layout(self.kv_stage, stride=k_stage_stride), + ) + k_pe_smem_layout_staged = cute.append( + cute.select(k_pe_smem_layout, mode=[0, 1, 2]), + cute.make_layout(self.kv_stage, stride=k_stage_stride), + ) + k_pe_smem_offset = cute.cosize(k_nope_smem_layout) + sK_nope = storage.sK.get_tensor( + k_nope_smem_layout_staged.outer, + swizzle=k_nope_smem_layout_staged.inner, + ) + sK_pe_ptr = cute.recast_ptr( + storage.sK.data_ptr() + k_pe_smem_offset, + k_pe_smem_layout_staged.inner, + ) + sK_pe = cute.make_tensor(sK_pe_ptr, k_pe_smem_layout_staged.outer) + # (MMA, MMA_K, MMA_D, PIPE) + # Reuse k's smem buffer for v. Recast element type so MMA descriptor matches v_dtype. + sV_ptr = cute.recast_ptr( + cute.recast_ptr(sK.iterator, dtype=self.v_dtype), + v_smem_layout_staged.inner, + ) + sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) + sO = storage.sO.get_tensor(o_smem_layout_staged.outer, swizzle=o_smem_layout_staged.inner) + s0_warp_wants_skip_softmax_exchange = ( + storage.s0_warp_wants_skip_softmax_exchange.get_tensor(cute.make_layout((4,))) + ) + s1_warp_wants_skip_softmax_exchange = ( + storage.s1_warp_wants_skip_softmax_exchange.get_tensor(cute.make_layout((4,))) + ) + + qk_thr_mma = qk_tiled_mma.get_slice(0) # default 1sm + pv_thr_mma = pv_tiled_mma.get_slice(0) # default 1sm + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tSrK = qk_thr_mma.make_fragment_B(sK) + tOrV = pv_thr_mma.make_fragment_B(sV) + + def make_tmem_tensors( + self, + qk_thr_mma: cute.TiledMma, + pv_thr_mma: cute.TiledMma, + p_tmem_layout_staged: cute.Layout, + tmem_ptr: cute.Pointer, + ): + qk_acc_shape = qk_thr_mma.partition_shape_C( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + tStS_fake = qk_thr_mma.make_fragment_C(qk_acc_shape) + tStS = cute.make_tensor(tmem_ptr + self.tmem_s0_offset, tStS_fake.layout) + pv_acc_shape = pv_thr_mma.partition_shape_C( + (self.pv_mma_tiler[0], self.pv_mma_tiler[1]) + ) + tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape) + tStS0 = cute.make_tensor(tmem_ptr + self.tmem_s0_offset, tStS.layout) + tStS1 = cute.make_tensor(tmem_ptr + self.tmem_s1_offset, tStS.layout) + tOtO0 = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout) + tOtO1 = cute.make_tensor(tmem_ptr + self.tmem_o1_offset, tOtO.layout) + tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) + tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] + tOrP0 = cute.make_tensor( + cute.recast_ptr( + tmem_ptr + self.tmem_p0_offset, + dtype=tOrP.dtype, + ), + tOrP.layout, + ) + tOrP1 = cute.make_tensor( + cute.recast_ptr( + tmem_ptr + self.tmem_p1_offset, + dtype=tOrP.dtype, + ), + tOrP.layout, + ) + return tStS, tStS0, tStS1, tOtO0, tOtO1, tOrP0, tOrP1 + + tile_sched = fmha_utils.create_fmha_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + pipeline_init_wait() + softmax_fn = partial( + self.softmax, + qk_thr_mma=qk_thr_mma, + value_args=( + mK_kdl.shape[0], + mQ_qdl.shape[0], + scale_softmax_log2, + skip_softmax_threshold_log2, + ), + mask_args=(window_size_left, window_size_right), + sched_args=(tile_sched, work_tile), + ) + # /////////////////////////////////////////////////////////////////////////////// + # EMPTY + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.empty_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + # /////////////////////////////////////////////////////////////////////////////// + # LOAD + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + cute.arch.griddepcontrol_wait() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = ( + not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + ) + if not continue_cond: + if cutlass.const_expr(cum_seqlen_k is not None): + # Chunk-gather metadata is int64; TMA coordinates and + # the existing sequence-length math remain int32. + seqlen_k = (cum_seqlen_k[batch_coord + 1] - cum_seqlen_k[batch_coord]).to( + Int32 + ) + continue_cond = seqlen_k <= 0 + if not continue_cond: + mQ_qdl_ = mQ_qdl + mK_kdl_ = mK_kdl + mK_pe_kdl_ = mK_pe_kdl + mV_dkl_ = mV_dkl + if cutlass.const_expr(cum_seqlen_q is not None): + mQ_qdl_ = cute.domain_offset( + (cum_seqlen_q[batch_coord], 0, ((0, 0), 0)), mQ_qdl + ) + if cutlass.const_expr(cum_seqlen_k is not None): + mK_kdl_ = cute.domain_offset( + (cum_seqlen_k[batch_coord].to(Int32), 0, ((0, 0), 0)), mK_kdl + ) + if cutlass.const_expr(self.split_k_dims is not None): + mK_pe_kdl_ = cute.domain_offset( + (cum_seqlen_k[batch_coord].to(Int32), 0, ((0, 0), 0)), + mK_pe_kdl, + ) + mV_dkl_ = cute.domain_offset( + (0, cum_seqlen_k[batch_coord].to(Int32), ((0, 0), 0)), mV_dkl + ) + # Local tile partition global tensors + gQ_qdl = cute.flat_divide(mQ_qdl_, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + tQgQ = tQgQ_qdl[None, None, 0, curr_block_coord[2]] + k_load_tiler = cute.select(self.qk_mma_tiler, mode=[1, 2]) + if cutlass.const_expr(self.split_k_dims is not None): + k_load_tiler = ( + self.qk_mma_tiler[1], + self.split_k_dims[0], + ) + gK_kdl = cute.flat_divide(mK_kdl_, k_load_tiler) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sK_nope, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + tKgK = tKgK_kdl[None, None, 0, curr_block_coord[2]] + tKsK_pe = None + tKgK_pe = None + if cutlass.const_expr(self.split_k_dims is not None): + gK_pe_kdl = cute.flat_divide( + mK_pe_kdl_, + (self.qk_mma_tiler[1], self.split_k_dims[1]), + ) + tSgK_pe_kdl = qk_thr_mma.partition_B(gK_pe_kdl) + tKsK_pe, tKgK_pe_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k_pe, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sK_pe, 0, 3), + cute.group_modes(tSgK_pe_kdl, 0, 3), + ) + tKgK_pe = tKgK_pe_kdl[None, None, 0, curr_block_coord[2]] + gV_dkl = cute.flat_divide(mV_dkl_, cute.select(self.pv_mma_tiler, mode=[1, 2])) + tSgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sV, 0, 3), + cute.group_modes(tSgV_dkl, 0, 3), + ) + tVgV = tVgV_dkl[None, 0, None, curr_block_coord[2]] + seqlen_kv_loop_start = fmha_utils.FusedMask.get_trip_start( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + ) + # Q0 + q0_coord = 2 * curr_block_coord[0] + q0_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, q0_coord], + tQsQ[None, q0_handle.index], + tma_bar_ptr=q0_handle.barrier, + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + # K0 + kv_coord = seqlen_kv_loop_start + k_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_coord], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + if cutlass.const_expr(self.split_k_dims is not None): + cute.copy( + tma_atom_k_pe, + tKgK_pe[None, kv_coord], + tKsK_pe[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + # Q1 + q1_coord = q0_coord + 1 + q1_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, q1_coord], + tQsQ[None, q1_handle.index], + tma_bar_ptr=q1_handle.barrier, + ) + kv_coord += 1 + + for i in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + # Ki + k_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_coord], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + if cutlass.const_expr(self.split_k_dims is not None): + cute.copy( + tma_atom_k_pe, + tKgK_pe[None, kv_coord], + tKsK_pe[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + # Vi-1 + v_handle = load_kv_producer.acquire_and_advance( + expected_tx=self.tma_copy_v_bytes + ) + cute.copy( + tma_atom_v, + tVgV[None, kv_coord - 1], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + kv_coord += 1 + # End of seqlen_kv loop + # Vi_end + v_handle = load_kv_producer.acquire_and_advance( + expected_tx=self.tma_copy_v_bytes + ) + cute.copy( + tma_atom_v, + tVgV[None, kv_coord - 1], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + # End of if not continue_cond + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + tStS, tStS0, tStS1, tOtO0, tOtO1, tOrP0, tOrP1 = make_tmem_tensors( + self, qk_thr_mma, pv_thr_mma, p_tmem_layout_staged, tmem_ptr + ) + enable_skip_softmax = skip_softmax_threshold_log2 is not None + tiled_tmem_load_v = None + tTMEM_LOADtS_v0, tTMEM_LOADtS_v1 = None, None + tTMEM_LOADrS_v0, tTMEM_LOADrS_v1 = None, None + if cutlass.const_expr(enable_skip_softmax): + cS = cute.make_identity_tensor(cute.select(self.qk_mma_tiler, mode=[0, 1])) + tScS = qk_thr_mma.partition_C(cS) + tStS_v = cute.composition(tStS, cute.make_layout((self.threads_per_warp, 1))) + tScS_v = cute.composition(tScS, cute.make_layout((self.threads_per_warp, 1))) + tmem_load_v_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(1)), + self.qk_acc_dtype, + ) + thread_idx = tidx % self.threads_per_warp + + tiled_tmem_load_v = tcgen05.make_tmem_copy(tmem_load_v_atom, tStS_v) + thr_tmem_load_v = tiled_tmem_load_v.get_slice(thread_idx) + tTMEM_LOADtS_v = thr_tmem_load_v.partition_S(tStS_v) + tTMEM_LOADcS_v = thr_tmem_load_v.partition_D(tScS_v) + tTMEM_LOADrS_v0 = cute.make_rmem_tensor(tTMEM_LOADcS_v.shape, self.qk_acc_dtype) + tTMEM_LOADrS_v1 = cute.make_rmem_tensor(tTMEM_LOADcS_v.shape, self.qk_acc_dtype) + tTMEM_LOADtS_v0 = cute.make_tensor( + tTMEM_LOADtS_v.iterator + self.tmem_skip_softmax0_offset, + tTMEM_LOADtS_v.layout, + ) + tTMEM_LOADtS_v1 = cute.make_tensor( + tTMEM_LOADtS_v.iterator + self.tmem_skip_softmax1_offset, + tTMEM_LOADtS_v.layout, + ) + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + continue_cond = False + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = ( + not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + ) + if not continue_cond: + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord].to(Int32) + seqlen_k = cum_seqlen_k[batch_coord + 1].to(Int32) - cuseqlen_k + continue_cond = seqlen_k <= 0 + if not continue_cond: + # Wait for Q0 + q0_handle = load_q_consumer.wait_and_advance() + tSrQ0 = tSrQ[None, None, None, q0_handle.index] + # Wait for K0 + k_handle = load_kv_consumer.wait_and_advance() + tSrK0 = tSrK[None, None, None, k_handle.index] + # GEMM_QK00 (Q0 * K0 -> S0) + mma_s0_producer, s0_corr_producer = self.mma_qk( + qk_tiled_mma, + (tSrQ0, tSrK0, tStS0), + (mma_s0_producer, s0_corr_producer), + ) + # Wait for Q1 + q1_handle = load_q_consumer.wait_and_advance() + tSrQ1 = tSrQ[None, None, None, q1_handle.index] + # GEMM_QK10 (Q1 * K0 -> S1), K0 is ready in GEMM_QK00 + mma_s1_producer, s1_corr_producer = self.mma_qk( + qk_tiled_mma, + (tSrQ1, tSrK0, tStS1), + (mma_s1_producer, s1_corr_producer), + ) + # Release K0 + k_handle.release() + # Note: Q0 & Q1 are still needed in the seqlen_kv loop + # so we need to release them after the seqlen_kv loop + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + # O1 hasn't been accumulated yet, its first MMA calculation doesn't need to accumulate + pv_whether_acc = False + for i in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + # Wait for Ki + k_handle = load_kv_consumer.wait_and_advance() + tSrKi = tSrK[None, None, None, k_handle.index] + # GEMM_QK0i (Q0 * Ki -> S0) + mma_s0_producer, s0_corr_producer = self.mma_qk( + qk_tiled_mma, + (tSrQ0, tSrKi, tStS0), + (mma_s0_producer, s0_corr_producer), + ) + # Wait for Vi-1 + v_handle = load_kv_consumer.wait_and_advance() + tOrVi = tOrV[None, None, None, v_handle.index] + # GEMM_PV0(i-1) (P0 * Vi-1 -> O0_partial) + mma_corr_producer, p0_mma_consumer = self.mma_pv( + pv_tiled_mma, + pv_whether_acc, + (tOrP0, tOrVi, tOtO0), + (mma_corr_producer, p0_mma_consumer), + ( + enable_skip_softmax, + tiled_tmem_load_v, + tTMEM_LOADtS_v0, + tTMEM_LOADrS_v0, + ), + ) + # GEMM_QK1i (Q1 * Ki -> S1) + mma_s1_producer, s1_corr_producer = self.mma_qk( + qk_tiled_mma, + (tSrQ1, tSrKi, tStS1), + (mma_s1_producer, s1_corr_producer), + ) + # Release Ki + k_handle.release() + # GEMM_PV1(i-1) (P1 * Vi-1 -> O1_partial) + mma_corr_producer, p1_mma_consumer = self.mma_pv( + pv_tiled_mma, + pv_whether_acc, + (tOrP1, tOrVi, tOtO1), + (mma_corr_producer, p1_mma_consumer), + ( + enable_skip_softmax, + tiled_tmem_load_v, + tTMEM_LOADtS_v1, + tTMEM_LOADrS_v1, + ), + ) + pv_whether_acc = True + # Release Vi-1 + v_handle.release() + # End of seqlen_kv loop + # release Q0 & Q1 + q0_handle.release() + q1_handle.release() + # Wait for Vi_end + v_handle = load_kv_consumer.wait_and_advance() + tOrVi = tOrV[None, None, None, v_handle.index] + # GEMM_PV0(i_end) (P0 * Vi_end -> O0) + mma_corr_producer, p0_mma_consumer = self.mma_pv( + pv_tiled_mma, + pv_whether_acc, + (tOrP0, tOrVi, tOtO0), + (mma_corr_producer, p0_mma_consumer), + ( + enable_skip_softmax, + tiled_tmem_load_v, + tTMEM_LOADtS_v0, + tTMEM_LOADrS_v0, + ), + ) + # GEMM_PV1(i_end) (P1 * Vi_end -> O1) + mma_corr_producer, p1_mma_consumer = self.mma_pv( + pv_tiled_mma, + pv_whether_acc, + (tOrP1, tOrVi, tOtO1), + (mma_corr_producer, p1_mma_consumer), + ( + enable_skip_softmax, + tiled_tmem_load_v, + tTMEM_LOADtS_v1, + tTMEM_LOADrS_v1, + ), + ) + # Release Vi_end + v_handle.release() + # Empty step for correction epilog + vec0_handle = s0_corr_producer.acquire_and_advance() + vec0_handle.commit() + vec1_handle = s1_corr_producer.acquire_and_advance() + vec1_handle.commit() + # End of if not continue_cond + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue (TMA store path only) + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.epilogue_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + if cutlass.const_expr(self.use_tma_store): + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q = mQ_qdl.shape[0] + + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = ( + not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + ) + if not continue_cond: + mO_qdl_ = mO_qdl + if cutlass.const_expr(cum_seqlen_q is not None): + mO_qdl_ = cute.domain_offset( + (cum_seqlen_q[batch_coord], 0, ((0, 0), 0)), mO_qdl + ) + + o0_coord = 2 * curr_block_coord[0] + o1_coord = o0_coord + 1 + gO_qdl = cute.flat_divide( + mO_qdl_, cute.select(self.pv_mma_tiler, mode=[0, 1]) + ) + gO = gO_qdl[None, None, None, 0, curr_block_coord[2]] + tOsO, tOgO = cute.nvgpu.cpasync.tma_partition( + tma_atom_o, + 0, + cute.make_layout(1), + cute.group_modes(sO, 0, 2), + cute.group_modes(gO, 0, 2), + ) + + # O0 O1 using the same pipeline + # wait from corr, issue tma store on smem + # O0 + # 1. Wait for O0 final + o0_handle = corr_epi_consumer.wait_and_advance() + # 2. Copy O0 to gmem + cute.copy(tma_atom_o, tOsO[None, 0], tOgO[None, o0_coord]) + cute.arch.cp_async_bulk_commit_group() + # O1 + # 1. Wait for O1 final + o1_handle = corr_epi_consumer.wait_and_advance() + # 2. Copy O1 to gmem + cute.copy(tma_atom_o, tOsO[None, 1], tOgO[None, o1_coord]) + cute.arch.cp_async_bulk_commit_group() + + # Ensure O0 buffer is ready to be released + cute.arch.cp_async_bulk_wait_group(1, read=True) + o0_handle.release() + # Ensure O1 buffer is ready to be released + cute.arch.cp_async_bulk_wait_group(0, read=True) + o1_handle.release() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cute.arch.griddepcontrol_launch_dependents() + # End of persistent scheduler loop + # /////////////////////////////////////////////////////////////////////////////// + # Softmax0 + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.softmax1_warp_ids[0]: + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + tStS, tStS0, tStS1, tOtO0, tOtO1, tOrP0, tOrP1 = make_tmem_tensors( + self, qk_thr_mma, pv_thr_mma, p_tmem_layout_staged, tmem_ptr + ) + softmax_fn( + stage=0, + tensor_args=( + tStS, + tStS0, + cum_seqlen_k, + cum_seqlen_q, + s0_warp_wants_skip_softmax_exchange, + skip_softmax_count, + total_softmax_count, + ), + pipeline_args=(mma_s0_consumer, s0_corr_producer, p0_mma_producer), + inplace_args=(s0_p1_inplace_producer, s1_p0_inplace_consumer), + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax1 + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.correction_warp_ids[0] and warp_idx >= self.softmax1_warp_ids[0]: + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + tStS, tStS0, tStS1, tOtO0, tOtO1, tOrP0, tOrP1 = make_tmem_tensors( + self, qk_thr_mma, pv_thr_mma, p_tmem_layout_staged, tmem_ptr + ) + softmax_fn( + stage=1, + tensor_args=( + tStS, + tStS1, + cum_seqlen_k, + cum_seqlen_q, + s1_warp_wants_skip_softmax_exchange, + skip_softmax_count, + total_softmax_count, + ), + pipeline_args=(mma_s1_consumer, s1_corr_producer, p1_mma_producer), + inplace_args=(s1_p0_inplace_producer, s0_p1_inplace_consumer), + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.correction_warp_ids[0] and warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_correction) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + tStS, tStS0, tStS1, tOtO0, tOtO1, tOrP0, tOrP1 = make_tmem_tensors( + self, qk_thr_mma, pv_thr_mma, p_tmem_layout_staged, tmem_ptr + ) + cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tScS = qk_thr_mma.partition_C(cS) + + tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) + + tStS_vec0 = cute.make_tensor(tStS.iterator + self.tmem_vec0_offset, tStS_vec_layout) + tStS_vec1 = cute.make_tensor(tStS.iterator + self.tmem_vec1_offset, tStS_vec_layout) + + tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) + tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) + tmem_load_v_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_v_atom, tStS_vec0) + thread_idx = tidx % (self.threads_per_warp * len(self.correction_warp_ids)) + thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(thread_idx) + tTMEM_LOAD_VECtS0 = thr_tmem_load_vec.partition_S(tStS_vec0) + tTMEM_LOAD_VECtS1 = thr_tmem_load_vec.partition_S(tStS_vec1) + tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) + old_o0_tma_phase = Int32(0) + old_o1_tma_phase = Int32(0) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + seqlen_k = mK_kdl.shape[0] + row_idx = Int32(0) + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q = mQ_qdl.shape[0] + + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = ( + not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + ) + if not continue_cond: + row_idx = curr_block_coord[0] * self.cta_tiler[0] + tTMEM_LOAD_VECcS[0][0] + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord].to(Int32) + seqlen_k = cum_seqlen_k[batch_coord + 1].to(Int32) - cuseqlen_k + continue_cond = seqlen_k <= 0 + if not continue_cond: + # Ignore first signal from softmax as no correction is required + vec0_handle = s0_corr_consumer.wait_and_advance() + vec0_handle.release() + vec1_handle = s1_corr_consumer.wait_and_advance() + vec1_handle.release() + # O0/O1 share the same mma_corr consumer state, so the Oi + # peek token rolls from O0 -> O1 -> next O0. Seed with a + # blocking token; the rescale helper refreshes it near the + # end of each iteration. + oi_peek_status = cutlass.Boolean(False) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + for i in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + # Rescale O0 + ( + (s0_corr_consumer, mma_corr_consumer), + oi_peek_status, + ) = self.correction_rescale( + pv_thr_mma, + tiled_tmem_load_vec, + scale_softmax_log2, + (tOtO0, tTMEM_LOAD_VECtS0, tTMEM_LOAD_VECcS), + (s0_corr_consumer, mma_corr_consumer), + oi_peek_status, + ) + # Rescale O1 + ( + (s1_corr_consumer, mma_corr_consumer), + oi_peek_status, + ) = self.correction_rescale( + pv_thr_mma, + tiled_tmem_load_vec, + scale_softmax_log2, + (tOtO1, tTMEM_LOAD_VECtS1, tTMEM_LOAD_VECcS), + (s1_corr_consumer, mma_corr_consumer), + oi_peek_status, + ) + # End of seqlen_corr_loop_steps + value_args = ( + cuseqlen_q, + seqlen_q, + curr_block_coord, + scale_softmax, + scale_output, + ) + if cutlass.const_expr(self.use_tma_store): + # TMA store path: write to sO, signal epilogue warp + # Normalize O0 + s0_corr_consumer, mma_corr_consumer, corr_epi_producer = ( + self.correction_epilog( + pv_thr_mma, + tiled_tmem_load_vec, + ( + tOtO0, + tTMEM_LOAD_VECtS0, + tTMEM_LOAD_VECcS, + sO[None, None, 0], + None, + mLSE, + mSink, + mMergeStats, + mMergeOp, + None, + ), + ( + s0_corr_consumer, + mma_corr_consumer, + corr_epi_producer, + ), + (row_idx, *value_args, Int32(0)), + ) + ) + row_idx += self.qk_mma_tiler[0] + # Normalize O1 + s1_corr_consumer, mma_corr_consumer, corr_epi_producer = ( + self.correction_epilog( + pv_thr_mma, + tiled_tmem_load_vec, + ( + tOtO1, + tTMEM_LOAD_VECtS1, + tTMEM_LOAD_VECcS, + sO[None, None, 1], + None, + mLSE, + mSink, + mMergeStats, + mMergeOp, + None, + ), + ( + s1_corr_consumer, + mma_corr_consumer, + corr_epi_producer, + ), + (row_idx, *value_args, Int32(0)), + ) + ) + else: + # st.global path: store directly to global memory + block_offset_o = Int32(0) + if cutlass.const_expr(cum_seqlen_q is not None): + block_offset_o = cum_seqlen_q[batch_coord] + mO_ = cute.make_tensor( + mO.iterator + block_offset_o * mO.stride[0], + cute.make_layout( + (seqlen_q, mO.shape[1], mO.shape[2]), + stride=mO.stride, + ), + ) + o0_coord = 2 * curr_block_coord[0] + o1_coord = o0_coord + 1 + gO_stg = cute.local_tile( + mO_, + (self.pv_mma_tiler[0], self.pv_mma_tiler[1]), + (None, None, None), + ) + gO0 = gO_stg[None, None, o0_coord, 0, curr_block_coord[2]] + gO1 = gO_stg[None, None, o1_coord, 0, curr_block_coord[2]] + tOldOsO = None + tOldOgO = None + if cutlass.const_expr(tma_atom_old_o is not None): + mOldO_qdl_ = mOldO_qdl + if cutlass.const_expr(cum_seqlen_q is not None): + mOldO_qdl_ = cute.domain_offset( + (cuseqlen_q, 0, ((0, 0), 0)), mOldO_qdl + ) + gOldO_qdl = cute.flat_divide( + mOldO_qdl_, + cute.select(self.pv_mma_tiler, mode=[0, 1]), + ) + gOldO = gOldO_qdl[None, None, None, 0, curr_block_coord[2]] + tOldOsO, tOldOgO = cute.nvgpu.cpasync.tma_partition( + tma_atom_old_o, + 0, + cute.make_layout(1), + cute.group_modes(sO, 0, 2), + cute.group_modes(gOldO, 0, 2), + ) + # Normalize O0 and store to global memory + old_o_tma_args = None + if cutlass.const_expr(tma_atom_old_o is not None): + old_o_tma_args = ( + tma_atom_old_o, + tOldOgO[None, o0_coord], + tOldOsO[None, 0], + old_o_tma_mbar_ptr, + ) + ( + s0_corr_consumer, + mma_corr_consumer, + old_o0_tma_phase, + ) = self.correction_epilog( + pv_thr_mma, + tiled_tmem_load_vec, + ( + tOtO0, + tTMEM_LOAD_VECtS0, + tTMEM_LOAD_VECcS, + gO0, + sO[None, None, 0], + mLSE, + mSink, + mMergeStats, + mMergeOp, + old_o_tma_args, + ), + (s0_corr_consumer, mma_corr_consumer), + (row_idx, *value_args, old_o0_tma_phase), + ) + row_idx += self.qk_mma_tiler[0] + # Normalize O1 and st.global to global memory + old_o_tma_args = None + if cutlass.const_expr(tma_atom_old_o is not None): + old_o_tma_args = ( + tma_atom_old_o, + tOldOgO[None, o1_coord], + tOldOsO[None, 1], + old_o_tma_mbar_ptr + 1, + ) + ( + s1_corr_consumer, + mma_corr_consumer, + old_o1_tma_phase, + ) = self.correction_epilog( + pv_thr_mma, + tiled_tmem_load_vec, + ( + tOtO1, + tTMEM_LOAD_VECtS1, + tTMEM_LOAD_VECcS, + gO1, + sO[None, None, 1], + mLSE, + mSink, + mMergeStats, + mMergeOp, + old_o_tma_args, + ), + (s1_corr_consumer, mma_corr_consumer), + (row_idx, *value_args, old_o1_tma_phase), + ) + # End of if not continue_cond + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + if cutlass.const_expr(not self.use_tma_store): + cute.arch.griddepcontrol_launch_dependents() + # End of persistent scheduler loop + tmem.relinquish_alloc_permit() + # Synchronize before TMEM dealloc (done by the caller) + self.tmem_dealloc_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + return + + @cute.jit + def get_skip_softmax_flag(self, tiled_tmem_load_v, tTMEM_LOADtS_v, tTMEM_LOADrS_v): + cute.copy(tiled_tmem_load_v, tTMEM_LOADtS_v, tTMEM_LOADrS_v) + tTMEM_LOADrS_v_i32 = cute.recast_tensor(tTMEM_LOADrS_v, dtype=cutlass.Int32) + skip_softmax_flag = cute.arch.make_warp_uniform(tTMEM_LOADrS_v_i32[0]) + return skip_softmax_flag + + @cute.jit + def mma_qk( + self, + tiled_mma: cute.TiledMma, + tensor_args: Tuple, + pipeline_args: Tuple, + pipeline_tokens: Tuple = (None, None), + ) -> Tuple[pipeline.PipelineProducer, pipeline.PipelineProducer]: + """Perform a single step of the QK GEMM computation on a block of attention scores. + + :param tiled_mma: Tiled MMA for QK GEMM + :type tiled_mma: cute.TiledMma + :param tensor_args: Tuple containing Qi, K, and Si + :type tensor_args: Tuple + :param pipeline_args: Tuple containing mma_si_producer and si_corr_producer + :type pipeline_args: Tuple + :param pipeline_tokens: Optional non-blocking peek tokens for the Si and + vec_i producers, in the form ``(si_peek_status, veci_peek_status)``. + ``None`` for either token falls back to a blocking acquire. + :type pipeline_tokens: Tuple + :return: Tuple containing mma_si_producer and si_corr_producer + :rtype: Tuple[pipeline.PipelineProducer, pipeline.PipelineProducer] + """ + tSrQi, tSrK, tStSi = tensor_args + mma_si_producer, si_corr_producer = pipeline_args + si_peek_status, veci_peek_status = pipeline_tokens + # 0. Make sure Qi & K are ready when calling mma_qk + # 1. acquire S0 + si_handle = mma_si_producer.acquire_and_advance(si_peek_status) + # 2. make sure vec is already released in corr + veci_handle = si_corr_producer.acquire_and_advance(veci_peek_status) + veci_handle.commit() + # 3. gemm + num_kphases = cute.size(tSrQi, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0) + cute.gemm( + tiled_mma, + tStSi, + tSrQi[kphase_coord], + tSrK[kphase_coord], + tStSi, + ) + # {$nv-internal-release begin} + # This is a trick to avoid grouped ALUs which might block the issue + # of MUFU sequences. + # To apply this trick, this example hasn't changed to blocked gemm api. + fmha_utils.cfence() # {$nv-internal-release} + # {$nv-internal-release end} + # 4. release S0 + si_handle.commit() + return mma_si_producer, si_corr_producer + + @cute.jit + def mma_pv( + self, + tiled_mma: cute.TiledMma, + whether_acc: bool, + tensor_args: Tuple, + pipeline_args: Tuple, + skip_pv_args: Tuple, + ) -> Tuple[ + pipeline.PipelineProducer, + pipeline.PipelineConsumer, + ]: + """Perform a single step of the PV GEMM computation on accumulating O. + + :param tiled_mma: Tiled MMA for PV GEMM + :type tiled_mma: cute.TiledMma + :param whether_acc: Whether to accumulate O + :type whether_acc: bool + :param tensor_args: Tuple containing Pi, Vi, and Oi + :type tensor_args: Tuple + :param pipeline_args: Tuple containing mma_corr_producer and pi_mma_consumer + :type pipeline_args: Tuple + :param skip_pv_args: Tuple containing enable_skip_softmax, tiled_tmem_load_v, tTMEM_LOADtS_v, tTMEM_LOADrS_v + :type skip_pv_args: Tuple + :return: Tuple containing mma_corr_producer and pi_mma_consumer + :rtype: Tuple[pipeline.PipelineProducer, pipeline.PipelineConsumer] + """ + tOrPi, tOrVi, tOtOi = tensor_args + mma_corr_producer, pi_mma_consumer = pipeline_args + enable_skip_softmax, tiled_tmem_load_v, tTMEM_LOADtS_v, tTMEM_LOADrS_v = skip_pv_args + # 0. Make sure Vi is ready when calling mma_pv + # 1. acquire Oi + oi_handle = mma_corr_producer.acquire_and_advance() + # 2. wait for Pi + pi_handle = pi_mma_consumer.wait_and_advance() + # 3. gemm + num_kphases = cute.size(tOrPi, mode=[2]) + if cutlass.const_expr(enable_skip_softmax): + skip_pv = self.get_skip_softmax_flag(tiled_tmem_load_v, tTMEM_LOADtS_v, tTMEM_LOADrS_v) + if not skip_pv: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + tiled_mma.set(tcgen05.Field.ACCUMULATE, whether_acc or kphase_idx != 0) + cute.gemm( + tiled_mma, + tOtOi, + tOrPi[kphase_coord], + tOrVi[kphase_coord], + tOtOi, + ) + # {$nv-internal-release begin} + # This is a trick to avoid grouped ALUs which might block the issue + # of MUFU sequences. + # To apply this trick, this example hasn't changed to blocked gemm api. + fmha_utils.cfence() # {$nv-internal-release} + # {$nv-internal-release end} + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + tiled_mma.set(tcgen05.Field.ACCUMULATE, whether_acc or kphase_idx != 0) + cute.gemm( + tiled_mma, + tOtOi, + tOrPi[kphase_coord], + tOrVi[kphase_coord], + tOtOi, + ) + # {$nv-internal-release begin} + # This is a trick to avoid grouped ALUs which might block the issue + # of MUFU sequences. + # To apply this trick, this example hasn't changed to blocked gemm api. + fmha_utils.cfence() # {$nv-internal-release} + # {$nv-internal-release end} + # 4. commit Pi + pi_handle.release() + # 5. commit Oi + oi_handle.commit() + return mma_corr_producer, pi_mma_consumer + + @cute.jit + def calculate_skip_softmax_flag( + self, + row_max, + tile_row_max, + scale_softmax_log2, + skip_softmax_threshold_log2, + seqlen_q, + thread_idx, + logical_offset, + warp_wants_skip_softmax_exchange, + stage, + skip_softmax_count, + total_softmax_count, + ) -> Tuple[bool, float]: + """Calculate the skip softmax flag and the row maximum. + + :param row_max: The row maximum. + :type row_max: float + :param tile_row_max: The tile row maximum. + :type tile_row_max: float + :param scale_softmax_log2: The scale softmax log2. + :type scale_softmax_log2: float + :param skip_softmax_threshold_log2: The skip softmax threshold log2. + :type skip_softmax_threshold_log2: float + :param seqlen_q: The sequence length q. + :type seqlen_q: int + :param thread_idx: The thread index. + :type thread_idx: int + :param logical_offset: The logical offset. + :type logical_offset: Tuple[int, int] + :param warp_wants_skip_softmax_exchange: The warp wants skip softmax exchange. + :type warp_wants_skip_softmax_exchange: cute.Tensor + :param stage: The stage. + :type stage: int + :param skip_softmax_count: The skip softmax count. + :type skip_softmax_count: cute.Tensor + :param total_softmax_count: The total softmax count. + :type total_softmax_count: cute.Tensor + :return: Tuple containing the skip softmax flag and the row maximum. + :rtype: Tuple[bool, float] + """ + thread_wants_skip = ( + tile_row_max * scale_softmax_log2 - row_max * scale_softmax_log2 + ) < skip_softmax_threshold_log2 + thread_wants_skip = thread_wants_skip or ((logical_offset[0] + thread_idx) >= seqlen_q) + warp_wants_skip = cute.arch.vote_all_sync(thread_wants_skip) + + with cute.arch.elect_one(): + warp_wants_skip_softmax_exchange[cute.arch.warp_idx() % 4] = warp_wants_skip + softmax_barrier = self.s0_warpgroup_barrier if stage == 0 else self.s1_warpgroup_barrier + softmax_barrier.arrive_and_wait() + warp_wants_skip_softmax_exchange_i32 = cute.make_tensor( + cute.recast_ptr(warp_wants_skip_softmax_exchange.iterator, dtype=cutlass.Int32), + cute.make_layout((1,)), + ) + skip_softmax = cute.arch.popc(warp_wants_skip_softmax_exchange_i32[0]) == 4 + + if not skip_softmax: + row_max = max(row_max, tile_row_max) + + if cutlass.const_expr(skip_softmax_count is not None): + if thread_idx == 0: + if skip_softmax: + cute.arch.atomic_add(skip_softmax_count.iterator.llvm_ptr, Int32(1)) + cute.arch.atomic_add(total_softmax_count.iterator.llvm_ptr, Int32(1)) + return skip_softmax, row_max + + @cute.jit + def apply_exp_and_cvt( + self, + tTMEM_LOADrS, + tTMEM_LOADrS_cvt, + tTMEM_STORErS_x4_e_cvt, + stage, + scale, + minus_row_max_scale, + local_row_sum, + inplace_consumer, + EXP2_EMULATION_OFFSET, + EXP2_EMULATION_COUNT, + CVT_COUNT, + CVT_PER_STEP, + FMA_COUNT, + ARV_COUNT, + ): + """Apply the exp and conversion to the P data type on fragment. + + :param tTMEM_LOADrS: The tTMEM_LOADrS tensor. + :type tTMEM_LOADrS: cute.Tensor + :param tTMEM_LOADrS_cvt: The tTMEM_LOADrS_cvt tensor. + :type tTMEM_LOADrS_cvt: cute.Tensor + :param tTMEM_STORErS_x4_e_cvt: The tTMEM_STORErS_x4_e_cvt tensor. + :type tTMEM_STORErS_x4_e_cvt: cute.Tensor + :param stage: The stage. + :type stage: int + :param scale: The scale. + :type scale: float + :param minus_row_max_scale: The minus row maximum scale. + :type minus_row_max_scale: float + :param local_row_sum: The local row sum. + :type local_row_sum: float + :param inplace_consumer: The inplace consumer. + :type inplace_consumer: cute.Tensor + :param EXP2_EMULATION_OFFSET: The exp2 emulation offset. + :type EXP2_EMULATION_OFFSET: int + :param EXP2_EMULATION_COUNT: The exp2 emulation count. + :type EXP2_EMULATION_COUNT: int + :param CVT_COUNT: The cvt count. + :type CVT_COUNT: int + :param CVT_PER_STEP: The cvt per step. + :type CVT_PER_STEP: int + :param FMA_COUNT: The fma count. + :type FMA_COUNT: int + :param ARV_COUNT: The arv count. + :type ARV_COUNT: int + :return: The local row sum and the inplace consumer. + :rtype: Tuple[float, cute.Tensor] + """ + for i in cutlass.range_constexpr(0, EXP2_EMULATION_OFFSET, 2): + if cutlass.const_expr(i >= CVT_COUNT): + if cutlass.const_expr(i % CVT_PER_STEP == 0): + if cutlass.const_expr(self.v_dtype.width == 8): + fmha_utils.cvt_f32x4_to_f8x4( + tTMEM_LOADrS_cvt[None, (i - CVT_COUNT) // CVT_PER_STEP], + tTMEM_STORErS_x4_e_cvt[None, (i - CVT_COUNT) // CVT_PER_STEP], + ) + else: + s_vec = tTMEM_LOADrS_cvt[None, (i - CVT_COUNT) // CVT_PER_STEP].load() + tTMEM_STORErS_x4_e_cvt[None, (i - CVT_COUNT) // CVT_PER_STEP].store( + s_vec.to(self.v_dtype) + ) + local_row_sum = cute.arch.add_packed_f32x2( + local_row_sum, + ( + tTMEM_LOADrS[i - CVT_COUNT], + tTMEM_LOADrS[i - CVT_COUNT + 1], + ), + ) + fmha_utils.cfence() # {$nv-internal-release} + tTMEM_LOADrS[i] = cute.math.exp2(tTMEM_LOADrS[i], fastmath=True) + fmha_utils.cfence() # {$nv-internal-release} + if cutlass.const_expr(i + FMA_COUNT < EXP2_EMULATION_OFFSET): + ( + tTMEM_LOADrS[i + FMA_COUNT], + tTMEM_LOADrS[i + FMA_COUNT + 1], + ) = cute.arch.fma_packed_f32x2( + ( + tTMEM_LOADrS[i + FMA_COUNT], + tTMEM_LOADrS[i + FMA_COUNT + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + ) + fmha_utils.cfence() # {$nv-internal-release} + tTMEM_LOADrS[i + 1] = cute.math.exp2(tTMEM_LOADrS[i + 1], fastmath=True) + fmha_utils.cfence() # {$nv-internal-release} + if cutlass.const_expr(i == EXP2_EMULATION_OFFSET - ARV_COUNT): + if cutlass.const_expr(self.enable_sequence_barrier): + if cutlass.const_expr(stage == 0): + self.sequence_s1_s0_barrier.arrive() + else: + self.sequence_s0_s1_barrier.arrive() + fmha_utils.cfence() # {$nv-internal-release} + fmha_utils.reset_sched_res_busy_xu64() # {$nv-internal-release} + + # The remaining conversion steps + for i in cutlass.range_constexpr( + EXP2_EMULATION_OFFSET - CVT_COUNT, + EXP2_EMULATION_OFFSET, + 2, + ): + if cutlass.const_expr(i % CVT_PER_STEP == 0): + if cutlass.const_expr(self.v_dtype.width == 8): + fmha_utils.cvt_f32x4_to_f8x4( + tTMEM_LOADrS_cvt[None, i // CVT_PER_STEP], + tTMEM_STORErS_x4_e_cvt[None, i // CVT_PER_STEP], + ) + else: + s_vec = tTMEM_LOADrS_cvt[None, i // CVT_PER_STEP].load() + tTMEM_STORErS_x4_e_cvt[None, i // CVT_PER_STEP].store(s_vec.to(self.v_dtype)) + fmha_utils.warp_switch() # {$nv-internal-release} + local_row_sum = cute.arch.add_packed_f32x2( + local_row_sum, (tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1]) + ) + fmha_utils.warp_switch() # {$nv-internal-release} + for i in cutlass.range_constexpr( + EXP2_EMULATION_OFFSET, EXP2_EMULATION_OFFSET + EXP2_EMULATION_COUNT // 2, 2 + ): + tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1] = cute.arch.fma_packed_f32x2( + (tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1]), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + ) + fmha_utils.warp_switch() # {$nv-internal-release} + tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1] = fmha_utils.ex2_emulation_packed_f32x2( + tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1] + ) + if cutlass.const_expr((i + 2) % CVT_PER_STEP == 0): + if cutlass.const_expr(self.v_dtype.width == 8): + fmha_utils.cvt_f32x4_to_f8x4( + tTMEM_LOADrS_cvt[None, i // CVT_PER_STEP], + tTMEM_STORErS_x4_e_cvt[None, i // CVT_PER_STEP], + ) + else: + s_vec = tTMEM_LOADrS_cvt[None, i // CVT_PER_STEP].load() + tTMEM_STORErS_x4_e_cvt[None, i // CVT_PER_STEP].store(s_vec.to(self.v_dtype)) + + inplace_peek_status = inplace_consumer.try_wait() + for i in cutlass.range_constexpr( + EXP2_EMULATION_OFFSET + EXP2_EMULATION_COUNT // 2, + EXP2_EMULATION_OFFSET + EXP2_EMULATION_COUNT, + 2, + ): + tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1] = cute.arch.fma_packed_f32x2( + (tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1]), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + ) + tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1] = fmha_utils.ex2_emulation_packed_f32x2( + tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1] + ) + if cutlass.const_expr((i + 2) % CVT_PER_STEP == 0): + if cutlass.const_expr(self.v_dtype.width == 8): + fmha_utils.cvt_f32x4_to_f8x4( + tTMEM_LOADrS_cvt[None, i // CVT_PER_STEP], + tTMEM_STORErS_x4_e_cvt[None, i // CVT_PER_STEP], + ) + else: + s_vec = tTMEM_LOADrS_cvt[None, i // CVT_PER_STEP].load() + tTMEM_STORErS_x4_e_cvt[None, i // CVT_PER_STEP].store(s_vec.to(self.v_dtype)) + inplace_consumer.wait_and_advance(inplace_peek_status) + return local_row_sum, inplace_consumer + + @cute.jit + def softmax_step( + self, + stage: int, + whether_apply_mask: bool, + iter_args: Tuple, + stats_args: Tuple, + pipeline_args: Tuple, + value_args: Tuple, + atom_args: Tuple, + tensor_args: Tuple, + ) -> Tuple[Tuple, Tuple]: + """Perform a single step of the softmax computation on a block of attention scores. + + This method processes one block of the attention matrix, computing numerically stable + softmax by first finding the row maximum, subtracting it from all elements, applying + exponential function, and then normalizing by the sum of exponentials. It also handles + optional masking of attention scores. + + The method involves several key operations: + 1. Loading attention scores from tensor memory + 2. Applying optional masking based on position + 3. Computing row-wise maximum values for numerical stability + 4. Transforming scores using exp2(x*scale - max*scale) + 5. Computing row sums for normalization + 6. Coordinating pipeline synchronization between different processing stages + + :param stage: Processing stage (0 for first half, 1 for second half) + :type stage: int + :param whether_apply_mask: Whether to apply attention masking + :type whether_apply_mask: bool + :param iter_args: Tuple containing the counting tensor, row_max, row_sum, + and vector-buffer handle for the current iteration + :type iter_args: Tuple + :param stats_args: Tuple containing row_sum and row_max + :type stats_args: Tuple + :param pipeline_args: Tuple containing pipeline arguments for MMA, + correction, and sequence synchronization + :type pipeline_args: Tuple + :param value_args: Tuple containing seqlen_k, seqlen_q, and scale_softmax_log2 + :type value_args: Tuple + :param atom_args: Tuple containing mma & copy atoms + :type atom_args: Tuple + :param tensor_args: Tuple containing softmax related tensors + :type tensor_args: Tuple + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask + :return: Updated stats_args and pipeline_args + :rtype: Tuple[Tuple, Tuple] + """ + row_sum, row_max = stats_args + cS, is_last_iter = iter_args + ( + seqlen_k, + seqlen_q, + scale_softmax_log2, + window_size_left, + window_size_right, + skip_softmax_threshold_log2, + thread_idx, + logical_offset, + ) = value_args + ( + si_peek_status, + mma_si_consumer, + si_corr_producer, + pi_mma_producer, + inplace_producer, + inplace_consumer, + ) = pipeline_args + ( + qk_thr_mma, + tiled_tmem_load, + tiled_tmem_store, + tiled_tmem_store_vec, + thr_tmem_load, + thr_tmem_store, + thr_tmem_store_vec, + ) = atom_args + ( + tTMEM_LOADtS, + tTMEM_STORE_VECtS, + tTMEM_STORE_SKIP_SOFTMAX, + tTMEM_STOREtS_x4, + warp_wants_skip_softmax_exchange, + skip_softmax_count, + total_softmax_count, + ) = tensor_args + tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width + tScS = qk_thr_mma.partition_C(cS) + enable_skip_softmax = skip_softmax_threshold_log2 is not None + tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) + tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) + tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) + tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) + tTMEM_LOADcS = thr_tmem_load.partition_D(tScS) + tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) + tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P) + # Wait for Si + si_handle = mma_si_consumer.wait_and_advance(si_peek_status) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + old_row_max = row_max + skip_softmax = cutlass.Boolean(False) + if whether_apply_mask: + if cutlass.const_expr(self.arch >= Arch.sm_100 and self.arch <= Arch.sm_100f): + cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) + else: + tTMEM_LOADrMax = cute.make_rmem_tensor( + cute.make_layout((1, cute.size(tTMEM_LOADrS, mode=[1]))), + self.qk_acc_dtype, + ) + for i in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS, mode=[1])): + cute.copy_atom_call( + tiled_tmem_load, + tTMEM_LOADtS[None, i, 0, 0], + (tTMEM_LOADrS[None, i, 0, 0], tTMEM_LOADrMax[None, i]), + ) + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + tile_row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, -cutlass.Float32.inf, 0) + if cutlass.const_expr(not enable_skip_softmax): + row_max = cute.arch.fmax(row_max, tile_row_max) + else: + skip_softmax, row_max = self.calculate_skip_softmax_flag( + row_max, + tile_row_max, + scale_softmax_log2, + skip_softmax_threshold_log2, + seqlen_q, + thread_idx, + logical_offset, + warp_wants_skip_softmax_exchange, + stage, + skip_softmax_count, + total_softmax_count, + ) + si_handle.release() + # S0 -> P1 / S1 -> P0 + inplace_producer.commit() + inplace_producer.advance() + else: + if cutlass.const_expr(self.arch >= Arch.sm_100 and self.arch <= Arch.sm_100f): + cute.copy( + tiled_tmem_load, + tTMEM_LOADtS[None, 0, None, None], + tTMEM_LOADrS[None, 0, None, None], + ) + cute.copy( + tiled_tmem_load, + tTMEM_LOADtS[None, 1, None, None], + tTMEM_LOADrS[None, 1, None, None], + ) + fmha_utils.cfence() # {$nv-internal-release} + tile_row_max = -cutlass.Float32.inf + tile_row_max_ = tile_row_max + for i in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS, mode=[0]), 4): + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i, 0, 0, 0]) + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i + 1, 0, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 2, 0, 0, 0]) + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 3, 0, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + fmha_utils.cfence() # {$nv-internal-release} + cute.copy( + tiled_tmem_load, + tTMEM_LOADtS[None, 2, None, None], + tTMEM_LOADrS[None, 2, None, None], + ) + fmha_utils.cfence() # {$nv-internal-release} + for i in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS, mode=[0]), 4): + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i, 1, 0, 0]) + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i + 1, 1, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 2, 1, 0, 0]) + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 3, 1, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + fmha_utils.cfence() # {$nv-internal-release} + cute.copy( + tiled_tmem_load, + tTMEM_LOADtS[None, 3, None, None], + tTMEM_LOADrS[None, 3, None, None], + ) + for i in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS, mode=[0]), 4): + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i, 2, 0, 0]) + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i + 1, 2, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 2, 2, 0, 0]) + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 3, 2, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + fmha_utils.cfence() # {$nv-internal-release} + cute.arch.fence_view_async_tmem_store() + si_handle.release() + fmha_utils.cfence() # {$nv-internal-release} + # S0 -> P1 / S1 -> P0 + inplace_producer.commit() + inplace_producer.advance() + fmha_utils.cfence() # {$nv-internal-release} + for i in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS, mode=[0]), 4): + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i, 3, 0, 0]) + tile_row_max = cute.arch.fmax(tile_row_max, tTMEM_LOADrS[i + 1, 3, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 2, 3, 0, 0]) + tile_row_max_ = cute.arch.fmax(tile_row_max_, tTMEM_LOADrS[i + 3, 3, 0, 0]) + fmha_utils.warp_switch() # {$nv-internal-release} + fmha_utils.cfence() # {$nv-internal-release} + tile_row_max = cute.arch.fmax(tile_row_max, tile_row_max_) + if cutlass.const_expr(not enable_skip_softmax): + row_max = cute.arch.fmax(tile_row_max, row_max) + else: + skip_softmax, row_max = self.calculate_skip_softmax_flag( + row_max, + tile_row_max, + scale_softmax_log2, + skip_softmax_threshold_log2, + seqlen_q, + thread_idx, + logical_offset, + warp_wants_skip_softmax_exchange, + stage, + skip_softmax_count, + total_softmax_count, + ) + else: + tTMEM_LOADrMax = cute.make_rmem_tensor( + cute.make_layout((1, cute.size(tTMEM_LOADrS, mode=[1]))), + self.qk_acc_dtype, + ) + for i in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS, mode=[1])): + cute.copy_atom_call( + tiled_tmem_load, + tTMEM_LOADtS[None, i, 0, 0], + (tTMEM_LOADrS[None, i, 0, 0], tTMEM_LOADrMax[None, i]), + ) + cute.arch.fence_view_async_tmem_store() + tile_row_max = tTMEM_LOADrMax.load().reduce( + cute.ReductionOp.MAX, -cutlass.Float32.inf, 0 + ) + if cutlass.const_expr(not enable_skip_softmax): + row_max = cute.arch.fmax(tile_row_max, row_max) + else: + skip_softmax, row_max = self.calculate_skip_softmax_flag( + row_max, + tile_row_max, + scale_softmax_log2, + skip_softmax_threshold_log2, + seqlen_q, + thread_idx, + logical_offset, + warp_wants_skip_softmax_exchange, + stage, + skip_softmax_count, + total_softmax_count, + ) + si_handle.release() + # S0 -> P1 / S1 -> P0 + inplace_producer.commit() + inplace_producer.advance() + + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + if cutlass.const_expr(self.rescale_threshold > 0.0): + if (row_max_safe - old_row_max) * scale_softmax_log2 <= self.rescale_threshold: + row_max_safe = old_row_max + tTMEM_STORE_VECrS = cute.make_rmem_tensor(tTMEM_STORE_VECcS.shape, self.qk_acc_dtype) + tTMEM_STORE_VECrS[0] = old_row_max + tTMEM_STORE_VECrS[1] = row_max_safe + vec_i_peek_status = si_corr_producer.try_acquire() + tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) + tTMEM_STORErS_x4_e = cute.make_tensor( + cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.v_dtype), + tTMEM_LOADrS.layout, + ) + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + if cutlass.const_expr(self.v_dtype.width == 8 and self.p_fp8_prescale_log2 > 0): + minus_row_max_scale = minus_row_max_scale + self.p_fp8_prescale_log2 + + ARV_COUNT = 4 + FMA_COUNT = 8 + CVT_COUNT = 8 if self.v_dtype.width == 8 else 4 + CVT_PER_STEP = 4 if self.v_dtype.width == 8 else 2 + assert CVT_COUNT % CVT_PER_STEP == 0, ( + f"CVT_COUNT {CVT_COUNT} must be divisible by CVT_PER_STEP {CVT_PER_STEP}" + ) + tTMEM_LOADrS_cvt = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(CVT_PER_STEP)) + tTMEM_STORErS_x4_e_cvt = cute.logical_divide( + tTMEM_STORErS_x4_e, cute.make_layout(CVT_PER_STEP) + ) + for i in cutlass.range_constexpr(0, FMA_COUNT, 2): + tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1] = cute.arch.fma_packed_f32x2( + (tTMEM_LOADrS[i], tTMEM_LOADrS[i + 1]), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + ) + vec_i_handle = si_corr_producer.acquire_and_advance(vec_i_peek_status) + cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) + cute.arch.fence_view_async_tmem_store() + # Notify correction wg that row_max is ready + vec_i_handle.commit() + + fmha_utils.sched_res_busy_xu64() # {$nv-internal-release} + EXP2_EMULATION_COUNT = 20 if self.enable_ex2_emulation and not whether_apply_mask else 0 + EXP2_EMULATION_OFFSET = cute.size(tTMEM_LOADrS) - EXP2_EMULATION_COUNT + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5 + fmha_utils.cfence() # {$nv-internal-release} + if cutlass.const_expr(self.enable_sequence_barrier): + if cutlass.const_expr(stage == 0): + self.sequence_s0_s1_barrier.arrive_and_wait() + else: + self.sequence_s1_s0_barrier.arrive_and_wait() + fmha_utils.cfence() # {$nv-internal-release} + + if cutlass.const_expr(enable_skip_softmax): + if not skip_softmax: + row_sum *= acc_scale + local_row_sum = (row_sum, row_sum) + local_row_sum, inplace_consumer = self.apply_exp_and_cvt( + tTMEM_LOADrS, + tTMEM_LOADrS_cvt, + tTMEM_STORErS_x4_e_cvt, + stage, + scale, + minus_row_max_scale, + local_row_sum, + inplace_consumer, + EXP2_EMULATION_OFFSET, + EXP2_EMULATION_COUNT, + CVT_COUNT, + CVT_PER_STEP, + FMA_COUNT, + ARV_COUNT, + ) + tTMEM_STORE_VECrS_i32 = cute.recast_tensor(tTMEM_STORE_VECrS, dtype=cutlass.Int32) + tTMEM_STORE_VECrS_i32[0] = 0 + pi_handle = pi_mma_producer.acquire_and_advance() + # store skip softmax flag + cute.copy( + tiled_tmem_store_vec, + tTMEM_STORE_VECrS_i32, + tTMEM_STORE_SKIP_SOFTMAX, + ) + # store P + cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) + cute.arch.fence_view_async_tmem_store() + pi_handle.commit() + fmha_utils.cfence() # {$nv-internal-release} + for j in cutlass.range_constexpr( + EXP2_EMULATION_OFFSET, + EXP2_EMULATION_OFFSET + EXP2_EMULATION_COUNT, + 2, + ): + local_row_sum = cute.arch.add_packed_f32x2( + (tTMEM_LOADrS[j], tTMEM_LOADrS[j + 1]), + local_row_sum, + ) + row_sum = local_row_sum[0] + local_row_sum[1] + fmha_utils.cfence() # {$nv-internal-release} + cute.arch.fence_view_async_tmem_store() + else: + if cutlass.const_expr(self.enable_sequence_barrier): + if cutlass.const_expr(stage == 0): + self.sequence_s1_s0_barrier.arrive() + else: + self.sequence_s0_s1_barrier.arrive() + inplace_peek_status = inplace_consumer.try_wait() + inplace_consumer.wait_and_advance(inplace_peek_status) + tTMEM_STORE_VECrS_i32 = cute.recast_tensor(tTMEM_STORE_VECrS, dtype=cutlass.Int32) + tTMEM_STORE_VECrS_i32[0] = 1 + pi_handle = pi_mma_producer.acquire_and_advance() + # store skip softmax flag + cute.copy( + tiled_tmem_store_vec, + tTMEM_STORE_VECrS_i32, + tTMEM_STORE_SKIP_SOFTMAX, + ) + cute.arch.fence_view_async_tmem_store() + pi_handle.commit() + fmha_utils.cfence() # {$nv-internal-release} + else: + row_sum *= acc_scale + local_row_sum = (row_sum, row_sum) + local_row_sum, inplace_consumer = self.apply_exp_and_cvt( + tTMEM_LOADrS, + tTMEM_LOADrS_cvt, + tTMEM_STORErS_x4_e_cvt, + stage, + scale, + minus_row_max_scale, + local_row_sum, + inplace_consumer, + EXP2_EMULATION_OFFSET, + EXP2_EMULATION_COUNT, + CVT_COUNT, + CVT_PER_STEP, + FMA_COUNT, + ARV_COUNT, + ) + pi_handle = pi_mma_producer.acquire_and_advance() + # store P + cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) + cute.arch.fence_view_async_tmem_store() + fmha_utils.cfence() # {$nv-internal-release} + for j in cutlass.range_constexpr( + EXP2_EMULATION_OFFSET, + EXP2_EMULATION_OFFSET + EXP2_EMULATION_COUNT, + 2, + ): + local_row_sum = cute.arch.add_packed_f32x2( + (tTMEM_LOADrS[j], tTMEM_LOADrS[j + 1]), + local_row_sum, + ) + row_sum = local_row_sum[0] + local_row_sum[1] + fmha_utils.cfence() # {$nv-internal-release} + cute.arch.fence_view_async_tmem_store() + # Notify tensor core warp that softmax(S->P) is ready + pi_handle.commit() + if not is_last_iter: + si_peek_status = mma_si_consumer.try_wait() + + stats_args = (row_sum, row_max_safe) + pipeline_args = ( + si_peek_status, + mma_si_consumer, + si_corr_producer, + pi_mma_producer, + inplace_producer, + inplace_consumer, + ) + return stats_args, pipeline_args + + # For both softmax0 and softmax1 warp group + @cute.jit + def softmax( + self, + stage: int, + tensor_args: Tuple, + pipeline_args: Tuple, + inplace_args: Tuple, + qk_thr_mma: cute.ThrMma, + value_args: Tuple, + mask_args: Tuple, + sched_args: Tuple, + ): + """Compute softmax on attention scores from QK matrix multiplication. + + This method handles the softmax computation for either the first or second half of the + attention matrix, depending on the 'stage' parameter. It calculates row-wise maximum + and sum values needed for stable softmax computation, applies optional masking, and + transforms raw attention scores into probability distributions. + + The implementation uses specialized memory access patterns and efficient math operations + for computing exp(x) using exp2 functions. It also coordinates pipeline + synchronization between MMA, correction, and sequence processing stages. + + :param stage: Processing stage (0 for first half, 1 for second half of attention matrix) + :type stage: int + :param seqlen_k: Length of the key sequence + :type seqlen_k: Int32 + :param seqlen_q: Length of the query sequence + :type seqlen_q: Int32 + :param cum_seqlen_q: Cumulative sequence lengths for queries + :type cum_seqlen_q: cute.Tensor | None + :param cum_seqlen_k: Cumulative sequence lengths for keys + :type cum_seqlen_k: cute.Tensor | None + :param scale_softmax_log2: Log2 scale factor for softmax operation + :type scale_softmax_log2: Float32 + :param qk_thr_mma: Thread MMA operation for QK matrix multiplication + :type qk_thr_mma: cute.ThrMma + :param tStS: Shared tensor for softmax input/output + :type tStS: cute.Tensor + :param tStSi: Input tensor containing attention scores + :type tStSi: cute.Tensor + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param mma_si_consumer: Pipeline for synchronizing with Si tensors + :type mma_si_consumer: pipeline.PipelineConsumer + :param si_corr_producer: Pipeline for synchronizing with correction operations + :type si_corr_producer: pipeline.PipelineProducer + :param pi_mma_producer: Pipeline for synchronizing with Pi tensors + :type pi_mma_producer: pipeline.PipelineProducer + :param tile_sched_params: Parameters for tile scheduling + :type tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask + """ + ( + tStS, + tStSi, + cum_seqlen_k, + cum_seqlen_q, + warp_wants_skip_softmax_exchange, + skip_softmax_count, + total_softmax_count, + ) = tensor_args + mma_si_consumer, si_corr_producer, pi_mma_producer = pipeline_args + inplace_producer, inplace_consumer = inplace_args + ( + seqlen_k, + seqlen_q, + scale_softmax_log2, + skip_softmax_threshold_log2, + ) = value_args + window_size_left, window_size_right = mask_args + tile_sched, work_tile = sched_args + + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax0_warp_ids)) + + cS_base = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tilePlikeFP32 = self.qk_mma_tiler[1] // 32 * self.o_dtype.width + tScS = qk_thr_mma.partition_C(cS_base) + tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) + tmem_vec_offset = self.tmem_vec0_offset if stage == 0 else self.tmem_vec1_offset + tStS_vec = cute.make_tensor(tStS.iterator + tmem_vec_offset, tStS_vec_layout) + tmem_skip_softmax_offset = ( + self.tmem_skip_softmax0_offset if stage == 0 else self.tmem_skip_softmax1_offset + ) + tStS_skip_softmax = cute.make_tensor( + tStS.iterator + tmem_skip_softmax_offset, tStS_vec_layout + ) + tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) + tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) + tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) + tmem_p_offset = self.tmem_p0_offset if stage == 0 else self.tmem_p1_offset + tStS_P = cute.make_tensor(tStS.iterator + tmem_p_offset, tStS_P_layout) + if cutlass.const_expr(self.arch >= Arch.sm_100 and self.arch <= Arch.sm_100f): + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.qk_acc_dtype, + ) + else: + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(32), redOp=tcgen05.TmemLoadRedOp.MAX + ), + self.qk_acc_dtype, + ) + + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStSi) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_tmem_load.partition_S(tStSi) + tmem_store_vec_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) + thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(thread_idx) + tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) + tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) + tTMEM_STORE_SKIP_SOFTMAX = thr_tmem_store_vec.partition_D(tStS_skip_softmax) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), + self.qk_acc_dtype, + ) + tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) + thr_tmem_store = tiled_tmem_store.get_slice(thread_idx) + tTMEM_STOREtS_x4 = thr_tmem_store.partition_D(tStS_P) + if cutlass.const_expr(self.enable_sequence_barrier): + if cutlass.const_expr(stage == 1): + self.sequence_s0_s1_barrier.arrive() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + seqlen_k_ = seqlen_k + seqlen_q_ = seqlen_q + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q_ = seqlen_q + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q_ = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = ( + not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q_, + ) + ) + if not continue_cond: + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord].to(Int32) + seqlen_k_ = cum_seqlen_k[batch_coord + 1].to(Int32) - cuseqlen_k + continue_cond = seqlen_k_ <= 0 + if not continue_cond: + logical_offset = ( + curr_block_coord[0] * self.cta_tiler[0] + stage * self.qk_mma_tiler[0], + 0, + ) + cS = cute.domain_offset(logical_offset, cS_base) + value_args_ = ( + seqlen_k_, + seqlen_q_, + scale_softmax_log2, + window_size_left, + window_size_right, + skip_softmax_threshold_log2, + thread_idx, + logical_offset, + ) + atom_args = ( + qk_thr_mma, + tiled_tmem_load, + tiled_tmem_store, + tiled_tmem_store_vec, + thr_tmem_load, + thr_tmem_store, + thr_tmem_store_vec, + ) + tensor_args_ = ( + tTMEM_LOADtS, + tTMEM_STORE_VECtS, + tTMEM_STORE_SKIP_SOFTMAX, + tTMEM_STOREtS_x4, + warp_wants_skip_softmax_exchange, + skip_softmax_count, + total_softmax_count, + ) + st_cnt, end_cnt, ld_mask_cnt, unmask_cnt, tl_mask_cnt = ( + fmha_utils.FusedMask.get_masked_info( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q_, + seqlen_k_, + window_size_left, + window_size_right, + ) + ) + row_max = -Float32.inf + row_sum = 0.0 + stats_args = (row_sum, row_max) + + def softmax_loop( + whether_apply_mask: bool, + loop_args: Tuple, + stats_args: Tuple, + pipeline_args: Tuple, + inner_fn: Callable, + value_args: Tuple, + atom_args: Tuple, + tensor_args: Tuple, + cS: cute.Tensor, + ) -> Tuple[Tuple, Tuple]: + start_index, iter_num, upper_bound = loop_args + for i in cutlass.range(start_index, start_index + iter_num, 1, unroll=1): + cS_iter = cute.domain_offset((0, i * self.qk_mma_tiler[1]), cS) + iter_args = (cS_iter, i == upper_bound - 1) + stats_args, pipeline_args = inner_fn( + stage, + whether_apply_mask, + iter_args, + stats_args, + pipeline_args, + value_args, + atom_args, + tensor_args, + ) + return stats_args, pipeline_args + + softmax_step_fn = self.softmax_step + softmax_loop_fn = partial( + softmax_loop, + inner_fn=softmax_step_fn, + value_args=value_args_, + atom_args=atom_args, + tensor_args=tensor_args_, + cS=cS, + ) + si_peek_status = mma_si_consumer.try_wait() + if cutlass.const_expr(stage == 1): + inplace_consumer.wait_and_advance() + pipeline_args_ = ( + si_peek_status, + mma_si_consumer, + si_corr_producer, + pi_mma_producer, + inplace_producer, + inplace_consumer, + ) + # 1. Leading mask loop + loop_args = (st_cnt, ld_mask_cnt, end_cnt) + stats_args, pipeline_args_ = softmax_loop_fn( + True, loop_args, stats_args, pipeline_args_ + ) + # 2. Unmasked loop + loop_args = (st_cnt + ld_mask_cnt, unmask_cnt, end_cnt) + stats_args, pipeline_args_ = softmax_loop_fn( + False, loop_args, stats_args, pipeline_args_ + ) + # 3. Trailing mask loop + loop_args = (st_cnt + ld_mask_cnt + unmask_cnt, tl_mask_cnt, end_cnt) + stats_args, pipeline_args_ = softmax_loop_fn( + True, loop_args, stats_args, pipeline_args_ + ) + + # Unpack pipeline_args + ( + _, + mma_si_consumer, + si_corr_producer, + pi_mma_producer, + inplace_producer, + inplace_consumer, + ) = pipeline_args_ + if cutlass.const_expr(stage == 0): + inplace_producer.commit() + inplace_producer.advance() + # 4. Copy the final stats for correction epilog + tTMEM_STORE_VECrS = cute.make_rmem_tensor( + tTMEM_STORE_VECcS.shape, self.qk_acc_dtype + ) + tTMEM_STORE_VECrS[0] = stats_args[0] + tTMEM_STORE_VECrS[1] = stats_args[1] + vec_i_handle = si_corr_producer.acquire_and_advance() + cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) + cute.arch.fence_view_async_tmem_store() + vec_i_handle.commit() + # End of if not continue_cond + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + + @cute.jit + def correction_rescale( + self, + thr_mma: cute.ThrMma, + tiled_tmem_load_vec: cute.TiledCopy, + scale_softmax_log2: Float32, + tensor_args: Tuple, + pipeline_args: Tuple, + oi_peek_status=None, + ): + """Rescale intermediate attention results based on softmax normalization factor. + + This method performs a crucial correction step in the attention computation pipeline. + When processing attention in blocks, the softmax normalization factors may change + as new blocks are processed. This method rescales previously computed partial + output values to account for updated normalization factors. + + The implementation uses efficient tensor memory operations to: + 1. Load existing partial attention output from tensor memory + 2. Apply the scaling factor to all elements + 3. Store the rescaled results back to tensor memory + + When ``self.enable_correction_double_buffer`` is True, the rescale loop + runs as a 2-buffer tensor-memory load / multiply / store pipeline. + + :param thr_mma: Thread MMA operation for the computation + :type thr_mma: cute.ThrMma + :param tiled_tmem_load_vec: Tiled memory load operation for the vectorized row-wise max + :type tiled_tmem_load_vec: cute.TiledCopy + :param scale_softmax_log2: Log2 of the softmax factor + :type scale_softmax_log2: Float32 + :param tensor_args: Tuple containing the tensors for the correction + :type tensor_args: Tuple[cute.Tensor, cute.Tensor, cute.Tensor] + :param pipeline_args: Tuple containing the pipeline arguments for the correction + :type pipeline_args: Tuple[pipeline.PipelineConsumer, pipeline.PipelineConsumer] + :param oi_peek_status: Optional non-blocking token for the Oi consumer + wait. ``None`` or ``False`` falls back to a blocking wait. + :return: ``((si_corr_consumer, mma_corr_consumer), next_oi_peek_status)`` + where ``next_oi_peek_status`` is the peek for the next Oi (only + refreshed when a rescale actually ran; otherwise stays + ``cutlass.Boolean(False)``). + """ + tOtO, tTMEM_LOAD_VECtSi, tTMEM_LOAD_VECcS = tensor_args + si_corr_consumer, mma_corr_consumer = pipeline_args + + pv_tiled_mma_shape = ( + self.pv_mma_tiler[0], + self.pv_mma_tiler[1], + ) + cO = cute.make_identity_tensor(pv_tiled_mma_shape) + tOcO = thr_mma.partition_C(cO) + corr_tile_size = 16 # tuneable parameter + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.pv_acc_dtype, + ) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.pv_acc_dtype, + ) + tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) + tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) + tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) + tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i) + tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_i) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.correction_warp_ids)) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + thr_tmem_store = tiled_tmem_store.get_slice(thread_idx) + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i) + tTMEM_LOADcO = thr_tmem_load.partition_D(tOcO_i) + tTMEM_STOREtO = thr_tmem_store.partition_D(tOtO_i) + num_tiles = self.cta_tiler[2] // corr_tile_size + if cutlass.const_expr(self.enable_correction_double_buffer): + # Double buffer: 2 register buffers to pipeline tensor-memory load / multiply / store. + tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, 2), self.pv_acc_dtype) + # Rank-matching views for cute.copy (TMEM partitions are rank-3, + # raw tTMrO[None, idx] is rank-1; composition restores the rank). + copy_layout = cute.make_layout(tTMrO.shape[0]) + view_0 = tTMrO[None, 0] + view_1 = tTMrO[None, 1] + tTMrO_copy = ( + cute.make_tensor( + view_0.iterator, + cute.composition(view_0.layout, copy_layout), + ), + cute.make_tensor( + view_1.iterator, + cute.composition(view_1.layout, copy_layout), + ), + ) + else: + tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, num_tiles), self.pv_acc_dtype) + tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) + # Wait for vec_i (row_wise current max & previous max) + vec_i_handle = si_corr_consumer.wait_and_advance() + cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtSi, tTMEM_LOAD_VECrS) + cute.arch.fence_view_async_tmem_load() + vec_i_handle.release() + # Wait for Oi (peek-token aware: None/False falls back to blocking). + oi_handle = mma_corr_consumer.wait_and_advance(oi_peek_status) + next_oi_peek_status = cutlass.Boolean(False) + vote_ballot_cnt = cute.arch.vote_ballot_sync(tTMEM_LOAD_VECrS[0] != tTMEM_LOAD_VECrS[1]) + should_rescale = vote_ballot_cnt != 0 + if should_rescale: + scale_ = scale_softmax_log2 * (tTMEM_LOAD_VECrS[0] - tTMEM_LOAD_VECrS[1]) + scale = cute.math.exp2(scale_, fastmath=True) + if cutlass.const_expr(self.enable_correction_double_buffer): + num_elems = cute.size(tTMrO, mode=[0]) + # Prologue: load first tile into buffer 0 + cute.copy(tiled_tmem_load, tTMEM_LOADtO, tTMrO_copy[0]) + # Steady state. Refresh the next Oi probe near the end of the + # current rescale pipeline. + for i in cutlass.range_constexpr(1, num_tiles): + cute.copy( + tiled_tmem_load, + cute.make_tensor( + tTMEM_LOADtO.iterator + i * corr_tile_size, + tTMEM_LOADtO.layout, + ), + tTMrO_copy[i % 2], + ) + for j in range(0, num_elems, 2): + tTMrO[j, (i - 1) % 2], tTMrO[j + 1, (i - 1) % 2] = ( + cute.arch.mul_packed_f32x2( + ( + tTMrO[j, (i - 1) % 2], + tTMrO[j + 1, (i - 1) % 2], + ), + (scale, scale), + ) + ) + cute.copy( + tiled_tmem_store, + tTMrO_copy[(i - 1) % 2], + cute.make_tensor( + tTMEM_STOREtO.iterator + (i - 1) * corr_tile_size, + tTMEM_STOREtO.layout, + ), + ) + next_oi_peek_status = mma_corr_consumer.try_wait() + # Epilogue: compute and store last tile + last = (num_tiles - 1) % 2 + for j in range(0, num_elems, 2): + tTMrO[j, last], tTMrO[j + 1, last] = cute.arch.mul_packed_f32x2( + (tTMrO[j, last], tTMrO[j + 1, last]), + (scale, scale), + ) + cute.copy( + tiled_tmem_store, + tTMrO_copy[last], + cute.make_tensor( + tTMEM_STOREtO.iterator + (num_tiles - 1) * corr_tile_size, + tTMEM_STOREtO.layout, + ), + ) + else: + for i in cutlass.range_constexpr(0, num_tiles): + tTMrO_i_ = tTMrO[None, i] + tTMrO_i = cute.make_tensor( + tTMrO_i_.iterator, + cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])), + ) + cute.copy( + tiled_tmem_load, + cute.make_tensor( + tTMEM_LOADtO.iterator + i * corr_tile_size, + tTMEM_LOADtO.layout, + ), + tTMrO_i, + ) + for j in range(0, cute.size(tTMrO_i), 2): + tTMrO_i[j], tTMrO_i[j + 1] = cute.arch.mul_packed_f32x2( + (tTMrO_i[j], tTMrO_i[j + 1]), + (scale, scale), + ) + cute.copy( + tiled_tmem_store, + tTMrO_i, + cute.make_tensor( + tTMEM_STOREtO.iterator + i * corr_tile_size, + tTMEM_STOREtO.layout, + ), + ) + next_oi_peek_status = mma_corr_consumer.try_wait() + # Release Oi + cute.arch.fence_view_async_tmem_store() + oi_handle.release() + return (si_corr_consumer, mma_corr_consumer), next_oi_peek_status + + @cute.jit + def correction_epilog( + self, + thr_mma: cute.ThrMma, + tiled_tmem_load_vec: cute.TiledCopy, + tensor_args: Tuple, + pipeline_args: Tuple, + value_args: Tuple, + ): + """Apply final scaling and transformation to attention output. + + When use_tma_store=True: writes to shared memory and signals epilogue warp for TMA store. + When use_tma_store=False: writes directly to global memory via st.global. + + :param thr_mma: Thread MMA operation for the computation + :type thr_mma: cute.ThrMma + :param tiled_tmem_load_vec: Tiled memory load operation for the vectorized row-wise max + :type tiled_tmem_load_vec: cute.TiledCopy + :param tensor_args: Tuple containing output, softmax state, destination, + optional LSE/sink, and optional in-place merge tensors. + :type tensor_args: Tuple + :param pipeline_args: When use_tma_store: (si_corr_consumer, mma_corr_consumer, corr_epi_producer). + When not use_tma_store: (si_corr_consumer, mma_corr_consumer). + :type pipeline_args: Tuple + :param value_args: Tuple containing (row_idx, cuseqlen_q, seqlen_q, blk_coord, scale_softmax, scale_output) + :type value_args: Tuple + """ + ( + tOtO, + tTMEM_LOAD_VECtSi, + tTMEM_LOAD_VECcS, + dest_O, + sOldO, + mLSE, + mSink, + mMergeStats, + mMergeOp, + old_o_tma_args, + ) = tensor_args + ( + row_idx, + cuseqlen_q, + seqlen_q, + blk_coord, + scale_softmax, + scale_output, + old_o_tma_phase, + ) = value_args + + pv_tiled_mma_shape = ( + self.pv_mma_tiler[0], + self.pv_mma_tiler[1], + ) + cO = cute.make_identity_tensor(pv_tiled_mma_shape) + + corr_tile_size = 32 * 8 // self.o_dtype.width + tOdO = thr_mma.partition_C(dest_O) + tOcO = thr_mma.partition_C(cO) + tOtO_i = cute.logical_divide(tOtO, cute.make_layout((128, corr_tile_size))) + tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size))) + tOdO_i = cute.logical_divide(tOdO, cute.make_layout((128, corr_tile_size))) + + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.correction_warp_ids)) + epi_subtile = (self.epi_tile[0], corr_tile_size) + tmem_copy_atom = tcgen05_utils.get_tmem_load_op( + self.pv_mma_tiler, + self.o_layout, + self.o_dtype, + self.pv_acc_dtype, + epi_subtile, + use_2cta_instrs=False, + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0]) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i[(None, None), None]) + tTMEM_LOADdO = thr_tmem_load.partition_D(tOdO_i[(None, None), None]) + tTMEM_LOADoO = thr_tmem_load.partition_D(tOcO_i[(None, None), None]) + + if cutlass.const_expr(self.use_tma_store): + si_corr_consumer, mma_corr_consumer, corr_epi_producer = pipeline_args + smem_copy_atom = tcgen05_utils.get_smem_store_op( + self.o_layout, self.o_dtype, self.pv_acc_dtype, tiled_tmem_load + ) + tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load) + else: + si_corr_consumer, mma_corr_consumer = pipeline_args + + merge_op = Int64(2) + previous_max = Float32(0.0) + previous_sum = Float32(0.0) + if cutlass.const_expr(mMergeStats is not None): + batch_idx = blk_coord[2][1] + merge_op = mMergeOp[batch_idx] + if merge_op == 1: + if row_idx < seqlen_q: + stats_row = row_idx + cuseqlen_q + previous_max = mMergeStats[stats_row, blk_coord[2], 0] + previous_sum = mMergeStats[stats_row, blk_coord[2], 1] + + # Wait for vec_i (row_wise global sum) + vec_i_handle = si_corr_consumer.wait_and_advance() + tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) + cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtSi, tTMEM_LOAD_VECrS) + cute.arch.fence_view_async_tmem_load() + vec_i_handle.release() + + # Wait for Oi + oi_handle = mma_corr_consumer.wait_and_advance() + if cutlass.const_expr(self.use_tma_store): + oi_final_handle = corr_epi_producer.acquire_and_advance() + row_sum = tTMEM_LOAD_VECrS[0] + current_output_scale = scale_output + if cutlass.const_expr(mMergeStats is not None): + if cutlass.const_expr(self.v_dtype.width == 8 and self.p_fp8_prescale_log2 > 0): + p_fp8_prescale_inv = Float32(2.0 ** (-self.p_fp8_prescale_log2)) + row_sum = row_sum * p_fp8_prescale_inv + current_output_scale = scale_output * p_fp8_prescale_inv + if cutlass.const_expr(mSink is not None): + sink_val = mSink[blk_coord[2]] + row_max_raw = tTMEM_LOAD_VECrS[1] + # sink is already in scaled logit space, row_max_raw is unscaled + # exp2((sink - max_scaled) * log2(e)) = exp(sink - max_scaled) + log2_e = Float32(1.4426950408889634) + sink_exp = cute.math.exp2( + (sink_val - row_max_raw * scale_softmax) * log2_e, fastmath=True + ) + row_sum = row_sum + sink_exp + scale = current_output_scale / row_sum + previous_scale = Float32(0.0) + merged_max = scale_softmax * tTMEM_LOAD_VECrS[1] + merged_sum = row_sum + merge_log2_e = Float32(1.4426950408889634) + if cutlass.const_expr(mMergeStats is not None): + if merge_op == 1: + if row_idx < seqlen_q: + merged_max = cute.arch.fmax(previous_max, merged_max) + previous_shift = cute.math.exp2( + (previous_max - merged_max) * merge_log2_e, + fastmath=True, + ) + current_shift = cute.math.exp2( + (scale_softmax * tTMEM_LOAD_VECrS[1] - merged_max) * merge_log2_e, + fastmath=True, + ) + previous_weight = previous_sum * previous_shift + merged_sum = previous_weight + row_sum * current_shift + previous_scale = previous_weight / merged_sum + scale = current_shift * current_output_scale / merged_sum + + if cutlass.const_expr(sOldO is not None and mMergeStats is not None): + if merge_op == 1: + full_old_o_tile = row_idx - thread_idx + self.qk_mma_tiler[0] <= seqlen_q + if full_old_o_tile: + correction_warp_idx = thread_idx // self.threads_per_warp + if correction_warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + old_o_tma_args[3], self.tma_copy_old_o_bytes + ) + cute.copy( + old_o_tma_args[0], + old_o_tma_args[1], + old_o_tma_args[2], + tma_bar_ptr=old_o_tma_args[3], + ) + cute.arch.mbarrier_wait(old_o_tma_args[3], phase=old_o_tma_phase) + old_o_tma_phase ^= 1 + else: + lane_idx = tidx % self.threads_per_warp + correction_warp_idx = thread_idx // self.threads_per_warp + old_o_vec_elems = 128 // self.o_dtype.width + old_o_row_chunks = self.head_dim_v // old_o_vec_elems + old_o_rows_per_issue = self.threads_per_warp // old_o_row_chunks + old_o_rows_per_warp = self.qk_mma_tiler[0] // len(self.correction_warp_ids) + old_o_copy_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS), + self.o_dtype, + num_bits_per_copy=128, + ) + for load_step in range(old_o_rows_per_warp // old_o_rows_per_issue): + old_o_row = ( + correction_warp_idx * old_o_rows_per_warp + + load_step * old_o_rows_per_issue + + lane_idx // old_o_row_chunks + ) + old_o_chunk = lane_idx % old_o_row_chunks + if row_idx - thread_idx + old_o_row < seqlen_q: + g_old_o_row = cute.flat_divide( + dest_O[old_o_row, None], (old_o_vec_elems,) + ) + s_old_o_row = cute.flat_divide( + sOldO[old_o_row, None], (old_o_vec_elems,) + ) + cute.copy( + old_o_copy_atom, + g_old_o_row[None, old_o_chunk], + s_old_o_row[None, old_o_chunk], + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_warp() + + for i in range(self.cta_tiler[2] // corr_tile_size): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i] + tTMEM_LOADdO_i = tTMEM_LOADdO[None, 0, 0, i] + tTMrO = cute.make_rmem_tensor(tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype) + tDMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + if cutlass.const_expr(mMergeStats is not None): + if merge_op == 1 and row_idx < seqlen_q: + # Reuse the narrow output register buffer for the old O + # load, and widen each pair only when it is consumed. A + # separate full-width previous-O tensor extends 16 FP32 + # values across the whole subtile and raises epilogue + # register pressure for the op1 path. + s_old_o_row = cute.flat_divide(sOldO[thread_idx, None], (corr_tile_size,)) + tDMrO.store(s_old_o_row[None, i].load().reshape(tDMrO.shape)) + for j in range(0, cute.size(tTMrO), 2): + current_pair = cute.arch.mul_packed_f32x2( + (tTMrO[j], tTMrO[j + 1]), + (scale, scale), + ) + previous_pair = cute.arch.mul_packed_f32x2( + (Float32(tDMrO[j]), Float32(tDMrO[j + 1])), + (previous_scale, previous_scale), + ) + tTMrO[j], tTMrO[j + 1] = cute.arch.add_packed_f32x2( + current_pair, previous_pair + ) + else: + for j in range(0, cute.size(tTMrO), 2): + tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( + (tTMrO[j], tTMrO[j + 1]), + (scale, scale), + ) + else: + for j in range(0, cute.size(tTMrO), 2): + tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( + (tTMrO[j], tTMrO[j + 1]), + (scale, scale), + ) + o_vec = tTMrO.load() + tDMrO.store(o_vec.to(self.o_dtype)) + if cutlass.const_expr(self.use_tma_store): + # TMA store path: write to shared memory + cute.copy(tiled_smem_store, tDMrO, tTMEM_LOADdO_i) + else: + # st.global path: write directly to global memory with bounds check + if cutlass.const_expr(mMergeStats is not None): + if row_idx < seqlen_q and merge_op != 0: + cute.autovec_copy(tDMrO, tTMEM_LOADdO_i) + else: + if row_idx < seqlen_q: + cute.autovec_copy(tDMrO, tTMEM_LOADdO_i) + + if cutlass.const_expr(mMergeStats is not None): + if row_idx < seqlen_q and merge_op != 0: + stats_row = row_idx + cuseqlen_q + mMergeStats[stats_row, blk_coord[2], 0] = merged_max + mMergeStats[stats_row, blk_coord[2], 1] = merged_sum + + if cutlass.const_expr(mLSE is not None): + scaled_tmp = scale_softmax * tTMEM_LOAD_VECrS[1] + # Convert LSE from natural log to log2 space, consistent with flashinfer trtllm backend + lse = (cute.math.log(row_sum, fastmath=True) + scaled_tmp) * Float32(1.4426950408889634) + # Pre-scale correction: row_sum was inflated by 2^offset, so the + # log2-space LSE is too large by exactly `p_fp8_prescale_log2`. + if cutlass.const_expr(self.v_dtype.width == 8 and self.p_fp8_prescale_log2 > 0): + lse = lse - self.p_fp8_prescale_log2 + if row_idx < seqlen_q: + mLSE[row_idx + cuseqlen_q, blk_coord[2]] = lse + if cutlass.const_expr(self.use_tma_store): + # fence view async shared + cute.arch.fence_view_async_shared() + oi_handle.release() + oi_final_handle.commit() + return (si_corr_consumer, mma_corr_consumer, corr_epi_producer) + else: + oi_handle.release() + return (si_corr_consumer, mma_corr_consumer, old_o_tma_phase) + + def check_supported_dtypes( + self, + qk_dtype: Type[cutlass.Numeric], + pv_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + ): + supported_input_dtypes = { + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float16, + cutlass.BFloat16, + } + if qk_dtype not in supported_input_dtypes: + raise cute_testing.CantImplementError( + f"Unsupported qk_dtype on Rubin SM107: {qk_dtype}" + ) + if pv_dtype not in supported_input_dtypes: + raise cute_testing.CantImplementError( + f"Unsupported pv_dtype on Rubin SM107: {pv_dtype}" + ) + # The output is cast from the FP32 accumulator independently of the MMA atom. + if out_dtype not in { + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float16, + cutlass.BFloat16, + }: + raise cute_testing.CantImplementError( + f"Unsupported out_dtype on Rubin SM107: {out_dtype}" + ) + if qk_acc_dtype is not cutlass.Float32: + raise cute_testing.CantImplementError("Unsupported qk_acc_dtype on Rubin SM107") + if pv_acc_dtype is not cutlass.Float32: + raise cute_testing.CantImplementError("Unsupported pv_acc_dtype on Rubin SM107") + + def check_invalid_shape( + self, + qk_dtype: Type[cutlass.Numeric], + q_shape: Tuple[int, int, int, int], + k_shape: Tuple[int, int, int, int], + ): + # Shapes are passed around this example as (batch, seq_len, num_heads, head_dim). + b, s_q, h_q, d = q_shape + b_, s_k, h_k, d_ = k_shape + + if b != b_: + raise testing.CantImplementError("q & k must have the same batch size") + if d != d_: + raise testing.CantImplementError("q & k must have the same head dimension") + if d not in {32, 64, 128, 192}: + raise testing.CantImplementError("Unsupported head dimension") + if h_q % h_k != 0: + raise testing.CantImplementError("h_q must be divisible by h_k") + if isinstance(s_q, tuple) and len(s_q) != b: + raise testing.CantImplementError( + "variable_seqlen s_q must have the length of batch size" + ) + if isinstance(s_k, tuple) and len(s_k) != b: + raise testing.CantImplementError( + "variable_seqlen s_k must have the length of batch size" + ) + + def can_implement( + self, + q_shape: Tuple[int, int, int, int], + k_shape: Tuple[int, int, int, int], + qk_dtype: Type[cutlass.Numeric], + pv_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + ) -> bool: + """ + :param q_shape: Shape of the query tensor. + :type q_shape: Tuple[int, int, int, int] + :param k_shape: Shape of the key tensor. + :type k_shape: Tuple[int, int, int, int] + :param qk_dtype: Data type for Q and K (Bmm1 inputs). + :type qk_dtype: Type[cutlass.Numeric] + :param pv_dtype: Data type for P and V (Bmm2 inputs). + :type pv_dtype: Type[cutlass.Numeric] + :param out_dtype: Data type of the output tensor. + :type out_dtype: Type[cutlass.Numeric] + :param qk_acc_dtype: Data type of the qk accumulator tensor. + :type qk_acc_dtype: Type[cutlass.Numeric] + :param pv_acc_dtype: Data type of the pv accumulator tensor. + :type pv_acc_dtype: Type[cutlass.Numeric] + :return: True if the kernel can be implemented, False otherwise. + :rtype: bool + """ + try: + # Skip unsupported types + self.check_supported_dtypes( + qk_dtype, + pv_dtype, + out_dtype, + qk_acc_dtype, + pv_acc_dtype, + ) + # Skip invalid shape + self.check_invalid_shape( + qk_dtype, + q_shape, + k_shape, + ) + except testing.CantImplementError: + return False + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha_helpers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha_helpers.py new file mode 100644 index 000000000000..56a0858a5f2d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha_helpers.py @@ -0,0 +1,1240 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Runtime helper subset for the Kimi K3 Rubin FP8 context FMHA kernel (``k3_context_fmha``).""" + +import enum +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir as _mlir_ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.typing import Boolean +from cutlass.cutlass_dsl import ( + Float32, + Int32, + T, + dsl_user_op, + extract_mlir_values, + min, + new_from_mlir_values, +) +from cutlass.utils.hardware_info import HardwareInfo + +# The internal CuTe DSL build's scheduling hints (``dsl_hints``), used here and by the kernel as +# ``fmha_utils.()``. They only shape the SASS schedule of the softmax / correction warps. +from .dsl_hints import ( # noqa: F401 + HAS_SCHEDULING_HINTS, + cfence, + reset_sched_res_busy_xu64, + sched_res_busy_xu64, + warp_switch, +) + + +class WorkTileInfo: + """A class to represent information about a work tile. + + :ivar tile_idx: The index of the tile. + :type tile_idx: cute.Coord + :ivar is_valid_tile: Whether the tile is valid. + :type is_valid_tile: Boolean + """ + + def __init__(self, tile_idx: cute.Coord, is_valid_tile: Boolean): + self._tile_idx = tile_idx + self._is_valid_tile = Boolean(is_valid_tile) + + def __extract_mlir_values__(self) -> list[_mlir_ir.Value]: + values = extract_mlir_values(self.tile_idx) + values.extend(extract_mlir_values(self.is_valid_tile)) + return values + + def __new_from_mlir_values__(self, values: list[_mlir_ir.Value]) -> "WorkTileInfo": + assert len(values) == 4 + new_tile_idx = new_from_mlir_values(self._tile_idx, values[:-1]) + new_is_valid_tile = new_from_mlir_values(self._is_valid_tile, [values[-1]]) + return WorkTileInfo(new_tile_idx, new_is_valid_tile) + + @property + @cute.jit + def is_valid_tile(self) -> Boolean: + """Check latest tile returned by the scheduler is valid or not. Any scheduling + requests after all tasks completed will return an invalid tile. + + :return: The validity of the tile. + :rtype: Boolean + """ + return self._is_valid_tile + + @property + @cute.jit + def tile_idx(self) -> cute.Coord: + """ + Get the index of the tile. + + :return: The index of the tile. + :rtype: cute.Coord + """ + return self._tile_idx + + +############################################################################## +# Fmha static tile scheduler +############################################################################## + + +class FmhaStaticTileSchedulerParams: + """A class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler. + + This class holds the configuration parameters needed to initialize and configure + the tile scheduler for FMHA operations. + + :ivar is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :ivar problem_shape_mhb: Problem shape in (M, H, B) format. + :type problem_shape_mhb: cute.Shape + """ + + def __init__( + self, + is_persistent: bool, + problem_shape_mhb: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the FmhaStaticTileSchedulerParams with the given parameters. + + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param problem_shape_mhb: Problem shape in (M, H, B) format. + :type problem_shape_mhb: cute.Shape + """ + self.is_persistent = is_persistent + self.problem_shape_mhb = problem_shape_mhb + self._loc = loc + self._ip = ip + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.problem_shape_mhb]: + obj_values = extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.problem_shape_mhb], self._values_pos): + obj_list.append(new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return FmhaStaticTileSchedulerParams(self.is_persistent, *(tuple(obj_list)), loc=self._loc) + + +class FmhaStaticTileScheduler: + """A static tile scheduler for FMHA (Fused Multi-Head Attention) operations. + + This class manages the scheduling of work tiles for FMHA kernels, supporting + both persistent and non-persistent kernel modes. It tracks the current work + position and advances through the problem space efficiently. + + :ivar _params: Scheduler parameters. + :type _params: FmhaStaticTileSchedulerParams + :ivar _blk_coord: Block coordinates. + :type _blk_coord: cute.Coord + :ivar _grid_shape: Grid shape for the kernel. + :type _grid_shape: cute.Shape + :ivar _is_persistent: Whether to use persistent kernel mode. + :type _is_persistent: bool + :ivar _current_work_linear_idx: Current linear work index. + :type _current_work_linear_idx: Int32 + :ivar _problem_shape_mhb: Problem shape in (M, H, B) format. + :type _problem_shape_mhb: cute.Layout + :ivar _num_blocks: Number of blocks in the problem. + :type _num_blocks: Int32 + :ivar _is_first_block: Whether this is the first block. + :type _is_first_block: bool + :ivar num_persistent_sm: Number of persistent SMs. + :type num_persistent_sm: Int32 + """ + + def __init__( + self, + params: FmhaStaticTileSchedulerParams, + current_work_linear_idx: Int32, + blk_coord: cute.Coord, + grid_shape: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the FmhaStaticTileScheduler with the given parameters. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + :param current_work_linear_idx: Current linear work index. + :type current_work_linear_idx: Int32 + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param grid_shape: Grid shape for the kernel. + :type grid_shape: cute.Shape + """ + self._params = params + self._blk_coord = blk_coord + self._grid_shape = grid_shape + self._is_persistent = params.is_persistent + self._current_work_linear_idx = current_work_linear_idx + self._problem_shape_mhb = cute.make_layout(params.problem_shape_mhb, loc=loc, ip=ip) + self._num_blocks = cute.size(self._problem_shape_mhb, loc=loc, ip=ip) + self._is_first_block = True + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + self._loc = loc + self._ip = ip + + # called by host + @staticmethod + def get_grid_shape( + params: FmhaStaticTileSchedulerParams, + *, + sm_count: int | None = None, + loc=None, + ip=None, + ) -> cute.Shape: + """ + Determine the grid shape for the FMHA kernel. + + For persistent kernels, the grid shape is limited by the number of SMs + (Streaming Multiprocessors) available on the device. For non-persistent + kernels, the grid shape matches the problem shape. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + + :return: Grid shape as (M, H, B) tuple. + :rtype: cute.Shape + """ + if params.is_persistent: + if sm_count is None: + hardware_info = HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + return ( + min(sm_count, cute.size(params.problem_shape_mhb, loc=loc, ip=ip)), + 1, + 1, + ) + else: + return params.problem_shape_mhb + + @staticmethod + def check_valid_work_for_seqlen_q( + q_tiler: int, + current_idx: Int32, + seqlen_q: Int32, + ) -> Boolean: + """ + Check if the current work index is valid for the given query sequence length. + + This method verifies that the current work tile index multiplied by the + query tiler size is within the bounds of the query sequence length. + + :param q_tiler: Query tiler size. + :type q_tiler: int + :param current_idx: Current work index. + :type current_idx: Int32 + :param seqlen_q: Query sequence length. + :type seqlen_q: Int32 + + :return: True if the work is valid, False otherwise. + :rtype: Boolean + """ + return current_idx * q_tiler < seqlen_q + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + """ + Get information about the current work tile. + + Determines if the current work is valid and computes the tile coordinates + based on whether the kernel is persistent or non-persistent. + + :return: WorkTileInfo containing tile coordinates and validity flag. + :rtype: WorkTileInfo + """ + is_valid = ( + self._current_work_linear_idx < self._num_blocks + if self._is_persistent + else self._is_first_block + ) + + blk_coord = (0, 0, 0) + if self._is_persistent: + blk_coord = self._problem_shape_mhb.get_hier_coord( + self._current_work_linear_idx, loc=loc, ip=ip + ) + else: + blk_coord = self._blk_coord + + # cur_tile_coord is (mid, 0, (bid, hid)) + cur_tile_coord = ( + blk_coord[0], + 0, + (blk_coord[1], blk_coord[2]), + ) + + return WorkTileInfo(cur_tile_coord, is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + """ + Get the initial work tile information. + + :return: Initial WorkTileInfo. + :rtype: WorkTileInfo + """ + return self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + """ + Advance to the next work tile. + + For persistent kernels, advances by the number of persistent SMs. + For non-persistent kernels, marks that the first block has been processed. + + :param advance_count: Number of steps to advance (default: 1). + :type advance_count: int + """ + if self._is_persistent: + self._current_work_linear_idx += advance_count * self.num_persistent_sm + self._is_first_block = False + + def __extract_mlir_values__(self): + # Only pass mutable per-iteration state as scf.while block arguments. + # _params and _grid_shape are loop-invariant and captured from outer + # scope, keeping block arg count low (4 instead of 10). + values = extract_mlir_values(self._current_work_linear_idx) + values.extend(extract_mlir_values(self._blk_coord)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 4 + new_current_work_linear_idx = new_from_mlir_values( + self._current_work_linear_idx, [values[0]] + ) + new_blk_coord = new_from_mlir_values(self._blk_coord, values[1:4]) + return FmhaStaticTileScheduler( + self._params, new_current_work_linear_idx, new_blk_coord, self._grid_shape + ) + + +def create_fmha_static_tile_scheduler( + params: FmhaStaticTileSchedulerParams, + blk_coord: cute.Coord, + grid_shape: cute.Shape, +) -> FmhaStaticTileScheduler: + """ + Create a new FMHA static tile scheduler. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param grid_shape: Grid shape. + :type grid_shape: cute.Shape + + :return: New FmhaStaticTileScheduler instance. + :rtype: FmhaStaticTileScheduler + """ + return FmhaStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape) + + +def create_fmha_static_tile_scheduler_params( + is_persistent: bool, + problem_shape_mhb: cute.Shape, +) -> FmhaStaticTileSchedulerParams: + """ + Create FMHA static tile scheduler parameters. + + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param problem_shape_mhb: Problem shape in (M, H, B) format. + :type problem_shape_mhb: cute.Shape + + :return: New FmhaStaticTileSchedulerParams instance. + :rtype: FmhaStaticTileSchedulerParams + """ + return FmhaStaticTileSchedulerParams(is_persistent, problem_shape_mhb) + + +def compute_grid( + o_shape: cute.Shape, + cta_tiler: Tuple[int, int, int], + is_persistent: bool, + is_2cta: bool = False, +) -> Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]: + """ + Compute grid parameters for FMHA operation. + + This function calculates the appropriate grid shape and scheduler parameters + based on the output tensor shape, CTA (Cooperative Thread Array) tiler, + and whether to use persistent kernel mode. + + The output tensor o has shape (s, d, ((h_r, h_k), b)) where: + - s: sequence length + - d: head dimension + - h_r: number of heads for query + - h_k: number of heads for key + - b: batch size + + :param o_shape: Output tensor shape for grid computation. + :type o_shape: cute.Shape + :param cta_tiler: CTA tiler dimensions (M, N, K). + :type cta_tiler: Tuple[int, int, int] + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param is_2cta: Whether to use 2CTA mode. + :type is_2cta: bool + + :return: Tuple of (scheduler_params, grid_shape). + :rtype: Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]] + """ + tile_sched_params = create_fmha_static_tile_scheduler_params( + is_persistent, + ( + cute.round_up(cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]), 2 if is_2cta else 1), + cute.size(o_shape[2][0]), + cute.size(o_shape[2][1]), + ), + ) + grid = FmhaStaticTileScheduler.get_grid_shape(tile_sched_params) + + return tile_sched_params, grid + + +############################################################################## +# Fused Mask +############################################################################## + + +class MaskEnum(enum.Enum): + """Enumeration of mask types for FMHA operations. + + - RESIDUAL_MASK: Residual mask for handling variable sequence lengths + - WINDOW_MASK: Window mask for attention which also includes causal and no mask + - WINDOW_MASK_INFERENCE: Same as the window mask, but requires the end of q + to align with the end of k + - WINDOW_MASK_BWD: Window mask for backward pass + - WINDOW_MASK_BWD_INFERENCE: Same as the backward window mask, but requires + the end of q to align with the end of k + """ + + RESIDUAL_MASK = enum.auto() + RESIDUAL_MASK_BWD = enum.auto() + WINDOW_MASK = enum.auto() + WINDOW_MASK_INFERENCE = enum.auto() + WINDOW_MASK_BWD = enum.auto() + WINDOW_MASK_BWD_INFERENCE = enum.auto() + + +class FusedMask: + """A fused mask implementation for FMHA operations. + + This class handles different types of attention masks including no mask, + residual mask for variable sequence lengths, and causal mask for + autoregressive attention patterns. + + The class provides methods to: + - Calculate trip counts for different mask types + - Apply masks to attention scores + - Handle masked and unmasked trip calculations + """ + + def get_trip_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of trips needed for the current block. + + The trip count depends on the mask type and the block coordinates. + For causal masks, it considers the autoregressive constraint. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of trips needed. + :rtype: Int32 + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr(mask_type == MaskEnum.RESIDUAL_MASK): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + if cutlass.const_expr(mask_type is MaskEnum.RESIDUAL_MASK_BWD): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + if cutlass.const_expr( + mask_type == MaskEnum.WINDOW_MASK or mask_type == MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_right is None): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + else: + max_idx_q = (blk_coord[0] + 1) * tile_shape[0] + idx_k = max_idx_q + offset + window_size_right + tmp_blocks_k = cute.ceil_div(idx_k, tile_shape[1]) + max_blocks_k = cute.ceil_div(seqlen_k, tile_shape[1]) + result = min(max_blocks_k, tmp_blocks_k) + if cutlass.const_expr( + mask_type == MaskEnum.WINDOW_MASK_BWD or mask_type == MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_left is None): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + else: + max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + idx_k = max_idx_k + offset + window_size_left + tmp_blocks_q = cute.ceil_div(idx_k, tile_shape[0]) + max_blocks_q = cute.ceil_div(seqlen_q, tile_shape[0]) + result = min(max_blocks_q, tmp_blocks_q) + start_block = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = result - start_block + return result + + @cute.jit + def get_trip_start( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Get the start of the trip for the current block. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK or mask_type is MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + idx_k = min_idx_q + offset - window_size_left + tmp_blocks_k = idx_k // tile_shape[1] + result = max(tmp_blocks_k, result) + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK_BWD or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + idx_q = min_idx_k + offset - window_size_right + tmp_blocks_q = idx_q // tile_shape[0] + result = max(tmp_blocks_q, result) + return result + + @cute.jit + def get_leading_mask_id( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Int32, Int32]: + """ + Get the begin and end tile idx for the leading mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the leading mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + leading_mask_begin = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + leading_mask_end = leading_mask_begin + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK or mask_type is MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = (blk_coord[0] + 1) * tile_shape[0] + offset - window_size_left + leading_mask_end = min( + cute.ceil_div(min_idx_q, tile_shape[1]) - 1, + trip_count + leading_mask_begin - 1, + ) + else: + leading_mask_end = leading_mask_begin - 1 + elif cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK_BWD or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset - window_size_right + leading_mask_end = cute.ceil_div(min_idx_k, tile_shape[0]) - 1 + else: + leading_mask_end = leading_mask_begin - 1 + return leading_mask_begin, leading_mask_end + + @cute.jit + def get_trailing_mask_id( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Optional[Int32], Optional[Int32]]: + """ + Get the begin and end tile idx for the trailing mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the trailing mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + trip_start = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + trailing_mask_begin, trailing_mask_end = None, None + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK or mask_type is MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + offset + window_size_right + trailing_mask_begin = min(min_idx_q // tile_shape[1], trip_count + trip_start - 1) + trailing_mask_end = trip_count + trip_start - 1 + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + else: + if cutlass.const_expr(window_size_left is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + offset + window_size_left + 1 + max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset + window_size_left + trailing_mask_begin = min( + cute.ceil_div(min_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + trailing_mask_end = min( + cute.ceil_div(max_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + + return trailing_mask_begin, trailing_mask_end + + @cute.jit + def get_masked_leading_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of masked trips for the leading mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + if cutlass.const_expr( + mask_type is not MaskEnum.RESIDUAL_MASK and mask_type is not MaskEnum.RESIDUAL_MASK_BWD + ): + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = max(leading_mask_end - leading_mask_begin + 1, 0) + + return result + + @cute.jit + def get_masked_trailing_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + rem_count: Optional[Int32] = 0, + ) -> Int32: + """ + Calculate the number of masked trips for the trailing mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param rem_count: Remaining count from previous calculations. + :type rem_count: Int32 + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + + if cutlass.const_expr( + mask_type is not MaskEnum.RESIDUAL_MASK and mask_type is not MaskEnum.RESIDUAL_MASK_BWD + ): + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + trailing_mask_begin, trailing_mask_end = FusedMask.get_trailing_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + if cutlass.const_expr( + trailing_mask_begin is not None and trailing_mask_end is not None + ): + if trailing_mask_begin <= leading_mask_end: + result = max(trailing_mask_end - leading_mask_end, 0) + else: + result = max(trailing_mask_end - trailing_mask_begin + 1, 0) + else: + if seqlen_k % tile_shape[1] != 0: + result = 1 + else: + result = 0 + + return result + rem_count + + @cute.jit + def get_unmasked_trip_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of unmasked trips for the current block. + + This represents the number of trips that don't require special + masking treatment. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of unmasked trips. + :rtype: Int32 + """ + result = ( + FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - FusedMask.get_masked_leading_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - FusedMask.get_masked_trailing_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + 0, + ) + ) + return result + + @cute.jit + def get_masked_info( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + rem_count: Optional[Int32] = 0, + ) -> Tuple[Int32, Int32, Int32, Int32, Int32]: + """ + Calculate the number of masked trips for the trailing mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param rem_count: Remaining count from previous calculations. + :type rem_count: Int32 + + :return: Number of masked info. + :rtype: Tuple[Int32, Int32, Int32, Int32, Int32] + """ + start_count = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + leading_mask_count = FusedMask.get_masked_leading_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + unmask_count = FusedMask.get_unmasked_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trailing_mask_count = FusedMask.get_masked_trailing_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + rem_count, + ) + end_count = start_count + leading_mask_count + unmask_count + trailing_mask_count + return ( + start_count, + end_count, + leading_mask_count, + unmask_count, + trailing_mask_count, + ) + + @cute.jit + def apply_mask( + mask_type: MaskEnum, + acc_qk: cute.Tensor, + index_qk: cute.Tensor, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + index_transform: cutlass.Constexpr = lambda index_q, index_k: ( + index_q, + index_k, + ), + ): + """ + Apply the appropriate mask to the attention scores. + + This method modifies the attention scores (acc_qk) based on the mask type + and the positions in the index tensor. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param acc_qk: Accumulated QK attention scores tensor. + :type acc_qk: cute.Tensor + :param index_qk: Index tensor containing position information. + :type index_qk: cute.Tensor + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Optional[int] + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[int] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[int] + """ + + tidx, tidy, tidx = cute.arch.thread_idx() + offset = 0 + offset = ( + seqlen_k - seqlen_q + if cutlass.const_expr( + mask_type is MaskEnum.WINDOW_MASK_INFERENCE + or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE + ) + else 0 + ) + for i in cutlass.range(cute.size(acc_qk), unroll_full=True): + index_q, index_k = index_transform(*index_qk[i]) + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + if cutlass.const_expr(window_size_left is None): + if index_q + offset + window_size_right < index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + elif cutlass.const_expr(window_size_right is None): + if index_q + offset - window_size_left > index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + else: + max_K_index = min(index_q + offset + window_size_right, seqlen_k) + min_K_index = max(0, index_q + offset - window_size_left) + if index_k > max_K_index or index_k < min_K_index: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + + if cutlass.const_expr( + mask_type == MaskEnum.RESIDUAL_MASK or mask_type == MaskEnum.RESIDUAL_MASK_BWD + ): + if index_k >= seqlen_k or index_q >= seqlen_q: + acc_qk[i] = -Float32.inf + + +@dsl_user_op +def ex2_emulation_packed_f32x2( + x: Float32, y: Float32, *, loc=None, ip=None +) -> Tuple[Float32, Float32]: + # Clamp the xy so they fit within the FP32 exponent range + # The upper side is ensured by the (s - row_max) + xy_clamped = (cute.arch.fmax(x, -127.0), cute.arch.fmax(y, -127.0)) + + fp32_round_int = float(2**23 + 2**22) + # | 0 | 10010110 | 10000000000000000000000 | + # | sign | exponent | mantissa | + # ^ + # | + # digit of ones place + # During FP32 addition, any number that smaller than this will be + # aligned the exponent to 2^23, so that the digit of ones + # is at the rightest place + # We want to round down here, so that the fractional part is in [0, 1) + xy_rounded = cute.arch.add_packed_f32x2(xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm") + warp_switch() # {$nv-internal-release} + # The integer floor of x & y are now in the last 8 bits of xy_rounded + # We want the next 2 ops to round to nearest even. The rounding mode is important. + xy_rounded_back = cute.arch.sub_packed_f32x2(xy_rounded, (fp32_round_int, fp32_round_int)) + warp_switch() # {$nv-internal-release} + xy_frac = cute.arch.sub_packed_f32x2(xy_clamped, xy_rounded_back) + + @dsl_user_op + @cute.jit + def polynomial_deg3_packed_f32x2( + x: Float32, y: Float32, *, loc=None, ip=None + ) -> Tuple[Float32, Float32]: + # 2^x ~= (0.077 * x + 0.228) * x + 0.695) * x + 1, for x in [0, 1) + coeff = ( + 1.0, # coeff of deg0 + 0.695146143436431884765625, # coeff of deg1 + 0.227564394474029541015625, # coeff of deg2 + 0.077119089663028717041015625, # coeff of deg3 + ) + deg = len(coeff) - 1 # started with highest degree + out = (coeff[deg], coeff[deg]) + for i in cutlass.range_constexpr(deg - 1, -1, -1): + out = cute.arch.fma_packed_f32x2(out, (x, y), (coeff[i], coeff[i]), loc=loc, ip=ip) + warp_switch() # {$nv-internal-release} + return out + + xy_frac_ex2 = polynomial_deg3_packed_f32x2(*xy_frac, loc=loc, ip=ip) + + @dsl_user_op + def combine_int_frac_ex2( + x_rounded: Float32, frac_ex2: Float32, *, loc=None, ip=None + ) -> Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + Float32(x_rounded).ir_value(loc=loc, ip=ip), + Float32(frac_ex2).ir_value(loc=loc, ip=ip), + ], + "{\n\t" + ".reg .s32 x_rounded_i, frac_ex_i, x_rounded_e, out_i;\n\t" + "mov.b32 x_rounded_i, $1;\n\t" + "mov.b32 frac_ex_i, $2;\n\t" + "shl.b32 x_rounded_e, x_rounded_i, 23;\n\t" + "add.s32 out_i, x_rounded_e, frac_ex_i;\n\t" + "mov.b32 $0, out_i;\n\t" + "}\n", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + x_out = combine_int_frac_ex2(xy_rounded[0], xy_frac_ex2[0], loc=loc, ip=ip) + y_out = combine_int_frac_ex2(xy_rounded[1], xy_frac_ex2[1], loc=loc, ip=ip) + + return x_out, y_out + + +@cute.jit +def cvt_f32x4_to_f8x4_pack_i32(fp32x4, fp8_type, *, loc=None, ip=None): + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + + src0 = Float32(vector.extract(src_vec4, [], [0])).ir_value(loc=loc, ip=ip) + src1 = Float32(vector.extract(src_vec4, [], [1])).ir_value(loc=loc, ip=ip) + src2 = Float32(vector.extract(src_vec4, [], [2])).ir_value(loc=loc, ip=ip) + src3 = Float32(vector.extract(src_vec4, [], [3])).ir_value(loc=loc, ip=ip) + + cvt_instruction = "" + if cutlass.const_expr(fp8_type == cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.e4m3x2.f32" + else: + assert False, "Unsupported fp8 element type" + + asm_tmpl = ( + "{\n" + " .reg .b16 lo;\n" + " .reg .b16 hi;\n" + f" {cvt_instruction} lo, $2, $1;\n" + f" {cvt_instruction} hi, $4, $3;\n" + " mov.b32 $0, {lo, hi};\n" + "}" + ) + packed_i32 = llvm.inline_asm( + T.i32(), + [src0, src1, src2, src3], + asm_tmpl, + "=r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + return packed_i32 + + +@cute.jit +def cvt_f32x4_to_f8x4(fp32x4, fp8x4, *, loc=None, ip=None): + packed_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, fp8x4.element_type) + fp8x4_i32 = cute.recast_tensor(fp8x4, cutlass.Int32) + fp8x4_i32[0] = cutlass.Int32(packed_i32) + return diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha_op.py new file mode 100644 index 000000000000..caf34de77d2f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_context_fmha_op.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""":func:`run_k3_context_fmha`: the Kimi K3 Rubin FP8 split-K context FMHA. + +One launch attends a ragged batch of FP8 Q ``[T_q, H, 192]`` to one KV partition given as FP8 +K-nope ``[T_kv, H, 128]``, the shared FP8 K-pe ``[T_kv, 1, 64]`` and FP8 V ``[T_kv, H, 128]`` +(``k3_context_fmha.RubinFusedMultiHeadAttentionForward`` with ``split_k_dims=(128, 64)``), and +merges the partition into the BF16 output in its epilogue. Per request, ``merge_op`` selects +what the epilogue does with the running output ``[T_q, H, 128]`` and its softmax stats +``[T_q, H, 2]`` (row max, row sum): 2 writes the partition's result, 1 merges it with them, 0 +leaves both unchanged; the codes are those of ``TrtllmAttentionMetadata.merge_op_tensor``. A +causal partition uses a top-left aligned causal mask. + +Kernels compile on first use, one per (Q dtype, heads, causal, cumulative-KV-length dtype). +""" + +import math +from typing import ClassVar, Optional + +import torch + +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger + +try: + from cuda.bindings import driver as cuda +except ImportError: + from cuda import cuda + +IS_K3_CONTEXT_FMHA_AVAILABLE = False +K3_CONTEXT_FMHA_IMPORT_ERROR: Optional[ImportError] = None +# Whether the compiled kernel carries the internal CuTe DSL build's scheduling hints. +K3_CONTEXT_FMHA_SCHEDULING_HINTS = False + +try: + import cutlass + import cutlass.cute as cute + import cutlass.torch as cutlass_torch + from cutlass.cute.runtime import from_dlpack + from cutlass.cute.typing import Float32, Int32 + + from . import k3_context_fmha_helpers as fmha_utils + from .k3_context_fmha import RubinFusedMultiHeadAttentionForward + + IS_K3_CONTEXT_FMHA_AVAILABLE = bool( + getattr( + RubinFusedMultiHeadAttentionForward, + "supports_split_k_epilogue_merge", + False, + ) + ) + K3_CONTEXT_FMHA_SCHEDULING_HINTS = fmha_utils.HAS_SCHEDULING_HINTS +except ImportError as error: + K3_CONTEXT_FMHA_IMPORT_ERROR = error + + +if IS_K3_CONTEXT_FMHA_AVAILABLE: + + class CuteDSLK3ContextFmhaRunner: + """Compile and launch the Rubin FP8 split-K context FMHA.""" + + kernel_cache: ClassVar[dict] = {} + + @classmethod + def run( + cls, + q: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + merge_stats: torch.Tensor, + merge_op: torch.Tensor, + max_q_len: int, + max_kv_len: int, + softmax_scale: float, + is_causal: bool, + ) -> None: + num_tokens_q, num_heads, head_dim = q.shape + num_tokens_kv = k_nope.shape[0] + batch_size = merge_op.numel() + + q_5d = q.view(1, num_tokens_q, num_heads, 1, head_dim) + k_nope_5d = k_nope.view(1, num_tokens_kv, num_heads, 1, 128) + k_pe_5d = k_pe.view(1, num_tokens_kv, 1, 1, 64) + v_5d = v.view(1, num_tokens_kv, num_heads, 1, 128) + output_5d = output.view(1, num_tokens_q, num_heads, 1, 128) + + cache_key = (q.dtype, num_heads, is_causal, cu_kv_seqlens.dtype) + compiled_fmha = cls.kernel_cache.get(cache_key) + problem_size = ( + batch_size, + max_q_len, + num_tokens_q, + max_kv_len, + num_heads, + num_heads, + 192, + 128, + ) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + scale = Float32(softmax_scale) + + if compiled_fmha is None: + logger.info( + "Compiling Rubin K3 FP8 context FMHA: split K=(128,64), " + "epilogue merge, causal=%s, heads=%d, scheduling hints=%s", + is_causal, + num_heads, + K3_CONTEXT_FMHA_SCHEDULING_HINTS, + ) + mask_type = ( + fmha_utils.MaskEnum.WINDOW_MASK + if is_causal + else fmha_utils.MaskEnum.RESIDUAL_MASK + ) + fmha = RubinFusedMultiHeadAttentionForward( + cutlass.Float32, + cutlass.Float32, + (128, 128), + (192, 128), + False, + mask_type, + False, + False, + use_tma_store=False, + split_k_dims=(128, 64), + ) + + def as_dynamic_tensor(tensor: torch.Tensor, align: int = 16): + return from_dlpack(tensor, assumed_align=align).mark_layout_dynamic( + leading_dim=cutlass_torch.get_leading_dim(tensor) + ) + + compiled_fmha = cute.compile( + fmha, + as_dynamic_tensor(q_5d), + as_dynamic_tensor(k_nope_5d), + as_dynamic_tensor(k_pe_5d), + as_dynamic_tensor(v_5d), + as_dynamic_tensor(output_5d), + problem_size, + as_dynamic_tensor(cu_q_seqlens, 4), + as_dynamic_tensor(cu_kv_seqlens, 4), + None, + None, + as_dynamic_tensor(merge_stats), + as_dynamic_tensor(merge_op, 8), + Float32(softmax_scale * math.log2(math.e)), + scale, + Float32(1.0), + None, + None, + Int32(0) if is_causal else None, + None, + None, + stream, + False, + options="--opt-level=2 --ptxas-options '--uumn'", + ) + cls.kernel_cache[cache_key] = compiled_fmha + + compiled_fmha( + q_5d, + k_nope_5d, + k_pe_5d, + v_5d, + output_5d, + problem_size, + cu_q_seqlens, + cu_kv_seqlens, + None, + None, + merge_stats, + merge_op, + softmax_scale * math.log2(math.e), + softmax_scale, + 1.0, + None, + None, + 0 if is_causal else None, + None, + None, + stream, + False, + ) + + def _validate_k3_context_fmha_inputs( + q: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + merge_stats: torch.Tensor, + merge_op: torch.Tensor, + ) -> None: + if get_sm_version() != 107: + raise ValueError("K3 context FMHA requires Rubin SM107.") + if any(tensor.dtype != torch.float8_e4m3fn for tensor in (q, k_nope, k_pe, v)): + raise ValueError("K3 context FMHA requires FP8 E4M3 Q/K/V inputs.") + if output.dtype != torch.bfloat16: + raise ValueError("K3 context FMHA requires BF16 output.") + if q.dim() != 3 or q.shape[1] <= 0 or q.shape[2] != 192: + raise ValueError(f"Expected q [T, H, 192] with H > 0, got {tuple(q.shape)}.") + num_heads = q.shape[1] + if k_nope.dim() != 3 or k_nope.shape[1:] != (num_heads, 128): + raise ValueError(f"Expected k_nope [T, {num_heads}, 128], got {tuple(k_nope.shape)}.") + if k_pe.dim() != 3 or k_pe.shape[1:] != (1, 64): + raise ValueError(f"Expected k_pe [T, 1, 64], got {tuple(k_pe.shape)}.") + if v.dim() != 3 or v.shape[1:] != (num_heads, 128): + raise ValueError(f"Expected v [T, {num_heads}, 128], got {tuple(v.shape)}.") + if output.shape != (q.shape[0], num_heads, 128): + raise ValueError( + f"Expected output {(q.shape[0], num_heads, 128)}, got {tuple(output.shape)}." + ) + if merge_stats.shape != (q.shape[0], num_heads, 2): + raise ValueError( + f"Expected merge_stats {(q.shape[0], num_heads, 2)}, got {tuple(merge_stats.shape)}." + ) + if merge_stats.dtype != torch.float32: + raise ValueError("merge_stats must be float32.") + if merge_op.dtype != torch.int64 or merge_op.dim() != 1: + raise ValueError("merge_op must be a one-dimensional int64 tensor.") + expected_cu_shape = (merge_op.numel() + 1,) + for name, cu_seqlens, dtypes in ( + ("cu_q_seqlens", cu_q_seqlens, (torch.int32,)), + ("cu_kv_seqlens", cu_kv_seqlens, (torch.int32, torch.int64)), + ): + if cu_seqlens.dtype not in dtypes or cu_seqlens.shape != expected_cu_shape: + raise ValueError( + f"{name} must have dtype in {dtypes} and shape {expected_cu_shape}, got " + f"dtype={cu_seqlens.dtype}, shape={tuple(cu_seqlens.shape)}." + ) + if k_nope.shape[0] != k_pe.shape[0] or k_nope.shape[0] != v.shape[0]: + raise ValueError("K-nope, K-pe, and V token counts must match.") + for name, tensor in (("k_nope", k_nope), ("v", v)): + if tensor.stride()[1:] != (128, 1): + raise ValueError( + f"{name} must have contiguous head dimensions, got stride {tensor.stride()}." + ) + if k_pe.stride(2) != 1: + raise ValueError( + f"k_pe must have a contiguous head dimension, got stride {k_pe.stride()}." + ) + for name, tensor in ( + ("q", q), + ("output", output), + ("merge_stats", merge_stats), + ): + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous.") + + def run_k3_context_fmha( + q: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + merge_stats: torch.Tensor, + merge_op: torch.Tensor, + max_q_len: int, + max_kv_len: int, + softmax_scale: float, + is_causal: bool, + ) -> None: + """Validate the inputs, then attend ``q`` to one KV partition and merge it into + ``output`` / ``merge_stats`` per ``merge_op`` (see the module docstring).""" + _validate_k3_context_fmha_inputs( + q, + k_nope, + k_pe, + v, + output, + cu_q_seqlens, + cu_kv_seqlens, + merge_stats, + merge_op, + ) + CuteDSLK3ContextFmhaRunner.run( + q, + k_nope, + k_pe, + v, + output, + cu_q_seqlens, + cu_kv_seqlens, + merge_stats, + merge_op, + max_q_len, + max_kv_len, + softmax_scale, + is_causal, + ) + + +__all__ = [ + "IS_K3_CONTEXT_FMHA_AVAILABLE", + "K3_CONTEXT_FMHA_IMPORT_ERROR", + "K3_CONTEXT_FMHA_SCHEDULING_HINTS", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fp8w_lowm_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fp8w_lowm_gemm.py new file mode 100644 index 000000000000..677a53b3f738 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fp8w_lowm_gemm.py @@ -0,0 +1,4493 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M FP8-weight GEMM for weight-bandwidth-bound decode projections (SM100/SM103/SM107). + +Computes ``out[M, N] = x[M, K] @ (W[N, K] * S)^T`` where ``W`` is FP8 E4M3 with one +power-of-two (UE8M0) scale per 128x128 block and ``x`` is BF16 with ``M <= MAX_M``. +``N`` only needs to be a multiple of :data:`N_ALIGN`: the weight TMA zero-fills the rows +of a partial last 128-row tile past ``N`` and the epilogues do not store them. + +Low-M kernel (``M <= low_m_max_for(N, K) <= LOW_M_MAX``; kernel coordinates are swapped: +the weight's N rows are the UMMA M side): + +* One CTA owns a 128-row weight tile and a contiguous K slice. The ``split_k`` CTAs that + share a weight tile form a cluster and reduce-scatter their FP32 partials through DSMEM + (one bulk copy per peer; each rank owns every ``split_k``-th 4-token column chunk). +* The MMA is ``tcgen05.mma.kind::mxf8f6f4.block_scale`` with 1x32 UE8M0 scale vectors. + The activation is quantized inside the kernel to MXFP8 (E4M3 + one power-of-two scale + per 1x32 group). Because a CTA's weight tile lies inside a single 128-row scale block, + the weight block scale only depends on the K tile and is folded into the activation + scale (``sfb = e_act + e_w - 127``); the weight scale operand is a constant 1.0. The + weight therefore streams as raw FP8 codes with no expanded scale tensor. +* Warp roles: warp 0 streams the weight with TMA and never waits on the previous grid + (PDL: with enough stages a CTA prefetches most of its weight slice while the previous + kernel finishes), warp 1 issues the MMAs, warp 2 loads BF16 activation tiles with TMA + after ``griddepcontrol.wait``, and ``quant_warps`` warps starting at warp 4 quantize + activation tiles (round-robin over ``quant_groups`` groups) into the swizzled UMMA B + layout. Warps 4-7 then run the epilogue (TMEM -> registers -> DSMEM reduce-scatter -> + BF16 store). +* Shared memory is reused across phases: the activation ring becomes the reduction mailbox + once every peer has finished quantizing, and the weight ring stages outgoing partials. + +Mid-M path (``low_m_max_for(N, K) < M <= MAX_M``): + +* A small kernel quantizes the activation to MXFP8 (FP8 codes + 1x32 UE8M0 scales already + in the UMMA scale-factor layout). The GEMM is chained to it with PDL: its weight TMA warp + starts streaming before the quantization finishes and only the activation producer waits. +* The GEMM uses the same swap-AB split-K cluster structure with UMMA N padded to a multiple + of 32 tokens. The activation tile arrives by TMA (FP8) plus one bulk copy of its scales; + the weight block exponent goes into the A scale operand, refreshed once per K tile. +* The epilogue double-buffers 32-column TMEM loads, reduce-scatters partials with + ``st.async`` into the owner's DSMEM mailbox, and writes BF16 rows through a token-major + shared staging buffer with 16-byte stores. +* :func:`can_run` only accepts M up to :func:`max_m_for` (the measured cross-over against + ``fp8_swap_ab_gemm`` on SM100/SM103, against the CuTe MXFP8 GEMM on SM107); the op itself + accepts any ``M <= MAX_M``. + +Both kernels launch one non-persistent grid (``N/128`` weight tiles x ``split_k`` cluster +ranks), so the per-shape tactics, cut-overs and gates come from per-architecture tables +(:func:`_tables`): GB300 for SM100/SM103 and Rubin (212 SMs) for SM107. + +:func:`run_k3_fp8w_lowm_gemm_fused` runs either kernel with a fused epilogue instead of the BF16 +store (whole 128-row weight tiles only): the owner of each token chunk stages its BF16 rows +token-major in the dead activation ring, then writes the NVFP4 quantization of them +(``fp4_quantize`` arithmetic); the results equal the BF16 GEMM at the same tactic followed by +``fp4_quantize``. The fused Kimi-K3 latent_down projection (``k3_latent_fp8_gemm``) runs it on +SM107. + +Sigmoid epilogue (:func:`k3_fp8w_lowm_gemm_sigmoid`): ``torch.sigmoid`` of the BF16 output, +bit for bit, for a gate projection. + +Gated low-M kernel (:func:`k3_fp8w_lowm_gemm_gated`): ``out = bf16(x * gate) @ W^T`` for an +output-gated projection. The activation producer loads a BF16 gate tile next to each +activation tile (same barrier) and the quantization warps multiply before quantizing, with +PyTorch's BF16 multiply rounding, so it is bitwise equal to the multiply followed by the +plain kernel. + +All-reduce push epilogue (:func:`run_k3_fp8w_lowm_gemm_ar_push`, low-M kernel): instead of +storing its BF16 output, the kernel stores it into the Lamport slots of the next +MNNVL all-reduce on a workspace (:mod:`k3_lamport_ar`): each token row goes to its owner's +scatter slot (two-shot) or to every rank's slot (one-shot, multicast), with ``-0`` sent as +``+0`` as the C++ all-reduce sends it. The epilogue stages the BF16 tile token-major in the dead +weight ring and pushes 16-byte vectors; with split-K (SM107 only, see :func:`ar_push_tactic`) each +cluster rank stages and pushes the token rows it owns after the reduction, in the dead +quantized-activation ring like the fused epilogues. :func:`k3_lamport_ar.run_lamport_finish` +completes the +all-reduce, so the pair is bitwise equal to this GEMM followed by the MNNVL all-reduce of its +output. PDL: this variant releases its dependent (the finish kernel) from the activation warp +right after its own ``griddepcontrol.wait`` instead of from the weight warp at kernel start. The +finish kernel reads the workspace's rotation state as soon as it is resident, which is then +safe whatever precedes this GEMM: once this GEMM has passed its wait, every earlier grid has +completed (each waits for its predecessor before it completes), the previous all-reduce on the +workspace included. The weight warp still streams before the wait. +""" + +from __future__ import annotations + +import dataclasses +import functools +from typing import Optional, Tuple + +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 +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass import Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +from . import k3_lamport_ar +from .k3_latent_fp8_gemm import ( + _abs_max_f32x16, + _e2m1x8, + _ldg_f32, + _nvfp4_scales_x2, + _rcp_approx_ftz, +) + +#: K extent of one pipeline stage (one 128x128 FP8 weight block). +CTA_K = 128 +#: Weight rows per CTA tile (UMMA M); equals the weight scale block height. +MMA_M = 128 +#: MXFP8 scale vector length. +SF_VEC = 32 +#: Granularity of N. The last 128-row weight tile may be partial: TMA zero-fills its missing +#: rows and the epilogue does not store them. A multiple of 16 keeps every BF16 output row +#: 16-byte aligned and the DeepGEMM-packed scale exactly N rows wide (TMA-aligned to 4). +N_ALIGN = 16 +#: Supported token tiles (UMMA N) of the low-M kernel; public M is padded up to one of these. +SUPPORTED_MMA_N = (8, 16, 32, 64) +#: Largest public M served by the low-M kernel (in-kernel activation quantization). +LOW_M_MAX = 64 +#: Largest public M the op accepts (larger M than the low-M kernel's runs the mid-M path). +MAX_M = 256 +#: ``ar_push`` modes of the all-reduce push epilogue (the C++ kernel the call stands in for). +AR_PUSH_TWO_SHOT = 1 +AR_PUSH_ONE_SHOT = 2 + +_TMA_WARP = 0 +_MMA_WARP = 1 +_X_WARP = 2 +_QUANT_WARP0 = 4 +_QUANT_BARRIER_ID = 1 +_QUANT_DONE_BARRIER_ID = 2 +_EPI_BARRIER_ID = 3 +_TMEM_COLS = 128 +_SF_TMEM_COLS = 16 +#: L2 eviction hint for the streamed weight (CacheHintSm90::EVICT_FIRST). +_EVICT_FIRST = 0x12F0000000000000 +#: BF16 bit pattern of ~1e-4; activation groups below this amax share its scale. +_AMAX_FLOOR_BF16 = 0x38D2 + +_SMEM_CAPACITY = utils.get_smem_capacity_in_bytes("sm_100") + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def _align_up(a: int, b: int) -> int: + return _ceil_div(a, b) * b + + +@dataclasses.dataclass(frozen=True) +class K3LowMTactic: + """One kernel specialization. + + ``mma_n`` is the padded token count, ``split_k`` the cluster size along K, and the + stage counts size the weight ring (``a``), activation BF16 ring (``x``) and + quantized activation ring (``b``). + """ + + mma_n: int + split_k: int + a_stages: int + x_stages: int = 4 + b_stages: int = 4 + quant_warps: int = 4 + quant_groups: int = 4 + + +def smem_bytes(n: int, k: int, tactic: K3LowMTactic, gated: bool = False) -> int: + """Conservative host-side mirror of the kernel's shared-memory footprint. + + ``gated`` adds the gate ring of the gated kernel (one BF16 gate tile per activation + stage). + """ + kt_max = _ceil_div(k // CTA_K, tactic.split_k) + a_st = min(tactic.a_stages, kt_max) + x_st = min(tactic.x_stages, kt_max) + b_st = min(tactic.b_stages, kt_max) + mbox = (tactic.split_k - 1) * _ceil_div(tactic.mma_n // 4, tactic.split_k) * MMA_M * 16 + total = 8 * (2 * (a_st + x_st + b_st) + 4) + 4 + 4 * kt_max + # The weight ring doubles as the staging buffer for outgoing split-K partials. + staging = tactic.split_k * _ceil_div(tactic.mma_n // 4, tactic.split_k) * MMA_M * 16 + total = _align_up(total, 1024) + max(a_st * MMA_M * CTA_K, staging if tactic.split_k > 1 else 0) + total = _align_up(total, 1024) + b_st * tactic.mma_n * CTA_K + # The reduction mailbox aliases the activation ring (dead once quantization is done). + total = _align_up(total, 128) + max(x_st * tactic.mma_n * CTA_K * 2, mbox) + total = _align_up(total, 128) + 512 + total = _align_up(total, 128) + b_st * 512 + if gated: + total = _align_up(total, 128) + x_st * tactic.mma_n * CTA_K * 2 + return total + 1024 # allocator slack + + +def validate_tactic(n: int, k: int, tactic: K3LowMTactic, gated: bool = False) -> None: + if tactic.mma_n not in SUPPORTED_MMA_N: + raise ValueError(f"mma_n must be one of {SUPPORTED_MMA_N}, got {tactic.mma_n}") + if n % N_ALIGN or k % CTA_K: + raise ValueError(f"N={n} must be a multiple of {N_ALIGN} and K={k} of {CTA_K}") + if not 1 <= tactic.split_k <= min(16, k // CTA_K): + raise ValueError(f"split_k={tactic.split_k} invalid for K={k}") + for name in ("a_stages", "x_stages", "b_stages"): + if getattr(tactic, name) < 1: + raise ValueError(f"{name} must be >= 1") + groups, warps = tactic.quant_groups, tactic.quant_warps + if warps < 1 or groups < 1 or warps % groups or (warps // groups) * 32 > tactic.mma_n * 4: + raise ValueError(f"invalid quant_warps={warps} / quant_groups={groups}") + if 32 * (4 + warps) > 1024: + raise ValueError(f"quant_warps={warps} exceeds the CTA size limit") + kt_max = _ceil_div(k // CTA_K, tactic.split_k) + for name in ("x_stages", "b_stages"): + # Each quantization group must own its ring slots, or a group could wait on a + # slot two phases ahead (mbarrier parity aliasing). + stages = getattr(tactic, name) + if stages < kt_max and stages % groups: + raise ValueError(f"{name}={stages} must be a multiple of quant_groups={groups}") + need = smem_bytes(n, k, tactic, gated) + if need > _SMEM_CAPACITY: + raise ValueError(f"tactic {tactic} needs {need} B of SMEM > {_SMEM_CAPACITY} B") + + +def max_a_stages(n: int, k: int, tactic: K3LowMTactic, gated: bool = False) -> int: + """Largest weight-ring depth that fits next to the other buffers.""" + kt_max = _ceil_div(k // CTA_K, tactic.split_k) + best = 1 + for a in range(1, kt_max + 1): + if smem_bytes(n, k, dataclasses.replace(tactic, a_stages=a), gated) <= _SMEM_CAPACITY: + best = a + return best + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +@dsl_user_op +def _lds_v4_b32(addr, *, loc=None, ip=None): + """Load 16 bytes from shared memory as four 32-bit words.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)) + + +@dsl_user_op +def _lds_v4_f32(addr, *, loc=None, ip=None): + """Load four FP32 values from shared memory.""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + loc=loc, + ip=ip, + ) + return tuple( + cutlass.Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(4) + ) + + +@dsl_user_op +def _sts_v2_b32(addr, v0, v1, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(v0, loc=loc, ip=ip), _iv(v1, loc=loc, ip=ip)], + "st.shared.v2.b32 [$0], {$1, $2};", + "r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_v4_b32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_iv(x, loc=loc, ip=ip) for x in (addr, v0, v1, v2, v3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _abs_max_bf16x32(words, *, loc=None, ip=None): + """|x| max of 32 BF16 values (16 packed words) as two packed BF16 magnitudes. + + ``max.xorsign.abs`` returns ``max(|a|, |b|)`` with a garbage sign bit, which + ``_ue8m0_of_amax`` masks off. + """ + body = "".join(f"max.xorsign.abs.bf16x2 m{i}, ${2 * i + 1}, ${2 * i + 2};\n" for i in range(8)) + body += ( + "max.xorsign.abs.bf16x2 m0, m0, m1;\n" + "max.xorsign.abs.bf16x2 m2, m2, m3;\n" + "max.xorsign.abs.bf16x2 m4, m4, m5;\n" + "max.xorsign.abs.bf16x2 m6, m6, m7;\n" + "max.xorsign.abs.bf16x2 m0, m0, m2;\n" + "max.xorsign.abs.bf16x2 m4, m4, m6;\n" + "max.xorsign.abs.bf16x2 $0, m0, m4;\n" + ) + return Int32( + _asm( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in words], + "{\n.reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n" + body + "}", + "=r," + ",".join(["r"] * 16), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ue8m0_of_amax(packed_amax, *, loc=None, ip=None): + """Biased exponent of ``2^ceil(log2(amax / 448))`` from two packed BF16 magnitudes. + + For ``amax = 2^E * (1 + f)``: ``amax / 448 <= 2^(E - 8)`` iff ``1 + f <= 1.75``, + so the exponent is ``E - 8`` plus one when the 7-bit mantissa exceeds ``0x60``. + """ + return Int32( + _asm( + T.i32(), + [_iv(packed_amax, loc=loc, ip=ip)], + "{\n" + ".reg .b32 lo, hi, a, e, mant;\n" + ".reg .pred p;\n" + "and.b32 lo, $1, 0x7fff;\n" + "shr.u32 hi, $1, 16;\n" + "and.b32 hi, hi, 0x7fff;\n" + "max.u32 a, lo, hi;\n" + f"max.u32 a, a, {_AMAX_FLOOR_BF16};\n" + "shr.u32 e, a, 7;\n" + "and.b32 mant, a, 0x7f;\n" + "setp.gt.u32 p, mant, 0x60;\n" + "sub.u32 e, e, 8;\n" + "@p add.u32 e, e, 1;\n" + "mov.b32 $0, e;\n" + "}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bf16x2_inv_scale(ue8m0, *, loc=None, ip=None): + """Packed BF16x2 ``2^(127 - ue8m0)``.""" + return Int32( + _asm( + T.i32(), + [_iv(ue8m0, loc=loc, ip=ip)], + "{\n" + ".reg .b32 s, t;\n" + "sub.u32 s, 254, $1;\n" + "shl.b32 s, s, 7;\n" + "shl.b32 t, s, 16;\n" + "or.b32 $0, s, t;\n" + "}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _quant_bf16x8_to_e4m3(w0, w1, w2, w3, inv_scale, *, loc=None, ip=None): + """Scale eight BF16 values by the packed BF16x2 ``inv_scale`` and round to E4M3. + + The scale is a power of two, so the BF16 multiply is exact and the only rounding + is the satfinite round-to-nearest conversion to E4M3. + """ + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_iv(x, loc=loc, ip=ip) for x in (w0, w1, w2, w3, inv_scale)], + "{\n" + ".reg .b32 p0, p1, p2, p3;\n" + ".reg .b16 h0, h1, h2, h3;\n" + "mul.rn.bf16x2 p0, $2, $6;\n" + "mul.rn.bf16x2 p1, $3, $6;\n" + "mul.rn.bf16x2 p2, $4, $6;\n" + "mul.rn.bf16x2 p3, $5, $6;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h0, p0;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h1, p1;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h2, p2;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h3, p3;\n" + "mov.b32 $0, {h0, h1};\n" + "mov.b32 $1, {h2, h3};\n" + "}", + "=r,=r,r,r,r,r,r", + loc=loc, + ip=ip, + ) + return ( + Int32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)), + Int32(llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _gate_mul_bf16x8(x0, x1, x2, x3, g0, g1, g2, g3, *, loc=None, ip=None): + """Eight BF16 products ``bf16(float(x) * float(g))`` as four packed words. + + The same arithmetic as PyTorch's BF16 ``x * g`` (FP32 multiply, round to nearest even), + so the gated activation is bit-identical to the one a separate multiply would produce. + """ + body = "" + for i in range(4): + body += ( + f"shl.b32 xl, ${i + 4}, 16;\n" + f"and.b32 xh, ${i + 4}, 0xffff0000;\n" + f"shl.b32 gl, ${i + 8}, 16;\n" + f"and.b32 gh, ${i + 8}, 0xffff0000;\n" + "mov.b32 fxl, xl;\n" + "mov.b32 fxh, xh;\n" + "mov.b32 fgl, gl;\n" + "mov.b32 fgh, gh;\n" + "mul.rn.f32 fxl, fxl, fgl;\n" + "mul.rn.f32 fxh, fxh, fgh;\n" + f"cvt.rn.bf16x2.f32 ${i}, fxh, fxl;\n" + ) + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(v, loc=loc, ip=ip) for v in (x0, x1, x2, x3, g0, g1, g2, g3)], + "{\n.reg .b32 xl, xh, gl, gh;\n.reg .f32 fxl, fxh, fgl, fgh;\n" + body + "}", + "=r,=r,=r,=r,r,r,r,r,r,r,r,r", + loc=loc, + ip=ip, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)) + + +@dsl_user_op +def _div_rn_f32(a, b, *, loc=None, ip=None): + """IEEE ``a / b`` (``div.rn.f32``).""" + return cutlass.Float32( + _asm( + T.f32(), + [ + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), + ], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +def sigmoid_bf16(value): + """``torch.sigmoid`` of a BF16 value, bit for bit: ``bf16(1 / (1 + expf(-float(v))))``. + + PyTorch evaluates the BF16 sigmoid in FP32 with the accurate ``expf`` (libdevice + ``__nv_expf``) and an IEEE division, then rounds to nearest even; ``cute.math.exp`` without + ``fastmath`` lowers to the same libdevice function. Exhaustively checked over all BF16 + inputs by the unit tests. + """ + g = value.to(cutlass.Float32) + e = cute.math.exp(-g, fastmath=False) + return _div_rn_f32(cutlass.Float32(1.0), cutlass.Float32(1.0) + e).to(cutlass.BFloat16) + + +@dsl_user_op +def _sts_scale(addr, act_exp, w_exp, *, loc=None, ip=None): + """Store the combined UE8M0 scale ``act_exp + w_exp - 127`` (clamped) as one byte.""" + _asm( + None, + [_iv(x, loc=loc, ip=ip) for x in (addr, act_exp, w_exp)], + "{\n" + ".reg .s32 v;\n" + "add.s32 v, $1, $2;\n" + "sub.s32 v, v, 127;\n" + "max.s32 v, v, 0;\n" + "min.s32 v, v, 254;\n" + "st.shared.u8 [$0], v;\n" + "}", + "r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_v4_f32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [ + _iv(addr, loc=loc, ip=ip), + *(cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)), + ], + "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", + "r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_copy_to_peer(dst_local, src, nbytes, mbar_local, peer, *, loc=None, ip=None): + """Bulk-copy ``nbytes`` of local SMEM into a peer CTA's SMEM, crediting its mbarrier. + + ``dst_local`` and ``mbar_local`` are offsets in this CTA's SMEM window; they are mapped + to the peer's window, which has the same layout. + """ + _asm( + None, + [_iv(x, loc=loc, ip=ip) for x in (dst_local, src, nbytes, mbar_local, peer)], + "{\n" + ".reg .b32 rd, rb;\n" + "mapa.shared::cluster.u32 rd, $0, $4;\n" + "mapa.shared::cluster.u32 rb, $3, $4;\n" + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes " + "[rd], [$1], $2, [rb];\n" + "}", + "r,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_commit(*, loc=None, ip=None): + _asm(None, [], "cp.async.bulk.commit_group;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _bulk_wait_read(*, loc=None, ip=None): + """Wait until committed bulk copies have finished reading their SMEM sources.""" + _asm(None, [], "cp.async.bulk.wait_group.read 0;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _mbarrier_arrive_remote(local_mbar, peer, *, loc=None, ip=None): + """Relaxed arrive on the mbarrier at the same SMEM offset in a peer CTA. + + Only used to signal that reads of the local activation ring are done; those loads + have already returned (their values were consumed), so no release fence is needed. + """ + _asm( + None, + [_iv(local_mbar, loc=loc, ip=ip), _iv(peer, loc=loc, ip=ip)], + "{\n" + ".reg .b32 rb;\n" + "mapa.shared::cluster.u32 rb, $0, $1;\n" + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [rb];\n" + "}", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.wait;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.launch_dependents;", "", loc=loc, ip=ip) + + +# --------------------------------------------------------------------------- +# Output epilogues +# --------------------------------------------------------------------------- + +#: BF16 ``out [M, N]`` (the GEMM itself). +EPI_BF16 = "bf16" +#: NVFP4 of the BF16 result, as ``fp4_quantize(out, global_sf, 16)`` with the linear scale +#: layout: E2M1 codes ``[M, N / 2]`` (low nibble first) and E4M3 scales ``[M, N / 16]``. +EPI_NVFP4 = "nvfp4" +_EPILOGUES = (EPI_BF16, EPI_NVFP4) + + +@dsl_user_op +def _bf16x2_to_f32x2(word, *, loc=None, ip=None): + """The two BF16 halves of ``word`` (low half first) widened to FP32 (exact).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_iv(word, loc=loc, ip=ip)], + "{\n" + ".reg .b32 lo, hi;\n" + "shl.b32 lo, $2, 16;\n" + "and.b32 hi, $2, 0xffff0000;\n" + "mov.b32 $0, lo;\n" + "mov.b32 $1, hi;\n" + "}", + "=f,=f,r", + loc=loc, + ip=ip, + ) + return tuple( + cutlass.Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(2) + ) + + +@dsl_user_op +def _stg_u16(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u16 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _add_bf16x2(a, b, *, loc=None, ip=None): + """Lane-wise ``bf16_rn(float(a) + float(b))`` of two BF16 pairs (an FP32 round-to-nearest add, + as PyTorch adds BF16 tensors).""" + return Int32( + _asm( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "{\n.reg .b32 al, ah, bl, bh;\n.reg .f32 fal, fah, fbl, fbh, sl, sh;\n" + "shl.b32 al, $1, 16;\nand.b32 ah, $1, 0xffff0000;\n" + "shl.b32 bl, $2, 16;\nand.b32 bh, $2, 0xffff0000;\n" + "mov.b32 fal, al;\nmov.b32 fah, ah;\nmov.b32 fbl, bl;\nmov.b32 fbh, bh;\n" + "add.rn.f32 sl, fal, fbl;\nadd.rn.f32 sh, fah, fbh;\n" + "cvt.rn.bf16x2.f32 $0, sh, sl;\n}", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +class _FusedStoreMixin: + """Final stores of the NVFP4 epilogue from a token-major SMEM tile. + + Both GEMMs stage the BF16 result of their 128-row weight tile as token rows of 256 bytes + (``stage_base + t * 256``, column ``c`` of the tile at byte ``2 c``) before this runs. + Row ``r < nrows`` of the work list is token ``r`` (``split_k == 1``) or the ``r``-th + token this cluster rank owns (chunk ``rank + (r // 4) * split_k``, 4 tokens per chunk). + """ + + @cute.jit + def _owned_token(self, r, rank): + t = r + if cutlass.const_expr(self.split_k > 1): + t = 4 * (rank + (r >> 2) * self.split_k) + (r & 3) + return t + + @cute.jit + def fused_store(self, stage_base, nrows, rank, m, n_tile, q, out0, out1, aux0, aux1): + """NVFP4: ``out0`` codes, ``out1`` scales, ``aux0`` global scale (FP32); ``aux1`` unused.""" + n = self.n + col0 = n_tile * MMA_M + gsf = _ldg_f32(aux0.toint()) + rgsf = _rcp_approx_ftz(gsf) + rcp6 = _rcp_approx_ftz(cutlass.Float32(6.0)) + # One item = 32 columns of one token: two 16-element NVFP4 groups. + for p in cutlass.range(q, nrows * 4, MMA_M, unroll=1): + t = self._owned_token(p >> 2, rank) + g = p & 3 + if t < m: + src = stage_base + t * (MMA_M * 2) + g * 64 + vals = [] + for v in cutlass.range_constexpr(4): + words = _lds_v4_b32(src + v * 16) + for i in cutlass.range_constexpr(4): + lo, hi = _bf16x2_to_f32x2(words[i]) + vals.extend((lo, hi)) + amax0 = _abs_max_f32x16(vals[:16]) + amax1 = _abs_max_f32x16(vals[16:]) + sf2, os0, os1 = _nvfp4_scales_x2(amax0, amax1, gsf, rgsf, rcp6) + col = col0 + g * 32 + _stg_v4_b32( + out0.toint() + Int64(t) * (n // 2) + Int64(col // 2), + _e2m1x8(vals[0:8], os0), + _e2m1x8(vals[8:16], os0), + _e2m1x8(vals[16:24], os1), + _e2m1x8(vals[24:32], os1), + ) + _stg_u16(out1.toint() + Int64(t) * (n // 16) + Int64(col // 16), sf2) + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +class K3Fp8wLowMGemmKernel(_FusedStoreMixin): + """Swap-AB MXFP8 GEMM with in-kernel activation quantization and cluster split-K. + + ``epilogue`` selects the output (:data:`EPI_BF16`, :data:`EPI_NVFP4`); the NVFP4 one + stages the BF16 tile token-major in the dead quantized-activation ring and needs whole + 128-row weight tiles. + + ``gated`` builds the gated variant (:meth:`call_gated`): the activation producer also + loads a BF16 gate tile per stage, and the quantization warps quantize + ``bf16(x * gate)`` instead of ``x``. ``sigmoid`` stores ``sigmoid_bf16`` of each BF16 + output instead of the output. Both go with the BF16 epilogue only. + ``ar_push`` (``AR_PUSH_TWO_SHOT`` / ``AR_PUSH_ONE_SHOT``, ``ar_world`` ranks) replaces the + BF16 store by the all-reduce push (:meth:`call_ar_push`, ``_push_lamport``, or with split-K + ``_push_owned_rows``); it needs the BF16 epilogue and no ``sigmoid``. + ``split_src = (ldw, ldx)`` (:meth:`call_split_ar_push`) reads the weight with row stride + ``ldw`` and the activation with row stride ``ldx``: column slices of wider tensors, read in + place. It needs the BF16 epilogue (or the push) without ``gated`` / ``sigmoid``. + ``push_addend`` (with the push): every pushed BF16 value ``y`` becomes ``bf16(y + a)`` + with ``a`` from a BF16 ``[M, N]`` tensor (the PyTorch BF16 add). + """ + + def __init__( + self, + n: int, + k: int, + tactic: K3LowMTactic, + use_pdl: bool, + epilogue: str = EPI_BF16, + gated: bool = False, + sigmoid: bool = False, + ar_push: int = 0, + ar_world: int = 0, + split_src: Optional[Tuple[int, int]] = None, + push_addend: bool = False, + ar_slot_n: int = 0, + ): + validate_tactic(n, k, tactic, gated) + _validate_epilogue(n, epilogue) + if epilogue != EPI_BF16 and (gated or sigmoid): + raise ValueError("the fused epilogues are not combined with gated / sigmoid") + if ar_push: + if ar_push not in (AR_PUSH_TWO_SHOT, AR_PUSH_ONE_SHOT) or ar_world < 2: + raise ValueError(f"invalid all-reduce push {ar_push} over {ar_world} ranks") + if sigmoid or epilogue != EPI_BF16: + raise ValueError("the all-reduce push needs the BF16 epilogue without sigmoid") + # Row width (features) of one Lamport slot of the push: N, or more where ranks with + # different N push into slots of one size (the latent all-gather of k3_ull_front). + if ar_slot_n and (ar_slot_n < n or ar_slot_n % 8): + raise ValueError(f"ar_slot_n={ar_slot_n} must be a multiple of 8 and >= N={n}") + self.ar_slot_n = ar_slot_n or n + if split_src is not None: + ldw, ldx = split_src + if ldw < k or ldw % 16 or ldx < k or ldx % 8: + raise ValueError(f"split_src: invalid row strides ldw={ldw} ldx={ldx}") + if gated or sigmoid or epilogue != EPI_BF16: + raise ValueError("split_src goes with the BF16 epilogue or the push only") + if push_addend and not ar_push: + raise ValueError("push_addend goes with the all-reduce push") + self.split_src = split_src + self.push_addend = push_addend + self.epilogue_kind = epilogue + self.ar_push = ar_push + self.ar_world = ar_world + self.gated = gated + self.sigmoid = sigmoid + self.n = n + self.k = k + self.mma_n = tactic.mma_n + self.split_k = tactic.split_k + self.kt_total = k // CTA_K + self.kt_max = _ceil_div(self.kt_total, self.split_k) + self.a_stages = min(tactic.a_stages, self.kt_max) + self.x_stages = min(tactic.x_stages, self.kt_max) + self.b_stages = min(tactic.b_stages, self.kt_max) + self.n_tiles = _ceil_div(n, MMA_M) + self.use_pdl = use_pdl + self.chunks_cap = _ceil_div(self.mma_n // 4, self.split_k) + mailbox_bytes = (self.split_k - 1) * self.chunks_cap * MMA_M * 16 + self.xm_bytes = max(self.x_stages * self.mma_n * CTA_K * 2, mailbox_bytes) + self.staging_bytes = self.split_k * self.chunks_cap * MMA_M * 16 if self.split_k > 1 else 0 + self.mma_tiler = (MMA_M, self.mma_n, CTA_K) + self.a_stage_bytes = MMA_M * CTA_K + self.b_stage_bytes = self.mma_n * CTA_K + self.x_stage_bytes = self.mma_n * CTA_K * 2 + self.sf_stage_bytes = 512 + self.quant_warps = tactic.quant_warps + self.quant_groups = tactic.quant_groups + self.warps_per_group = tactic.quant_warps // tactic.quant_groups + self.quant_threads = 32 * tactic.quant_warps + self.threads = 32 * _QUANT_WARP0 + self.quant_threads + if epilogue != EPI_BF16 and self.b_stages * self.b_stage_bytes < self.mma_n * MMA_M * 2: + raise ValueError("the fused epilogues stage the output tile in >= 2 b_stages") + staged = self.b_stages * self.b_stage_bytes >= self.mma_n * MMA_M * 2 + if ar_push and self.split_k > 1 and not staged: + raise ValueError("the split-K all-reduce push stages its rows in >= 2 b_stages") + + @cute.jit + def __call__( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ): + """``out_ptr``: the BF16 output, or the NVFP4 codes (``out1_ptr`` scales, ``aux0_ptr`` + global scale).""" + self._launch( + w_ptr, x_ptr, None, None, ws_ptr, out_ptr, out1_ptr, aux0_ptr, aux1_ptr, m, stream + ) + + @cute.jit + def call_gated( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + g_ptr: cute.Pointer, + g_ld: Int32, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ): + """Gated launch: ``out = bf16(x * g) @ W^T`` with ``g`` a BF16 ``[m, K]`` tensor of + row stride ``g_ld`` elements (a multiple of 8). The BF16 epilogue does not read the + fused-epilogue pointers, so ``out_ptr`` stands in for them.""" + self._launch( + w_ptr, x_ptr, g_ptr, g_ld, ws_ptr, out_ptr, out_ptr, out_ptr, out_ptr, m, stream + ) + + @cute.jit + def call_ar_push( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + m: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + stream: cuda.CUstream, + ): + """All-reduce push launch: the output ``x @ W^T`` goes into the Lamport slots of the + next call on the workspace (:class:`k3_lamport_ar.LamportWorkspace` addresses: peer + pointer table, multicast base, state flags; this is rank ``rank`` of ``ar_world``). + Nothing is stored through the output pointers, so ``x_ptr`` stands in for them.""" + self._launch( + w_ptr, + x_ptr, + None, + None, + ws_ptr, + x_ptr, + x_ptr, + x_ptr, + x_ptr, + m, + stream, + ar_peers=peer_ptrs, + ar_mc=mc_base, + ar_flags=flags, + ar_rank=rank, + ) + + @cute.jit + def call_ar_push_gated( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + g_ptr: cute.Pointer, + g_ld: Int32, + ws_ptr: cute.Pointer, + m: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + stream: cuda.CUstream, + ): + """:meth:`call_ar_push` of the gated kernel: pushes ``bf16(x * g) @ W^T``.""" + self._launch( + w_ptr, + x_ptr, + g_ptr, + g_ld, + ws_ptr, + x_ptr, + x_ptr, + x_ptr, + x_ptr, + m, + stream, + ar_peers=peer_ptrs, + ar_mc=mc_base, + ar_flags=flags, + ar_rank=rank, + ) + + @cute.jit + def call_split_ar_push( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + add_ptr: cute.Pointer, + m: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + stream: cuda.CUstream, + ): + """Strided-source launch (``split_src``) with the all-reduce push epilogue (see + :meth:`call_ar_push`); ``add_ptr``: the BF16 ``[M, N]`` addend of ``push_addend``.""" + self._launch( + w_ptr, + x_ptr, + None, + None, + ws_ptr, + x_ptr, + x_ptr, + add_ptr, + x_ptr, + m, + stream, + ar_peers=peer_ptrs, + ar_mc=mc_base, + ar_flags=flags, + ar_rank=rank, + ) + + @cute.jit + def _launch( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + g_ptr, + g_ld, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ar_peers=None, + ar_mc=None, + ar_flags=None, + ar_rank=None, + ): + args = self.kernel_args( + w_ptr, + x_ptr, + g_ptr, + g_ld, + ws_ptr, + out_ptr, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ) + self.kernel(*args).launch( + grid=(self.n_tiles * self.split_k, 1, 1), + block=(self.threads, 1, 1), + cluster=(self.split_k, 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.jit + def kernel_args( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + g_ptr, + g_ld, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + ar_peers=None, + ar_mc=None, + ar_flags=None, + ar_rank=None, + ) -> tuple: + """The arguments of :meth:`kernel` (TMA atoms, layouts, tensors) for one launch, traced + inside a ``cute.jit`` caller; also sets ``self.shared_storage``. A kernel that runs this + GEMM's CTAs next to other roles (``k3_ull_front``) passes them to :meth:`cta`.""" + n, k = self.n, self.k + if cutlass.const_expr(self.split_src is not None): + ldw, ldx = self.split_src + mW = cute.make_tensor(w_ptr, cute.make_layout((n, k, 1), stride=(ldw, 1, n * ldw))) + mX = cute.make_tensor(x_ptr, cute.make_layout((m, k), stride=(ldx, 1))) + else: + mW = cute.make_tensor(w_ptr, cute.make_layout((n, k, 1), stride=(k, 1, n * k))) + mX = cute.make_tensor(x_ptr, cute.make_layout((m, k), stride=(k, 1))) + # DeepGEMM packed UE8M0 weight scale: int32 [N, ceil(K/512)], MN-major; byte + # (kb % 4) of word (n, kb // 4) holds the exponent of block (n // 128, kb). + mWS = cute.make_tensor( + ws_ptr, + cute.make_layout( + (_ceil_div(n, MMA_M), (4, _ceil_div(k // CTA_K, 4))), + stride=(MMA_M * 4, (1, 4 * n)), + ), + ) + mOut = cute.make_tensor(out_ptr, cute.make_layout((m, n), stride=(n, 1))) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float8E8M0FNU, + SF_VEC, + tcgen05.CtaGroup.ONE, + (MMA_M, self.mma_n), + ) + a_layout = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.a_stages + ) + b_layout = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.b_stages + ) + sfa_layout = blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler, SF_VEC, 1) + sfb_layout = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, self.mma_tiler, SF_VEC, self.b_stages + ) + x_layout = cute.make_layout( + (self.mma_n, CTA_K, self.x_stages), stride=(CTA_K, 1, self.mma_n * CTA_K) + ) + + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) + ) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mW, + cute.slice_(a_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_x, tma_tensor_x = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mX, + cute.slice_(x_layout, (None, None, 0)), + (self.mma_n, CTA_K), + ) + tma_atom_g, tma_tensor_g = None, None + if cutlass.const_expr(self.gated): + mG = cute.make_tensor(g_ptr, cute.make_layout((m, k), stride=(g_ld, 1))) + tma_atom_g, tma_tensor_g = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mG, + cute.slice_(x_layout, (None, None, 0)), + (self.mma_n, CTA_K), + ) + g_ring_bytes = self.x_stages * self.x_stage_bytes if self.gated else 128 + + @cute.struct + class SharedStorage: + a_full: cute.struct.MemRange[cutlass.Int64, self.a_stages] + a_empty: cute.struct.MemRange[cutlass.Int64, self.a_stages] + x_full: cute.struct.MemRange[cutlass.Int64, self.x_stages] + x_empty: cute.struct.MemRange[cutlass.Int64, self.x_stages] + b_full: cute.struct.MemRange[cutlass.Int64, self.b_stages] + b_empty: cute.struct.MemRange[cutlass.Int64, self.b_stages] + acc_full: cutlass.Int64 + tmem_ready: cutlass.Int64 + tmem_free: cutlass.Int64 + mbox_full: cutlass.Int64 + peers_quantized: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + w_exp: cute.struct.MemRange[cutlass.Int32, self.kt_max] + # Weight ring; reused to stage outgoing split-K partials after the mainloop. + sA: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float8E4M3FN, max(cute.cosize(a_layout.outer), self.staging_bytes) + ], + 1024, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E4M3FN, cute.cosize(b_layout.outer)], 1024 + ] + # Activation ring; reused as the split-K reduction mailbox after quantization. + sXM: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, self.xm_bytes], 128] + sSFA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfa_layout)], 128 + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfb_layout)], 128 + ] + + @cute.struct + class GatedSharedStorage: + base: SharedStorage + # Gate ring (gated kernel): one BF16 gate tile per activation stage. + sG: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, g_ring_bytes], 128] + + self.shared_storage = GatedSharedStorage if self.gated else SharedStorage + + return ( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_x, + tma_tensor_x, + mWS, + mOut, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + x_layout, + tma_atom_g, + tma_tensor_g, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_x: cute.CopyAtom, + mXt: cute.Tensor, + mWS: cute.Tensor, + mOut: cute.Tensor, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + x_layout: cute.Layout, + tma_atom_g: Optional[cute.CopyAtom], + mGt: Optional[cute.Tensor], + ar_peers: Optional[Int64], + ar_mc: Optional[Int64], + ar_flags: Optional[Int64], + ar_rank: Optional[Int32], + ): + self.cta( + tiled_mma, + tma_atom_a, + mA, + tma_atom_x, + mXt, + mWS, + mOut, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + x_layout, + tma_atom_g, + mGt, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ) + + @cute.jit + def cta( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_x: cute.CopyAtom, + mXt: cute.Tensor, + mWS: cute.Tensor, + mOut: cute.Tensor, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + x_layout: cute.Layout, + tma_atom_g: Optional[cute.CopyAtom], + mGt: Optional[cute.Tensor], + ar_peers: Optional[Int64], + ar_mc: Optional[Int64], + ar_flags: Optional[Int64], + ar_rank: Optional[Int32], + storage=None, + ): + """One CTA of the GEMM grid (block ``n_tile * split_k + rank``), the body of + :meth:`kernel`. ``storage``: the ``shared_storage`` instance if the caller allocated it + (a kernel with other roles shares one allocation between them).""" + S = self.split_k + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + rank = Int32(0) + if cutlass.const_expr(S > 1): + rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_tile = bidx // S + + kt_base = self.kt_total // S + kt_rem = self.kt_total % S + kt_count = Int32(kt_base) + kt_start = rank * kt_base + if cutlass.const_expr(kt_rem > 0): + if rank < kt_rem: + kt_count = kt_count + 1 + kt_start = kt_start + rank + else: + kt_start = kt_start + kt_rem + + if cutlass.const_expr(storage is None): + storage = utils.SmemAllocator().allocate(self.shared_storage) + if cutlass.const_expr(self.gated): + g_smem = storage.sG.data_ptr() + storage = storage.base + else: + g_smem = storage.sXM.data_ptr() # unused + a_full = storage.a_full.data_ptr() + a_empty = storage.a_empty.data_ptr() + x_full = storage.x_full.data_ptr() + x_empty = storage.x_empty.data_ptr() + b_full = storage.b_full.data_ptr() + b_empty = storage.b_empty.data_ptr() + acc_full = storage.acc_full + tmem_ready = storage.tmem_ready + tmem_free = storage.tmem_free + mbox_full = storage.mbox_full + peers_quantized = storage.peers_quantized + tmem_holding_buf = storage.tmem_holding_buf + + if warp_idx == _TMA_WARP: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_x) + if cutlass.const_expr(self.gated): + cpasync.prefetch_descriptor(tma_atom_g) + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(self.a_stages): + cute.arch.mbarrier_init(a_full + i, 1) + cute.arch.mbarrier_init(a_empty + i, 1) + for i in cutlass.range_constexpr(self.x_stages): + cute.arch.mbarrier_init(x_full + i, 1) + cute.arch.mbarrier_init(x_empty + i, self.warps_per_group) + for i in cutlass.range_constexpr(self.b_stages): + cute.arch.mbarrier_init(b_full + i, self.warps_per_group) + cute.arch.mbarrier_init(b_empty + i, 1) + cute.arch.mbarrier_init(acc_full, 1) + cute.arch.mbarrier_init(tmem_ready, 1) + cute.arch.mbarrier_init(tmem_free, 4) + cute.arch.mbarrier_init(mbox_full, 1) + cute.arch.mbarrier_init(peers_quantized, max(S - 1, 1)) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + if cutlass.const_expr(S > 1): + cute.arch.cluster_arrive_relaxed() + + if warp_idx == _TMA_WARP: + # Weight producer: never waits on the previous grid. The all-reduce push variant + # releases its dependent after its own wait (below). + if cutlass.const_expr(self.use_pdl and not self.ar_push): + _griddepcontrol_launch_dependents() + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + gA = cute.local_tile( + mA, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + tCgA = tiled_mma.get_slice(0).partition_A(gA) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + for kt in cutlass.range(kt_count, unroll=1): + stage = kt % self.a_stages + if kt >= self.a_stages: + cute.arch.mbarrier_wait(a_empty + stage, ((kt // self.a_stages) & 1) ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(a_full + stage, self.a_stage_bytes) + cute.copy( + tma_atom_a, + tAgA[(None, n_tile, kt_start + kt, 0)], + tAsA[(None, stage)], + tma_bar_ptr=a_full + stage, + cache_policy=Int64(_EVICT_FIRST), + ) + elif warp_idx == _X_WARP: + # Activation producer: the only TMA that depends on the previous grid. + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + if cutlass.const_expr(self.ar_push != 0): + # The all-reduce finish kernel reads the Lamport rotation state as soon as + # it is resident (k3_lamport_ar): release it only once every earlier grid + # (the previous all-reduce on the workspace included) has completed. + _griddepcontrol_launch_dependents() + sX = cute.make_tensor( + cute.recast_ptr(storage.sXM.data_ptr(), dtype=cutlass.BFloat16), x_layout + ) + gX = cute.local_tile(mXt, (self.mma_n, CTA_K), (0, None)) + tXsX, tXgX = cpasync.tma_partition( + tma_atom_x, + 0, + cute.make_layout(1), + cute.group_modes(sX, 0, 2), + cute.group_modes(gX, 0, 2), + ) + tGsG, tGgG = None, None + if cutlass.const_expr(self.gated): + # Each gate tile is loaded with the activation tile it multiplies, on the + # same barrier, so it is also read only after the wait above. + sG = cute.make_tensor(cute.recast_ptr(g_smem, dtype=cutlass.BFloat16), x_layout) + gG = cute.local_tile(mGt, (self.mma_n, CTA_K), (0, None)) + tGsG, tGgG = cpasync.tma_partition( + tma_atom_g, + 0, + cute.make_layout(1), + cute.group_modes(sG, 0, 2), + cute.group_modes(gG, 0, 2), + ) + stage_tx = self.x_stage_bytes * 2 if self.gated else self.x_stage_bytes + for kt in cutlass.range(kt_count, unroll=1): + stage = kt % self.x_stages + if kt >= self.x_stages: + cute.arch.mbarrier_wait(x_empty + stage, ((kt // self.x_stages) & 1) ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(x_full + stage, stage_tx) + cute.copy( + tma_atom_x, + tXgX[(None, kt_start + kt)], + tXsX[(None, stage)], + tma_bar_ptr=x_full + stage, + ) + if cutlass.const_expr(self.gated): + cute.copy( + tma_atom_g, + tGgG[(None, kt_start + kt)], + tGsG[(None, stage)], + tma_bar_ptr=x_full + stage, + ) + elif warp_idx == _MMA_WARP: + self.mma_warp( + tiled_mma, + storage, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + a_full, + a_empty, + b_full, + b_empty, + acc_full, + tmem_ready, + tmem_free, + tmem_holding_buf, + kt_count, + ) + elif warp_idx >= _QUANT_WARP0: + q = tidx - _QUANT_WARP0 * 32 + self.quantize_activations( + storage, + g_smem, + mWS, + q, + n_tile, + kt_start, + kt_count, + x_full, + x_empty, + b_full, + b_empty, + ) + if cutlass.const_expr(S > 1): + # The activation ring becomes the split-K mailbox: tell every peer once + # all quantization warps are done reading it. + cute.arch.barrier( + barrier_id=_QUANT_DONE_BARRIER_ID, number_of_threads=self.quant_threads + ) + if warp_idx < _QUANT_WARP0 + 4: + # Peers' mbarriers are initialized once the start-up cluster barrier + # completes (``barrier.cluster.wait`` is warp-aligned). + cute.arch.cluster_wait() + if q == 0: + for p in cutlass.range_constexpr(1, S): + _mbarrier_arrive_remote(peers_quantized.toint(), (rank + p) % S) + if warp_idx < _QUANT_WARP0 + 4: + self.epilogue( + tiled_mma, + storage, + mOut, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + q, + n_tile, + rank, + acc_full, + tmem_ready, + tmem_free, + mbox_full, + peers_quantized, + tmem_holding_buf, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ) + + @cute.jit + def mma_warp( + self, + tiled_mma: cute.TiledMma, + storage, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + a_full, + a_empty, + b_full, + b_empty, + acc_full, + tmem_ready, + tmem_free, + tmem_holding_buf, + kt_count: Int32, + ): + cute.arch.alloc_tmem(_TMEM_COLS, tmem_holding_buf) + cute.arch.relinquish_tmem_alloc_permit() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(tmem_ready) + + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tiled_mma.make_fragment_C(acc_shape).layout) + sfa_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + self.mma_n, dtype=cutlass.Float8E8M0FNU) + tCtSFA = cute.make_tensor( + sfa_tmem_ptr, + blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, self.mma_tiler, SF_VEC, cute.slice_(sfa_layout, (None, None, None, 0)) + ), + ) + sfb_tmem_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, self.mma_tiler, SF_VEC, cute.slice_(sfb_layout, (None, None, None, 0)) + ) + sSFA = storage.sSFA.get_tensor(sfa_layout) + sSFB = storage.sSFB.get_tensor(sfb_layout) + s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = self._s2t_partition(sSFA, tCtSFA) + # SFB: TMEM columns [mma_n + 16, +16). + tCtSFB = cute.make_tensor( + cute.recast_ptr( + acc_tmem_ptr + self.mma_n + _SF_TMEM_COLS, + dtype=cutlass.Float8E8M0FNU, + ), + sfb_tmem_layout, + ) + s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = self._s2t_partition(sSFB, tCtSFB) + + # The weight scale is folded into SFB, so SFA is a constant 1.0 (UE8M0 127). + lane = cute.arch.lane_idx() + one4 = Int32(0x7F7F7F7F) + _sts_v4_b32(storage.sSFA.data_ptr().toint() + lane * 16, one4, one4, one4, one4) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + cute.copy(s2t_sfa, tCsSFA_s2t[(None, None, None, None, 0)], tCtSFA_s2t) + + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + num_kphases = cute.size(tCrA, mode=[2]) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for kt in cutlass.range(kt_count, unroll=1): + a_stage = kt % self.a_stages + b_stage = kt % self.b_stages + cute.arch.mbarrier_wait(a_full + a_stage, (kt // self.a_stages) & 1) + cute.arch.mbarrier_wait(b_full + b_stage, (kt // self.b_stages) & 1) + cute.copy(s2t_sfb, tCsSFB_s2t[(None, None, None, None, b_stage)], tCtSFB_s2t) + for kp in cutlass.range_constexpr(num_kphases): + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator) + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[(None, None, kp, a_stage)], + tCrB[(None, None, kp, b_stage)], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + with cute.arch.elect_one(): + tcgen05.commit(a_empty + a_stage) + tcgen05.commit(b_empty + b_stage) + with cute.arch.elect_one(): + tcgen05.commit(acc_full) + + cute.arch.mbarrier_wait(tmem_free, 0) + cute.arch.dealloc_tmem(acc_tmem_ptr, _TMEM_COLS) + + def _s2t_partition(self, sSF: cute.Tensor, tSF: cute.Tensor): + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), cutlass.Float8E8M0FNU + ) + tiled_copy = tcgen05.make_s2t_copy(copy_atom, tCtSF_compact) + thr_copy = tiled_copy.get_slice(0) + tCsSF_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy, thr_copy.partition_S(tCsSF_compact) + ) + tCtSF_s2t = thr_copy.partition_D(tCtSF_compact) + return tiled_copy, tCsSF_s2t, tCtSF_s2t + + @cute.jit + def quantize_activations( + self, + storage, + g_smem: cute.Pointer, + mWS: cute.Tensor, + q: Int32, + n_tile: Int32, + kt_start: Int32, + kt_count: Int32, + x_full, + x_empty, + b_full, + b_empty, + ): + """Quantize BF16 activation tiles into the swizzled MXFP8 UMMA B layout. + + K tiles are dealt round-robin to ``quant_groups`` groups of warps, so several tiles + are in flight at once; ring slot ``kt % stages`` is then always served by the same + group. Within a tile, one work item is one (row, 32-element K group): four 16-byte + loads, one UE8M0 scale, 32 E4M3 bytes, split evenly over the group's lanes. The gated + kernel first replaces the 32 values by their BF16 products with the gate. + """ + nthr = self.quant_threads + w_exp = storage.w_exp.get_tensor(cute.make_layout(self.kt_max)) + if q < kt_count: + w_exp[q] = Int32(mWS[n_tile, kt_start + q]) + cute.arch.barrier(barrier_id=_QUANT_BARRIER_ID, number_of_threads=nthr) + + x_base = storage.sXM.data_ptr().toint() + g_base = g_smem.toint() + b_base = storage.sB.data_ptr().toint() + sf_base = storage.sSFB.data_ptr().toint() + group_threads = 32 * self.warps_per_group + gq = q % group_threads + for kt in cutlass.range(q // group_threads, kt_count, self.quant_groups, unroll=1): + x_stage = kt % self.x_stages + b_stage = kt % self.b_stages + cute.arch.mbarrier_wait(x_full + x_stage, (kt // self.x_stages) & 1) + if kt >= self.b_stages: + cute.arch.mbarrier_wait(b_empty + b_stage, ((kt // self.b_stages) & 1) ^ 1) + e_w = w_exp[kt] + xs = x_base + x_stage * self.x_stage_bytes + bs = b_base + b_stage * self.b_stage_bytes + ss = sf_base + b_stage * self.sf_stage_bytes + for j in cutlass.range_constexpr(self.mma_n * 4 // group_threads): + item = gq + group_threads * j + r = item >> 2 + g = item & 3 + # Rotate the 16-byte chunk order so a quarter-warp hits 8 distinct banks. + rot = item >> 1 + xrow = xs + r * (CTA_K * 2) + g * 64 + ws = [] + for v in cutlass.range_constexpr(4): + ws.append(_lds_v4_b32(xrow + ((v + rot) & 3) * 16)) + if cutlass.const_expr(self.gated): + grow = g_base + x_stage * self.x_stage_bytes + r * (CTA_K * 2) + g * 64 + for v in cutlass.range_constexpr(4): + ws[v] = _gate_mul_bf16x8(*ws[v], *_lds_v4_b32(grow + ((v + rot) & 3) * 16)) + amax = _abs_max_bf16x32([w for quad in ws for w in quad]) + e_a = _ue8m0_of_amax(amax) + inv = _bf16x2_inv_scale(e_a) + brow = bs + r * CTA_K + swz = r & 7 + for v in cutlass.range_constexpr(4): + sub = (v + rot) & 3 + p0, p1 = _quant_bf16x8_to_e4m3(*ws[v], inv) + unit = g * 2 + (sub >> 1) + _sts_v2_b32(brow + ((unit ^ swz) << 4) + (sub & 1) * 8, p0, p1) + _sts_scale(ss + (r & 31) * 16 + (r >> 5) * 4 + g, e_a, e_w) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(x_empty + x_stage) + cute.arch.mbarrier_arrive(b_full + b_stage) + + def _out_value(self, acc): + """The stored BF16 value of one FP32 accumulator.""" + out = acc.to(cutlass.BFloat16) + if self.sigmoid: + out = sigmoid_bf16(out) + return out + + @cute.jit + def epilogue( + self, + tiled_mma: cute.TiledMma, + storage, + mOut: cute.Tensor, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + q: Int32, + n_tile: Int32, + rank: Int32, + acc_full, + tmem_ready, + tmem_free, + mbox_full, + peers_quantized, + tmem_holding_buf, + ar_peers: Optional[Int64] = None, + ar_mc: Optional[Int64] = None, + ar_flags: Optional[Int64] = None, + ar_rank: Optional[Int32] = None, + ): + S = self.split_k + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + cute.arch.mbarrier_wait(tmem_ready, 0) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tiled_mma.make_fragment_C(acc_shape).layout) + tAcc = tCtAcc[((None, None), 0, 0)] + copy_atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(self.mma_n)), cutlass.Float32 + ) + tiled_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc) + thr_t2r = tiled_t2r.get_slice(q) + tTR_tAcc = thr_t2r.partition_S(tAcc) + tTR_cAcc = thr_t2r.partition_D(cute.make_identity_tensor((MMA_M, self.mma_n))) + rAcc = cute.make_rmem_tensor(tTR_cAcc.shape, cutlass.Float32) + + cute.arch.mbarrier_wait(acc_full, 0) + cute.copy(tiled_t2r, tTR_tAcc, rAcc) + cute.arch.fence_view_async_tmem_load() + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(tmem_free) + + n_base = n_tile * MMA_M + nchunks = cute.size(rAcc) // 4 + # Every thread holds one accumulator row. Rows past N (the zero-filled tail of a + # partial last weight tile) are reduced like the others but never stored. + m_store = m + if cutlass.const_expr(self.n % MMA_M != 0): + row0, _ = tTR_cAcc[0] + if n_base + row0 >= self.n: + m_store = Int32(0) + # Fused epilogues: BF16 token rows staged in the dead quantized-activation ring. + stage_out = storage.sB.data_ptr().toint() + if cutlass.const_expr(self.ar_push != 0 and S == 1): + self._push_lamport( + storage, + rAcc, + tTR_cAcc, + m, + q, + n_base, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + aux0_ptr, + ) + elif cutlass.const_expr(S == 1): + if cutlass.const_expr(self.epilogue_kind == EPI_BF16): + for i in cutlass.range_constexpr(nchunks * 4): + row, col = tTR_cAcc[i] + if col < m_store: + mOut[col, n_base + row] = self._out_value(rAcc[i]) + else: + for i in cutlass.range_constexpr(nchunks * 4): + row, col = tTR_cAcc[i] + if col < m: + _sts_u16( + stage_out + col * (MMA_M * 2) + row * 2, _f32_to_bf16_bits(rAcc[i]) + ) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + self.fused_store( + stage_out, m, rank, m, n_tile, q, mOut.iterator, out1_ptr, aux0_ptr, aux1_ptr + ) + else: + # Reduce-scatter over 4-column chunks: chunk c is owned by rank c % S. Every + # thread holds one accumulator row. Partials for each owner are staged in the + # (now dead) weight ring as [chunk][row] and pushed with one bulk DSMEM copy per + # owner into its mailbox (its dead activation ring), laid out [peer][chunk][row]. + mbox_base = storage.sXM.data_ptr().toint() + stage_base = storage.sA.data_ptr().toint() + mbar_addr = mbox_full.toint() + slot_bytes = self.chunks_cap * MMA_M * 16 + my_chunks = Int32(0) + for c in cutlass.range_constexpr(nchunks): + _, col0 = tTR_cAcc[4 * c] + if col0 < m: + if rank == c % S: + my_chunks = my_chunks + 1 + else: + _sts_v4_f32( + stage_base + (c % S) * slot_bytes + ((c // S) * MMA_M + q) * 16, + rAcc[4 * c], + rAcc[4 * c + 1], + rAcc[4 * c + 2], + rAcc[4 * c + 3], + ) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + if q == 0: + cute.arch.mbarrier_wait(peers_quantized, 0) + for o in cutlass.range_constexpr(S): + if rank != o: + sent = Int32(0) + for c in cutlass.range_constexpr(o, nchunks, S): + _, col0 = tTR_cAcc[4 * c] + if col0 < m: + sent = sent + 1 + if sent > 0: + slot = rank + if rank > o: + slot = rank - 1 + _bulk_copy_to_peer( + mbox_base + slot * slot_bytes, + stage_base + o * slot_bytes, + sent * (MMA_M * 16), + mbar_addr, + o, + ) + _bulk_commit() + if my_chunks > 0: + if q == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + mbox_full, my_chunks * ((S - 1) * MMA_M * 16) + ) + cute.arch.mbarrier_wait(mbox_full, 0) + for c in cutlass.range_constexpr(nchunks): + _, col0 = tTR_cAcc[4 * c] + if col0 < m: + if rank == c % S: + for p in cutlass.range_constexpr(S - 1): + v0, v1, v2, v3 = _lds_v4_f32( + mbox_base + p * slot_bytes + ((c // S) * MMA_M + q) * 16 + ) + rAcc[4 * c] = rAcc[4 * c] + v0 + rAcc[4 * c + 1] = rAcc[4 * c + 1] + v1 + rAcc[4 * c + 2] = rAcc[4 * c + 2] + v2 + rAcc[4 * c + 3] = rAcc[4 * c + 3] + v3 + for j in cutlass.range_constexpr(4): + row, col = tTR_cAcc[4 * c + j] + if cutlass.const_expr( + self.epilogue_kind == EPI_BF16 and self.ar_push == 0 + ): + if col < m_store: + mOut[col, n_base + row] = self._out_value(rAcc[4 * c + j]) + else: + if col < m: + _sts_u16( + stage_out + col * (MMA_M * 2) + row * 2, + _f32_to_bf16_bits(rAcc[4 * c + j]), + ) + if cutlass.const_expr(self.ar_push != 0): + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + self._push_owned_rows( + stage_out, + my_chunks * 4, + rank, + m, + n_base, + q, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + aux0_ptr, + ) + elif cutlass.const_expr(self.epilogue_kind != EPI_BF16): + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + self.fused_store( + stage_out, + my_chunks * 4, + rank, + m, + n_tile, + q, + mOut.iterator, + out1_ptr, + aux0_ptr, + aux1_ptr, + ) + if q == 0: + # The staging buffer must stay valid until the copies have read it. + _bulk_wait_read() + + @cute.jit + def _add_addend(self, w, addend: cute.Pointer, t, col): + """``bf16(w + a)`` per BF16 lane of the 16-byte vector ``w`` (features ``[col, col + 8)`` + of token ``t``), ``a`` from the BF16 ``[M, N]`` addend (the PyTorch BF16 add).""" + a = _ldg_v4_b32(addend.toint() + (Int64(t) * self.n + Int64(col)) * 2) + return tuple(_add_bf16x2(w[i], a[i]) for i in range(4)) + + @cute.jit + def _push_lamport( + self, + storage, + rAcc: cute.Tensor, + tTR_cAcc: cute.Tensor, + m: Int32, + q: Int32, + n_base: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + ar_rank: Int32, + addend: cute.Pointer, + ): + """All-reduce push of this CTA's output tile (``split_k == 1``): round every value to + BF16 (as the BF16 store would), stage the tile token-major in the dead weight ring, then + store 16-byte vectors (8 features of one token) into the Lamport slots of the next call + on the workspace (``k3_lamport_ar.push_slot``; ``-0`` is sent as ``+0``), with + ``push_addend`` after the BF16 add of ``addend``. The epilogue has passed + ``griddepcontrol.wait``: the previous call on the workspace has completed, so its + rotation state is final, and the addend is complete.""" + stage_base = storage.sA.data_ptr().toint() + tile_row = MMA_M * 2 + for i in cutlass.range_constexpr(cute.size(rAcc)): + row, col = tTR_cAcc[i] + _sts_u16(stage_base + col * tile_row + row * 2, _f32_to_bf16_bits(rAcc[i])) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + cur, _, nbytes, _, _, _, _, _ = k3_lamport_ar.lamport_state(flags) + one_shot = self.ar_push == AR_PUSH_ONE_SHOT + vecs = MMA_M // 8 + for p in cutlass.range_constexpr(self.mma_n * vecs // MMA_M): + j = q + p * MMA_M + t = j // vecs + c = j - t * vecs + if t < m and n_base + c * 8 < self.n: + w = _lds_v4_b32(stage_base + t * tile_row + c * 16) + if cutlass.const_expr(self.push_addend): + w = self._add_addend(w, addend, t, n_base + c * 8) + dst = k3_lamport_ar.push_slot( + peer_ptrs, + mc_base, + cur, + nbytes, + t, + ar_rank, + n_base + c * 8, + self.ar_world, + self.ar_slot_n, + one_shot, + ) + k3_lamport_ar.push_v4(dst, *w, one_shot) + + @cute.jit + def _push_owned_rows( + self, + stage_base, + nrows: Int32, + rank: Int32, + m: Int32, + n_base: Int32, + q: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + ar_rank: Int32, + addend: cute.Pointer, + ): + """All-reduce push after the split-K reduction: the ``nrows`` token rows this cluster + rank owns are staged as BF16 at ``stage_base + t * 256`` (as for the fused epilogues; + the ``r``-th owned row is token ``_owned_token(r, rank)``) and go out in 16-byte + vectors like :meth:`_push_lamport`'s, so every (token, feature) vector is pushed by + exactly one CTA of the cluster.""" + cur, _, nbytes, _, _, _, _, _ = k3_lamport_ar.lamport_state(flags) + one_shot = self.ar_push == AR_PUSH_ONE_SHOT + # One item = 8 features (16 bytes) of one owned row; MMA_M // 8 = 16 items per row. + for p in cutlass.range(q, nrows * 16, MMA_M, unroll=1): + t = self._owned_token(p >> 4, rank) + c = p & 15 + if t < m and n_base + c * 8 < self.n: + w = _lds_v4_b32(stage_base + t * (MMA_M * 2) + c * 16) + if cutlass.const_expr(self.push_addend): + w = self._add_addend(w, addend, t, n_base + c * 8) + dst = k3_lamport_ar.push_slot( + peer_ptrs, + mc_base, + cur, + nbytes, + t, + ar_rank, + n_base + c * 8, + self.ar_world, + self.ar_slot_n, + one_shot, + ) + k3_lamport_ar.push_v4(dst, *w, one_shot) + + +# --------------------------------------------------------------------------- +# Mid-M path (M up to 256): standalone MXFP8 activation quantization + GEMM with a +# TMA-loaded FP8 activation operand. +# --------------------------------------------------------------------------- + +#: Token tiles (UMMA N) of the mid-M GEMM: M rounds up to a multiple of 32 (the epilogue +#: drains TMEM in 32-column blocks); MMA time scales with the padded N. +SUPPORTED_MID_MMA_N = (64, 96, 128, 160, 192, 224, 256) +#: Leave room for the quantization kernel's CTAs (1 KB reserved SMEM each) next to a GEMM +#: CTA, so the GEMM can launch and stream weights while the quantization still runs. +_MID_SMEM_CAPACITY = _SMEM_CAPACITY - 2048 + + +@dataclasses.dataclass(frozen=True) +class K3MidMTactic: + """One mid-M specialization: ``mma_n`` padded tokens, cluster ``split_k``, ring depths.""" + + mma_n: int + split_k: int + a_stages: int + b_stages: int = 3 + + +def _sf_blocks(mma_n: int) -> int: + """128-row scale-factor atoms covering ``mma_n`` tokens.""" + return _ceil_div(mma_n, MMA_M) + + +def _mid_mailbox_bytes(tactic: K3MidMTactic) -> int: + if tactic.split_k == 1: + return 0 + return (tactic.split_k - 1) * _ceil_div(tactic.mma_n // 4, tactic.split_k) * MMA_M * 16 + + +def mid_smem_bytes(n: int, k: int, tactic: K3MidMTactic) -> int: + """Conservative host-side mirror of the mid-M kernel's shared-memory footprint.""" + kt_max = _ceil_div(k // CTA_K, tactic.split_k) + a_st = min(tactic.a_stages, kt_max) + b_st = min(tactic.b_stages, kt_max) + total = 8 * (2 * (a_st + b_st) + 5) + 4 + # The split-K mailbox reuses the weight ring once the mainloop is done. + total = _align_up(total, 1024) + max(a_st * MMA_M * CTA_K, _mid_mailbox_bytes(tactic)) + total = _align_up(total, 1024) + b_st * tactic.mma_n * CTA_K + total = _align_up(total, 128) + kt_max * 512 + total = _align_up(total, 128) + b_st * _sf_blocks(tactic.mma_n) * 512 + return total + 1024 # allocator slack + + +def validate_mid_tactic(n: int, k: int, tactic: K3MidMTactic) -> None: + if tactic.mma_n not in SUPPORTED_MID_MMA_N: + raise ValueError(f"mma_n must be one of {SUPPORTED_MID_MMA_N}, got {tactic.mma_n}") + if n % N_ALIGN or k % CTA_K: + raise ValueError(f"N={n} must be a multiple of {N_ALIGN} and K={k} of {CTA_K}") + if not 1 <= tactic.split_k <= min(16, k // CTA_K): + raise ValueError(f"split_k={tactic.split_k} invalid for K={k}") + if tactic.a_stages < 1 or tactic.b_stages < 1: + raise ValueError("a_stages and b_stages must be >= 1") + need = mid_smem_bytes(n, k, tactic) + if need > _MID_SMEM_CAPACITY: + raise ValueError(f"tactic {tactic} needs {need} B of SMEM > {_SMEM_CAPACITY} B") + + +def max_mid_a_stages(n: int, k: int, tactic: K3MidMTactic) -> int: + kt_max = _ceil_div(k // CTA_K, tactic.split_k) + best = 1 + for a in range(1, kt_max + 1): + if mid_smem_bytes(n, k, dataclasses.replace(tactic, a_stages=a)) <= _MID_SMEM_CAPACITY: + best = a + return best + + +@dsl_user_op +def _ldg_v4_b32(addr, *, loc=None, ip=None): + """Load 16 bytes from global memory as four 32-bit words.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)) + + +@dsl_user_op +def _stg_v4_b32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u8(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u8 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_load_g2s(dst_smem, src_gmem, nbytes, mbar, *, loc=None, ip=None): + """Bulk-copy ``nbytes`` from global memory into local SMEM, crediting ``mbar``.""" + _asm( + None, + [ + _iv(dst_smem, loc=loc, ip=ip), + Int64(src_gmem).ir_value(loc=loc, ip=ip), + _iv(nbytes, loc=loc, ip=ip), + _iv(mbar, loc=loc, ip=ip), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];", + "r,l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_async_v4_f32(local_addr, v0, v1, v2, v3, local_mbar, peer, *, loc=None, ip=None): + """Store four FP32 values into a peer CTA's SMEM, crediting its mbarrier.""" + _asm( + None, + [ + _iv(local_addr, loc=loc, ip=ip), + *(cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)), + _iv(local_mbar, loc=loc, ip=ip), + _iv(peer, loc=loc, ip=ip), + ], + "{\n" + ".reg .b32 ra, rb;\n" + "mapa.shared::cluster.u32 ra, $0, $6;\n" + "mapa.shared::cluster.u32 rb, $5, $6;\n" + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [ra], {$1, $2, $3, $4}, [rb];\n" + "}", + "r,f,f,f,f,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_u16(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.u16 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _f32_to_bf16_bits(value, *, loc=None, ip=None): + """Round-to-nearest-even FP32 -> BF16, returned in the low 16 bits of an Int32.""" + return Int32( + _asm( + T.i32(), + [cutlass.Float32(value).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\ncvt.u32.u16 $0, h;\n}", + "=r,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _min_s32(a, b, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "min.s32 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +class K3MxFp8QuantKernel: + """``x`` BF16 ``[M, K]`` -> E4M3 ``[M, K]`` + UE8M0 1x32 scales for the mid-M GEMM. + + One work item is one (row, 32-element group), quantized exactly like the low-M + kernel's in-kernel path; each thread issues the loads of its ``_ITEMS`` items before + converting any. CTAs are small so they fit next to a resident GEMM CTA (which + launches early via PDL and streams its weights meanwhile). + Scales are written in the tcgen05 scale-factor atom layout ``[K/128][sf_blocks][512 B]`` + with atom byte ``(r % 32) * 16 + (r // 32) * 4 + k32 % 4``, so the GEMM moves one + contiguous ``sf_blocks * 512`` byte block per K tile. Rows past ``M`` (up to + ``128 * sf_blocks``) get the neutral scale 1.0. + """ + + _THREADS = 128 + _ITEMS = 2 + + def __init__(self, k: int, sf_blocks: int, use_pdl: bool): + self.k = k + self.sf_blocks = sf_blocks + self.groups = k // SF_VEC + self.items = MMA_M * sf_blocks * self.groups + self.items_per_thread = self._ITEMS + self.ctas = _ceil_div(self.items, self._THREADS * self._ITEMS) + self.use_pdl = use_pdl + + @cute.jit + def __call__( + self, + x_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ): + self.kernel(x_ptr, xq_ptr, xs_ptr, m).launch( + grid=(self.ctas, 1, 1), + block=(self._THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel(self, x_ptr: cute.Pointer, xq_ptr: cute.Pointer, xs_ptr: cute.Pointer, m: Int32): + if cutlass.const_expr(self.use_pdl): + # The GEMM only reads this grid's results after its own griddepcontrol.wait; + # let it launch right away so its weight stream overlaps this kernel. + _griddepcontrol_launch_dependents() + _griddepcontrol_wait() + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + ipt = self.items_per_thread + rX = cute.make_rmem_tensor((ipt * 16,), Int32) + # Issue every load first; out-of-range items re-read a valid row and are not stored. + for it in cutlass.range_constexpr(ipt): + item = _min_s32( + bidx * self._THREADS + tidx + it * (self.ctas * self._THREADS), self.items - 1 + ) + r = _min_s32(item // self.groups, m - 1) + g = item - (item // self.groups) * self.groups + src = x_ptr.toint() + Int64(r * self.k + g * SF_VEC) * 2 + for v in cutlass.range_constexpr(4): + w0, w1, w2, w3 = _ldg_v4_b32(src + v * 16) + rX[it * 16 + v * 4] = w0 + rX[it * 16 + v * 4 + 1] = w1 + rX[it * 16 + v * 4 + 2] = w2 + rX[it * 16 + v * 4 + 3] = w3 + for it in cutlass.range_constexpr(ipt): + item = bidx * self._THREADS + tidx + it * (self.ctas * self._THREADS) + if item < self.items: + r = item // self.groups + g = item - r * self.groups + rr = r & (MMA_M - 1) + sf_off = ((g >> 2) * self.sf_blocks + (r >> 7)) * 512 + sf_off = sf_off + (rr & 31) * 16 + (rr >> 5) * 4 + (g & 3) + sf_addr = xs_ptr.toint() + Int64(sf_off) + if r < m: + ws = [rX[it * 16 + i] for i in range(16)] + e_a = _ue8m0_of_amax(_abs_max_bf16x32(ws)) + inv = _bf16x2_inv_scale(e_a) + packed = [] + for v in cutlass.range_constexpr(4): + p0, p1 = _quant_bf16x8_to_e4m3(*ws[4 * v : 4 * v + 4], inv) + packed.extend((p0, p1)) + dst = xq_ptr.toint() + Int64(r * self.k + g * SF_VEC) + _stg_v4_b32(dst, *packed[:4]) + _stg_v4_b32(dst + 16, *packed[4:]) + _stg_u8(sf_addr, e_a) + else: + _stg_u8(sf_addr, Int32(127)) + + +def _s2t_partition_sf(sSF: cute.Tensor, tSF: cute.Tensor): + """S2T (SMEM -> TMEM) copy of a scale-factor operand, partitioned for one thread.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), cutlass.Float8E8M0FNU + ) + tiled_copy = tcgen05.make_s2t_copy(copy_atom, tCtSF_compact) + thr_copy = tiled_copy.get_slice(0) + tCsSF_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy, thr_copy.partition_S(tCsSF_compact)) + tCtSF_s2t = thr_copy.partition_D(tCtSF_compact) + return tiled_copy, tCsSF_s2t, tCtSF_s2t + + +class K3Fp8wMidMGemmKernel(_FusedStoreMixin): + """Swap-AB MXFP8 GEMM fed by pre-quantized FP8 activations (64 < M <= 256). + + Same weight path and split-K cluster as :class:`K3Fp8wLowMGemmKernel`, but the + activation operand is the output of :class:`K3MxFp8QuantKernel`: warp 2 loads the + E4M3 tile with TMA (128B swizzle, i.e. directly in the UMMA B layout) and its scale + block with a bulk copy. The weight block scale goes into SFA (one UE8M0 per K tile, + replicated over the 128 rows), written once to SMEM before the mainloop. Warps 4-7 run + the epilogue in 32-column TMEM blocks; with ``split_k > 1`` each rank owns every + ``split_k``-th 4-token chunk and peers push partials with ``st.async`` into its + weight ring once every CTA of the cluster has finished its mainloop. ``epilogue`` selects + the output as in :class:`K3Fp8wLowMGemmKernel`. + + ``ar_push`` (``AR_PUSH_TWO_SHOT`` / ``AR_PUSH_ONE_SHOT``, ``ar_world`` ranks; c240, + :meth:`call_ar_push`): the staged BF16 rows go into the Lamport slots of the next all-reduce + on a workspace instead of the output (:meth:`_push_staged_rows`, the low-M kernel's push), + and the dependent grid is released from the activation warp after its + ``griddepcontrol.wait``. + """ + + _THREADS = 256 + _B_WARP = 2 + _EPI_WARP0 = 4 + + def __init__( + self, + n: int, + k: int, + tactic: K3MidMTactic, + use_pdl: bool, + epilogue: str = EPI_BF16, + ar_push: int = 0, + ar_world: int = 0, + ): + validate_mid_tactic(n, k, tactic) + _validate_epilogue(n, epilogue) + if ar_push: + if ar_push not in (AR_PUSH_TWO_SHOT, AR_PUSH_ONE_SHOT) or ar_world < 2: + raise ValueError(f"invalid all-reduce push {ar_push} over {ar_world} ranks") + if epilogue != EPI_BF16: + raise ValueError("the all-reduce push needs the BF16 epilogue") + self.ar_push = ar_push + self.ar_world = ar_world + self.epilogue_kind = epilogue + self.n = n + self.k = k + self.mma_n = tactic.mma_n + self.split_k = tactic.split_k + self.kt_total = k // CTA_K + self.kt_max = _ceil_div(self.kt_total, self.split_k) + self.a_stages = min(tactic.a_stages, self.kt_max) + self.b_stages = min(tactic.b_stages, self.kt_max) + self.n_tiles = _ceil_div(n, MMA_M) + self.use_pdl = use_pdl + self.mma_tiler = (MMA_M, self.mma_n, CTA_K) + self.a_stage_bytes = MMA_M * CTA_K + self.b_stage_bytes = self.mma_n * CTA_K + self.sf_blocks = _sf_blocks(self.mma_n) + self.sfb_stage_bytes = self.sf_blocks * 512 + self.chunks = self.mma_n // 4 + self.chunks_cap = _ceil_div(self.chunks, self.split_k) + self.a_region_bytes = max(self.a_stages * self.a_stage_bytes, _mid_mailbox_bytes(tactic)) + self.out_stage_bytes = self.mma_n * MMA_M * 2 + # Accumulator columns plus SFA / SFB columns, rounded to a power of two. + cols = self.mma_n + 2 * _SF_TMEM_COLS + self.tmem_cols = 1 << (cols - 1).bit_length() + + @cute.jit + def __call__( + self, + w_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ): + self._launch( + w_ptr, xq_ptr, xs_ptr, ws_ptr, out_ptr, out1_ptr, aux0_ptr, aux1_ptr, m, stream + ) + + @cute.jit + def call_ar_push( + self, + w_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + m: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + stream: cuda.CUstream, + ): + """All-reduce push launch (see :meth:`K3Fp8wLowMGemmKernel.call_ar_push`). Nothing is + stored through the output pointers, so ``xs_ptr`` stands in for them.""" + self._launch( + w_ptr, + xq_ptr, + xs_ptr, + ws_ptr, + xs_ptr, + xs_ptr, + xs_ptr, + xs_ptr, + m, + stream, + ar_peers=peer_ptrs, + ar_mc=mc_base, + ar_flags=flags, + ar_rank=rank, + ) + + @cute.jit + def _launch( + self, + w_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ar_peers=None, + ar_mc=None, + ar_flags=None, + ar_rank=None, + ): + n, k = self.n, self.k + mW = cute.make_tensor(w_ptr, cute.make_layout((n, k, 1), stride=(k, 1, n * k))) + mXq = cute.make_tensor(xq_ptr, cute.make_layout((m, k, 1), stride=(k, 1, m * k))) + mWS = cute.make_tensor( + ws_ptr, + cute.make_layout( + (_ceil_div(n, MMA_M), (4, _ceil_div(k // CTA_K, 4))), + stride=(MMA_M * 4, (1, 4 * n)), + ), + ) + mOut = cute.make_tensor(out_ptr, cute.make_layout((m, n), stride=(n, 1))) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float8E8M0FNU, + SF_VEC, + tcgen05.CtaGroup.ONE, + (MMA_M, self.mma_n), + ) + a_layout = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.a_stages + ) + b_layout = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.b_stages + ) + sfa_layout = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, self.mma_tiler, SF_VEC, self.kt_max + ) + sfb_layout = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, self.mma_tiler, SF_VEC, self.b_stages + ) + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) + ) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mW, + cute.slice_(a_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mXq, + cute.slice_(b_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + + @cute.struct + class SharedStorage: + a_full: cute.struct.MemRange[cutlass.Int64, self.a_stages] + a_empty: cute.struct.MemRange[cutlass.Int64, self.a_stages] + b_full: cute.struct.MemRange[cutlass.Int64, self.b_stages] + b_empty: cute.struct.MemRange[cutlass.Int64, self.b_stages] + acc_full: cutlass.Int64 + tmem_ready: cutlass.Int64 + tmem_free: cutlass.Int64 + mbox_full: cutlass.Int64 + peers_done: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # Weight ring; reused as the split-K mailbox after the mainloop. + sA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E4M3FN, self.a_region_bytes], 1024 + ] + # Activation ring; reused to stage the BF16 output tile token-major. + sB: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float8E4M3FN, max(cute.cosize(b_layout.outer), self.out_stage_bytes) + ], + 1024, + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfa_layout)], 128 + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfb_layout)], 128 + ] + + self.shared_storage = SharedStorage + + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + xs_ptr, + mWS, + mOut, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ).launch( + grid=(self.n_tiles * self.split_k, 1, 1), + block=(self._THREADS, 1, 1), + cluster=(self.split_k, 1, 1), + smem=self.shared_storage.size_in_bytes(), + # Caps registers at 128/thread so a quantization CTA can be co-resident. + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB: cute.Tensor, + xs_ptr: cute.Pointer, + mWS: cute.Tensor, + mOut: cute.Tensor, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + ar_peers: Optional[Int64], + ar_mc: Optional[Int64], + ar_flags: Optional[Int64], + ar_rank: Optional[Int32], + ): + S = self.split_k + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + rank = Int32(0) + if cutlass.const_expr(S > 1): + rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_tile = bidx // S + + kt_base = self.kt_total // S + kt_rem = self.kt_total % S + kt_count = Int32(kt_base) + kt_start = rank * kt_base + if cutlass.const_expr(kt_rem > 0): + if rank < kt_rem: + kt_count = kt_count + 1 + kt_start = kt_start + rank + else: + kt_start = kt_start + kt_rem + + storage = utils.SmemAllocator().allocate(self.shared_storage) + a_full = storage.a_full.data_ptr() + a_empty = storage.a_empty.data_ptr() + b_full = storage.b_full.data_ptr() + b_empty = storage.b_empty.data_ptr() + + if warp_idx == _TMA_WARP: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(self.a_stages): + cute.arch.mbarrier_init(a_full + i, 1) + cute.arch.mbarrier_init(a_empty + i, 1) + for i in cutlass.range_constexpr(self.b_stages): + cute.arch.mbarrier_init(b_full + i, 1) + cute.arch.mbarrier_init(b_empty + i, 1) + cute.arch.mbarrier_init(storage.acc_full, 1) + cute.arch.mbarrier_init(storage.tmem_ready, 1) + cute.arch.mbarrier_init(storage.tmem_free, 4) + cute.arch.mbarrier_init(storage.mbox_full, 1) + cute.arch.mbarrier_init(storage.peers_done, max(S - 1, 1)) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + if cutlass.const_expr(S > 1): + cute.arch.cluster_arrive_relaxed() + + if warp_idx == _TMA_WARP: + # Weight producer: never waits on the previous grid. The all-reduce push variant + # releases its dependent after its own wait (below). + if cutlass.const_expr(self.use_pdl and not self.ar_push): + _griddepcontrol_launch_dependents() + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + gA = cute.local_tile( + mA, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + tCgA = tiled_mma.get_slice(0).partition_A(gA) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + for kt in cutlass.range(kt_count, unroll=1): + stage = kt % self.a_stages + if kt >= self.a_stages: + cute.arch.mbarrier_wait(a_empty + stage, ((kt // self.a_stages) & 1) ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(a_full + stage, self.a_stage_bytes) + cute.copy( + tma_atom_a, + tAgA[(None, n_tile, kt_start + kt, 0)], + tAsA[(None, stage)], + tma_bar_ptr=a_full + stage, + cache_policy=Int64(_EVICT_FIRST), + ) + elif warp_idx == self._B_WARP: + # Activation producer: E4M3 tile by TMA plus its scale block by a bulk copy. + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + if cutlass.const_expr(self.ar_push != 0): + # A finish kernel reads the Lamport rotation state as soon as it is + # resident (k3_lamport_ar): release it only once every earlier grid (the + # previous all-reduce on the workspace included) has completed. + _griddepcontrol_launch_dependents() + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + gB = cute.local_tile( + mB, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + tCgB = tiled_mma.get_slice(0).partition_B(gB) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + sfb_base = storage.sSFB.data_ptr().toint() + for kt in cutlass.range(kt_count, unroll=1): + stage = kt % self.b_stages + if kt >= self.b_stages: + cute.arch.mbarrier_wait(b_empty + stage, ((kt // self.b_stages) & 1) ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + b_full + stage, self.b_stage_bytes + self.sfb_stage_bytes + ) + _bulk_load_g2s( + sfb_base + stage * self.sfb_stage_bytes, + xs_ptr.toint() + Int64((kt_start + kt) * self.sfb_stage_bytes), + self.sfb_stage_bytes, + (b_full + stage).toint(), + ) + cute.copy( + tma_atom_b, + tBgB[(None, 0, kt_start + kt, 0)], + tBsB[(None, stage)], + tma_bar_ptr=b_full + stage, + ) + elif warp_idx == _MMA_WARP: + self.mma_warp( + tiled_mma, + storage, + mWS, + n_tile, + kt_start, + kt_count, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + ) + elif warp_idx >= self._EPI_WARP0: + self.epilogue( + tiled_mma, + storage, + mOut, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + tidx - 32 * self._EPI_WARP0, + n_tile, + rank, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ) + + @cute.jit + def mma_warp( + self, + tiled_mma: cute.TiledMma, + storage, + mWS: cute.Tensor, + n_tile: Int32, + kt_start: Int32, + kt_count: Int32, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + ): + a_full = storage.a_full.data_ptr() + a_empty = storage.a_empty.data_ptr() + b_full = storage.b_full.data_ptr() + b_empty = storage.b_empty.data_ptr() + lane = cute.arch.lane_idx() + + # SFA for every K tile: the weight block exponent replicated over 128 rows x 4 + # groups. Weights are static, so this needs no wait on the previous grid. + sfa_base = storage.sSFA.data_ptr().toint() + exps = [] + for j in cutlass.range_constexpr(_ceil_div(self.kt_max, 32)): + e = Int32(127) + if lane + 32 * j < kt_count: + e = Int32(mWS[n_tile, kt_start + lane + 32 * j]) + exps.append(e) + for kt in cutlass.range(kt_count, unroll=1): + e = cute.arch.shuffle_sync(exps[0], kt & 31) + for j in cutlass.range_constexpr(1, len(exps)): + ej = cute.arch.shuffle_sync(exps[j], kt & 31) + if (kt >> 5) == j: + e = ej + word = e * Int32(0x01010101) + _sts_v4_b32(sfa_base + kt * 512 + lane * 16, word, word, word, word) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + + cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding_buf) + cute.arch.relinquish_tmem_alloc_permit() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(storage.tmem_ready) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tiled_mma.make_fragment_C(acc_shape).layout) + tCtSFA = cute.make_tensor( + cute.recast_ptr(acc_tmem_ptr + self.mma_n, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, self.mma_tiler, SF_VEC, cute.slice_(sfa_layout, (None, None, None, 0)) + ), + ) + tCtSFB = cute.make_tensor( + cute.recast_ptr(acc_tmem_ptr + self.mma_n + _SF_TMEM_COLS, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, self.mma_tiler, SF_VEC, cute.slice_(sfb_layout, (None, None, None, 0)) + ), + ) + s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = _s2t_partition_sf( + storage.sSFA.get_tensor(sfa_layout), tCtSFA + ) + s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = _s2t_partition_sf( + storage.sSFB.get_tensor(sfb_layout), tCtSFB + ) + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + num_kphases = cute.size(tCrA, mode=[2]) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for kt in cutlass.range(kt_count, unroll=1): + a_stage = kt % self.a_stages + b_stage = kt % self.b_stages + cute.arch.mbarrier_wait(a_full + a_stage, (kt // self.a_stages) & 1) + cute.arch.mbarrier_wait(b_full + b_stage, (kt // self.b_stages) & 1) + cute.copy(s2t_sfa, tCsSFA_s2t[(None, None, None, None, kt)], tCtSFA_s2t) + cute.copy(s2t_sfb, tCsSFB_s2t[(None, None, None, None, b_stage)], tCtSFB_s2t) + for kp in cutlass.range_constexpr(num_kphases): + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator) + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[(None, None, kp, a_stage)], + tCrB[(None, None, kp, b_stage)], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + with cute.arch.elect_one(): + tcgen05.commit(a_empty + a_stage) + tcgen05.commit(b_empty + b_stage) + with cute.arch.elect_one(): + tcgen05.commit(storage.acc_full) + + cute.arch.mbarrier_wait(storage.tmem_free, 0) + cute.arch.dealloc_tmem(acc_tmem_ptr, self.tmem_cols) + + @cute.jit + def epilogue( + self, + tiled_mma: cute.TiledMma, + storage, + mOut: cute.Tensor, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + q: Int32, + n_tile: Int32, + rank: Int32, + ar_peers: Optional[Int64], + ar_mc: Optional[Int64], + ar_flags: Optional[Int64], + ar_rank: Optional[Int32], + ): + """TMEM -> registers (32-column blocks) -> token-major SMEM tile -> 16-byte stores + (BF16), :meth:`fused_store` or :meth:`_push_staged_rows` (``ar_push``). + + With ``split_k > 1`` rank ``r`` reduces every chunk ``c`` (4 tokens) with + ``c % split_k == r``; peers push their FP32 partials with ``st.async`` into its + dead weight ring once every CTA of the cluster has finished its mainloop. + """ + S = self.split_k + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + cute.arch.mbarrier_wait(storage.tmem_ready, 0) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tiled_mma.make_fragment_C(acc_shape).layout) + tAcc = tCtAcc[((None, None), 0, 0)] + cAcc = cute.make_identity_tensor((MMA_M, self.mma_n)) + blk = 32 + tiled_t2r = tcgen05.make_tmem_copy( + cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(blk)), cutlass.Float32), + cute.local_tile(tAcc, (MMA_M, blk), (0, 0)), + ) + thr_t2r = tiled_t2r.get_slice(q) + out_base = storage.sB.data_ptr().toint() + out_gmem = mOut.iterator.toint() + Int64(n_tile * MMA_M * 2) + row_stride = Int64(self.n * 2) + # The output stores below give thread q the 8 weight rows starting at (q % 16) * 8 + # of the tile for every token; past N (partial last tile) it stores nothing. + store_cols = Int32(16) + if cutlass.const_expr(self.n % MMA_M != 0): + if n_tile * MMA_M + (q & 15) * 8 >= self.n: + store_cols = Int32(0) + + cute.arch.mbarrier_wait(storage.acc_full, 0) + if cutlass.const_expr(S > 1): + # Every CTA's rings are dead once its MMAs completed: the weight ring becomes + # the mailbox. Tell the peers, then wait for theirs before pushing partials. + cute.arch.cluster_wait() + if q == 0: + for p in cutlass.range_constexpr(1, S): + _mbarrier_arrive_remote(storage.peers_done.toint(), (rank + p) % S) + cute.arch.mbarrier_wait(storage.peers_done, 0) + + mbox_base = storage.sA.data_ptr().toint() + mbar_addr = storage.mbox_full.toint() + slot_bytes = self.chunks_cap * MMA_M * 16 + # This rank's own partials of the chunks it reduces, indexed by c // S. + rOwn = cute.make_rmem_tensor((self.chunks_cap * 4,), cutlass.Float32) + nblk = self.mma_n // blk + tTR_cAcc0 = thr_t2r.partition_D(cute.local_tile(cAcc, (MMA_M, blk), (0, 0))) + # Double-buffered TMEM loads: block j + 1 is in flight while block j is consumed. + rBuf = [cute.make_rmem_tensor(tTR_cAcc0.shape, cutlass.Float32) for _ in range(2)] + cute.copy( + tiled_t2r, thr_t2r.partition_S(cute.local_tile(tAcc, (MMA_M, blk), (0, 0))), rBuf[0] + ) + for j in cutlass.range_constexpr(nblk): + cute.arch.fence_view_async_tmem_load() + rAcc = rBuf[j % 2] + if cutlass.const_expr(j + 1 < nblk): + cute.copy( + tiled_t2r, + thr_t2r.partition_S(cute.local_tile(tAcc, (MMA_M, blk), (0, j + 1))), + rBuf[(j + 1) % 2], + ) + for c4 in cutlass.range_constexpr(blk // 4): + c = j * (blk // 4) + c4 + if cutlass.const_expr(S == 1): + if 4 * c < m: + for v in cutlass.range_constexpr(4): + _sts_u16( + out_base + (4 * c + v) * (MMA_M * 2) + q * 2, + _f32_to_bf16_bits(rAcc[4 * c4 + v]), + ) + else: + if rank == c % S: + for v in cutlass.range_constexpr(4): + rOwn[(c // S) * 4 + v] = rAcc[4 * c4 + v] + else: + if 4 * c < m: + slot = rank + if rank > c % S: + slot = rank - 1 + _st_async_v4_f32( + mbox_base + slot * slot_bytes + ((c // S) * MMA_M + q) * 16, + rAcc[4 * c4], + rAcc[4 * c4 + 1], + rAcc[4 * c4 + 2], + rAcc[4 * c4 + 3], + mbar_addr, + c % S, + ) + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(storage.tmem_free) + + # Staged BF16 token rows (256 B) go out with 16-byte stores; 16 consecutive threads + # cover one row, so every store is a full line. + if cutlass.const_expr(S == 1): + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + if cutlass.const_expr(self.ar_push != 0): + self._push_staged_rows( + out_base, m, rank, m, n_tile, q, store_cols, ar_peers, ar_mc, ar_flags, ar_rank + ) + elif cutlass.const_expr(self.epilogue_kind == EPI_BF16): + for p in cutlass.range(q, m * store_cols, MMA_M, unroll=1): + t = p >> 4 + col = (p & 15) * 16 + w0, w1, w2, w3 = _lds_v4_b32(out_base + t * (MMA_M * 2) + col) + _stg_v4_b32(out_gmem + Int64(t) * row_stride + col, w0, w1, w2, w3) + else: + self.fused_store( + out_base, m, rank, m, n_tile, q, mOut.iterator, out1_ptr, aux0_ptr, aux1_ptr + ) + else: + # Owned chunks c = rank + i * S; the column of element v of chunk c is 4c + v. + active = (m + 3) >> 2 + my_chunks = Int32(0) + if active > rank: + my_chunks = (active - rank + S - 1) // S + if my_chunks > 0: + if q == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + storage.mbox_full, my_chunks * ((S - 1) * MMA_M * 16) + ) + cute.arch.mbarrier_wait(storage.mbox_full, 0) + for i in cutlass.range_constexpr(self.chunks_cap): + if i < my_chunks: + acc = [rOwn[i * 4 + v] for v in range(4)] + for p in cutlass.range_constexpr(S - 1): + v0, v1, v2, v3 = _lds_v4_f32( + mbox_base + p * slot_bytes + (i * MMA_M + q) * 16 + ) + acc[0] = acc[0] + v0 + acc[1] = acc[1] + v1 + acc[2] = acc[2] + v2 + acc[3] = acc[3] + v3 + for v in cutlass.range_constexpr(4): + _sts_u16( + out_base + (4 * (rank + i * S) + v) * (MMA_M * 2) + q * 2, + _f32_to_bf16_bits(acc[v]), + ) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + if cutlass.const_expr(self.ar_push != 0): + self._push_staged_rows( + out_base, + my_chunks * 4, + rank, + m, + n_tile, + q, + store_cols, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ) + elif cutlass.const_expr(self.epilogue_kind == EPI_BF16): + # The r-th owned row is token 4 * (rank + (r // 4) * S) + r % 4. + for p in cutlass.range(q, my_chunks * 4 * store_cols, MMA_M, unroll=1): + r = p >> 4 + t = 4 * (rank + (r >> 2) * S) + (r & 3) + if t < m: + col = (p & 15) * 16 + w0, w1, w2, w3 = _lds_v4_b32(out_base + t * (MMA_M * 2) + col) + _stg_v4_b32(out_gmem + Int64(t) * row_stride + col, w0, w1, w2, w3) + else: + self.fused_store( + out_base, + my_chunks * 4, + rank, + m, + n_tile, + q, + mOut.iterator, + out1_ptr, + aux0_ptr, + aux1_ptr, + ) + + @cute.jit + def _push_staged_rows( + self, + stage_base, + nrows: Int32, + rank: Int32, + m: Int32, + n_tile: Int32, + q: Int32, + store_cols: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + ar_rank: Int32, + ): + """c240 all-reduce push of the staged BF16 rows (token ``t`` at ``stage_base + t * 256``): + rows ``r < nrows``, i.e. token ``r`` or with split-K this rank's ``r``-th owned token, go + out as the BF16 store's 16-byte vectors through ``k3_lamport_ar.push_slot`` / + ``push_v4`` (``-0`` sent as ``+0``), so every (token, feature) vector is pushed once. + The epilogue has passed ``griddepcontrol.wait``: the previous call on the workspace has + completed and its rotation state is final.""" + cur, _, nbytes, _, _, _, _, _ = k3_lamport_ar.lamport_state(flags) + one_shot = self.ar_push == AR_PUSH_ONE_SHOT + for p in cutlass.range(q, nrows * store_cols, MMA_M, unroll=1): + r = p >> 4 + t = r + if cutlass.const_expr(self.split_k > 1): + t = 4 * (rank + (r >> 2) * self.split_k) + (r & 3) + if t < m: + c = p & 15 + w0, w1, w2, w3 = _lds_v4_b32(stage_base + t * (MMA_M * 2) + c * 16) + dst = k3_lamport_ar.push_slot( + peer_ptrs, + mc_base, + cur, + nbytes, + t, + ar_rank, + n_tile * MMA_M + c * 8, + self.ar_world, + self.n, + one_shot, + ) + k3_lamport_ar.push_v4(dst, w0, w1, w2, w3, one_shot) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +@functools.cache +def _num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +@functools.cache +def _max_active_clusters(cluster_size: int) -> int: + return utils.HardwareInfo().get_max_active_clusters(cluster_size) + + +def mma_n_for(m: int) -> int: + for mma_n in SUPPORTED_MMA_N: + if m <= mma_n: + return mma_n + raise ValueError(f"M={m} exceeds {LOW_M_MAX}") + + +def mid_mma_n_for(m: int) -> int: + for mma_n in SUPPORTED_MID_MMA_N: + if m <= mma_n: + return mma_n + raise ValueError(f"M={m} exceeds {MAX_M}") + + +def _tile_n(n: int) -> int: + """N rounded up to whole 128-row weight tiles: the key of the per-shape tables below. + + A partial last tile costs the same CTAs, cluster shape and per-CTA weight slice as a + full one, so e.g. the 2112-row MLA ``[q_a; kv_a; k_pe]`` projection uses the 2176 entries. + """ + return _align_up(n, MMA_M) + + +#: Tuned on GB300 for the Kimi-K3 Helix decode projections (CUDA graph, weights not +#: L2-resident): (N, K, mma_n) -> (split_k, x_stages, b_stages, quant_warps, +#: quant_groups), N in whole tiles (:func:`_tile_n`). The weight ring then takes all +#: remaining shared memory. +_TUNED = { + # kda_qkvg + (3072, 7168, 8): (5, 12, 4, 4, 4), + (3072, 7168, 16): (4, 8, 4, 4, 2), + (3072, 7168, 32): (5, 8, 4, 8, 4), + (3072, 7168, 64): (5, 4, 4, 8, 4), + # latent_down + (3584, 7168, 8): (4, 14, 4, 4, 4), + (3584, 7168, 16): (4, 8, 4, 4, 2), + (3584, 7168, 32): (4, 8, 4, 8, 2), + (3584, 7168, 64): (4, 4, 4, 8, 4), + # latent_up + (7168, 3584, 8): (2, 14, 4, 4, 4), + (7168, 3584, 16): (2, 8, 4, 4, 2), + (7168, 3584, 32): (2, 8, 4, 8, 2), + (7168, 3584, 64): (2, 4, 4, 8, 4), + # kda_o / mla_o + (7168, 768, 8): (1, 6, 4, 4, 4), + (7168, 768, 16): (1, 6, 4, 8, 4), + (7168, 768, 32): (1, 6, 4, 8, 4), + (7168, 768, 64): (1, 4, 4, 8, 2), + # mla_q_a_kv_a (2112 rows: q_lora 1536 + kv_lora 512 + rope 64) + (2176, 7168, 8): (6, 10, 4, 4, 4), + (2176, 7168, 16): (6, 8, 4, 8, 4), + (2176, 7168, 32): (6, 8, 4, 8, 4), + (2176, 7168, 64): (6, 4, 4, 8, 4), + # mla_q_b + (18432, 1536, 8): (1, 12, 4, 4, 4), + (18432, 1536, 16): (1, 12, 4, 4, 2), + (18432, 1536, 32): (1, 8, 4, 8, 2), + (18432, 1536, 64): (1, 4, 4, 8, 2), +} + +#: Minimum weight-ring depth the heuristic keeps when growing the activation ring. +_MIN_A_STAGES = 6 + + +def _heuristic_split(n: int, k: int) -> int: + """Largest cluster split-K that still runs every cluster in a single wave.""" + kt_total = k // CTA_K + n_tiles = _ceil_div(n, MMA_M) + best, split = 0, 1 + for s in range(1, min(8, kt_total) + 1): + if n_tiles > _max_active_clusters(s) or n_tiles * s > _num_sms(): + continue + if s > 1 and kt_total // s < 3: + continue + if n_tiles * s > best: + best, split = n_tiles * s, s + return split + + +@functools.cache +def default_tactic(m: int, n: int, k: int) -> K3LowMTactic: + """Deterministic tactic for ``(M, N, K)``: tuned table first, heuristic otherwise.""" + mma_n = mma_n_for(m) + tuned = _tables().tuned.get((_tile_n(n), k, mma_n)) + if tuned is not None: + split, x_stages, b_stages, quant_warps, quant_groups = tuned + base = K3LowMTactic( + mma_n=mma_n, + split_k=split, + a_stages=1, + x_stages=x_stages, + b_stages=b_stages, + quant_warps=quant_warps, + quant_groups=quant_groups, + ) + return dataclasses.replace(base, a_stages=max_a_stages(n, k, base)) + split = _heuristic_split(n, k) + kt_max = _ceil_div(k // CTA_K, split) + quant_warps = 4 if mma_n <= 8 else 8 + tactic = None + # Keep every activation tile in flight if the weight ring stays deep enough. + for x_stages in (kt_max, 8, 4): + base = K3LowMTactic( + mma_n=mma_n, + split_k=split, + a_stages=1, + x_stages=min(x_stages, kt_max), + b_stages=min(4, kt_max), + quant_warps=quant_warps, + quant_groups=4, + ) + try: + validate_tactic(n, k, base) + except ValueError: + continue + tactic = dataclasses.replace(base, a_stages=max_a_stages(n, k, base)) + if tactic.a_stages >= min(_MIN_A_STAGES, kt_max): + break + if tactic is None: + raise ValueError(f"no valid tactic for M={m}, N={n}, K={k}") + return tactic + + +@functools.cache +def _compiled_kernel( + n: int, + k: int, + tactic: K3LowMTactic, + use_pdl: bool, + epilogue: str = EPI_BF16, + sigmoid: bool = False, +): + kernel = K3Fp8wLowMGemmKernel(n, k, tactic, use_pdl, epilogue, sigmoid=sigmoid) + dummy = 1 << 20 + w_ptr = make_ptr(cutlass.Float8E4M3FN, dummy, cute.AddressSpace.gmem, assumed_align=16) + x_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + ws_ptr = make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16) + out_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + aux = [make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16)] * 3 + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile(kernel, w_ptr, x_ptr, ws_ptr, out_ptr, *aux, Int32(1), stream) + + +@functools.cache +def gated_tactic(m: int, n: int, k: int) -> K3LowMTactic: + """:func:`default_tactic` with the weight ring shortened, if needed, to fit the gate ring. + + Split-K and the quantization layout are unchanged, so the gated GEMM rounds exactly like + the plain GEMM on ``bf16(x * gate)``. + """ + tactic = default_tactic(m, n, k) + a_stages = min(tactic.a_stages, max_a_stages(n, k, tactic, gated=True)) + tactic = dataclasses.replace(tactic, a_stages=a_stages) + validate_tactic(n, k, tactic, gated=True) + return tactic + + +@functools.cache +def _gated_tactic_fits(m: int, n: int, k: int) -> bool: + """Whether :func:`gated_tactic` has a tactic for this shape (the gate ring and a weight ring + of at least the minimum depth fit in shared memory).""" + try: + gated_tactic(m, n, k) + except ValueError: + return False + return True + + +@functools.cache +def _compiled_gated_kernel(n: int, k: int, tactic: K3LowMTactic, use_pdl: bool): + kernel = K3Fp8wLowMGemmKernel(n, k, tactic, use_pdl, gated=True) + dummy = 1 << 20 + w_ptr = make_ptr(cutlass.Float8E4M3FN, dummy, cute.AddressSpace.gmem, assumed_align=16) + x_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + g_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + ws_ptr = make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16) + out_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + kernel.call_gated, w_ptr, x_ptr, g_ptr, Int32(k), ws_ptr, out_ptr, Int32(1), stream + ) + + +#: Mid-M tactics tuned on GB300 (CUDA graph, weights not L2-resident; score = alone + +#: after a 3 us latency-bound kernel): (N, K, mma_n) -> (split_k, b_stages). Buckets not +#: listed use the (N, K) default below. +_MID_TUNED = { + (3072, 7168, 224): (4, 4), + (3072, 7168, 256): (4, 3), + (3584, 7168, 224): (3, 4), + (3584, 7168, 256): (3, 3), + (7168, 3584, 224): (2, 3), + (7168, 3584, 256): (2, 3), + (2176, 7168, 128): (6, 4), + (2176, 7168, 192): (6, 4), + (2176, 7168, 224): (4, 4), + (2176, 7168, 256): (6, 3), + (12288, 7168, 192): (1, 3), + (12288, 7168, 224): (1, 3), + (12288, 7168, 256): (1, 3), + (7168, 6144, 256): (2, 3), +} +_MID_DEFAULT = { + (3072, 7168): (4, 4), + (3584, 7168): (4, 4), + (7168, 3584): (2, 4), + (7168, 768): (1, 4), + (2176, 7168): (5, 4), + (18432, 1536): (1, 4), + (12288, 7168): (1, 4), + (7168, 6144): (2, 4), +} + + +#: Largest M per (N, K) for which this kernel (with its quantization kernel) beats +#: ``fp8_swap_ab_gemm`` (with its quantization kernel) on GB300, both standalone and behind +#: a latency-bound kernel with PDL. Used by :func:`can_run`. +_GATED_MAX_M = { + (3072, 7168): 224, + (3584, 7168): 224, + (7168, 3584): 144, + (7168, 768): 192, + (2176, 7168): 224, + (18432, 1536): 224, + (12288, 7168): 128, + (7168, 6144): 128, +} +#: Gate for shapes not in the table (every measured shape wins up to M = 128). +_DEFAULT_GATED_MAX_M = 128 +#: Shapes whose low-M kernel is slower than the mid-M path above this M (large weights on +#: fewer than all SMs, where the in-kernel quantization of a 64-token tile is exposed). +_LOW_M_CUTOVER = { + (12288, 7168): 32, + (7168, 6144): 32, +} + + +# Rubin (SM107, 212 SMs) tables for the Kimi-K3 decode projections, keyed like the GB300 +# ones above: the dep16 (attention-DP) weights, the replicated MLA / MoE latent weights and +# the helix8 / helix16 head shards (KDA fused qkvg, KDA / MLA o_proj, MLA output gate). +# Measured on a Rubin GPU with cold weights (rotated over 1.5 GiB): each tactic scores the +# mean CUDA-graph per-call time of back-to-back launches with and without PDL. Low-M path, +# per mma_n bucket timed at M = mma_n: a split_k sweep with the heuristic ring sizes, then +# quant warps / groups and x / b ring depths around the three best splits; mid-M path, per +# mma_n bucket (the 160-token bucket timed at M = 144): split_k x b_stages. Both grids are one +# wave of N/128 * split_k CTAs (not persistent), so the persistent-grid SM reserve +# (cute_dsl_utils.sm_reserve) does not apply. +# With 212 SMs the best split is larger than on GB300 for the small-N weights (latent down 8 +# vs 4, latent up 6-7 vs 2, KDA qkvg helix16 6-7 vs 4-5, MLA q_a|kv_a 7-10 vs 6, MLA g helix16 +# 10-16 clusters). +_SM107_TUNED = { + # KDA fused q/k/v/g (dep16) + (49152, 7168, 8): (1, 8, 4, 4, 4), + (49152, 7168, 16): (1, 8, 4, 8, 4), + (49152, 7168, 32): (1, 8, 4, 8, 4), + (49152, 7168, 64): (1, 4, 4, 8, 4), + # KDA / MLA o_proj (dep16) + (7168, 12288, 8): (3, 8, 4, 4, 4), + (7168, 12288, 16): (3, 8, 4, 8, 4), + (7168, 12288, 32): (3, 8, 4, 8, 4), + (7168, 12288, 64): (3, 4, 4, 8, 4), + # MLA output gate / shared-expert gate_up (dep16) + (12288, 7168, 8): (2, 8, 4, 4, 4), + (12288, 7168, 16): (2, 8, 4, 8, 4), + (12288, 7168, 32): (2, 8, 4, 8, 4), + (12288, 7168, 64): (2, 4, 4, 8, 4), + # shared-expert down (dep16) + (7168, 6144, 8): (6, 8, 4, 4, 4), + (7168, 6144, 16): (3, 8, 4, 8, 4), + (7168, 6144, 32): (3, 8, 4, 8, 4), + (7168, 6144, 64): (3, 4, 4, 8, 4), + # MLA q_b + (18432, 1536, 8): (2, 4, 4, 4, 4), + (18432, 1536, 16): (2, 4, 4, 8, 4), + (18432, 1536, 32): (1, 8, 4, 8, 4), + (18432, 1536, 64): (1, 4, 4, 8, 4), + # MoE latent down + (3584, 7168, 8): (8, 7, 4, 4, 4), + (3584, 7168, 16): (8, 4, 4, 8, 4), + (3584, 7168, 32): (5, 8, 4, 8, 4), + (3584, 7168, 64): (6, 4, 4, 8, 2), + # MoE latent up + (7168, 3584, 8): (7, 4, 4, 4, 4), + (7168, 3584, 16): (6, 4, 4, 8, 4), + (7168, 3584, 32): (6, 5, 4, 8, 4), + (7168, 3584, 64): (3, 4, 4, 8, 4), + # KDA fused q/k/v/g (helix16 shard) + (3072, 7168, 8): (7, 8, 4, 4, 4), + (3072, 7168, 16): (7, 4, 4, 8, 4), + (3072, 7168, 32): (6, 8, 4, 8, 4), + (3072, 7168, 64): (6, 4, 4, 8, 2), + # MLA fused q_a / kv_a (2112 rows) + (2176, 7168, 8): (10, 4, 4, 4, 4), + (2176, 7168, 16): (7, 8, 4, 8, 4), + (2176, 7168, 32): (8, 4, 4, 8, 4), + (2176, 7168, 64): (8, 7, 4, 8, 4), + # MLA output gate (helix16 shard) + (768, 7168, 8): (14, 4, 4, 4, 4), + (768, 7168, 16): (14, 4, 4, 8, 4), + (768, 7168, 32): (10, 6, 4, 8, 4), + (768, 7168, 64): (16, 4, 4, 8, 4), + # KDA / MLA o_proj (helix16 shard) + (7168, 768, 8): (3, 2, 2, 4, 4), + (7168, 768, 16): (3, 2, 2, 8, 4), + (7168, 768, 32): (1, 6, 4, 8, 4), + (7168, 768, 64): (2, 3, 3, 8, 4), + # KDA fused q/k/v/g (helix8 shard) + (6144, 7168, 8): (3, 4, 4, 4, 4), + (6144, 7168, 16): (3, 8, 4, 8, 4), + (6144, 7168, 32): (3, 8, 4, 8, 4), + (6144, 7168, 64): (3, 4, 4, 8, 4), + # KDA / MLA o_proj (helix8 shard) + (7168, 1536, 8): (4, 3, 3, 4, 4), + (7168, 1536, 16): (4, 3, 3, 8, 4), + (7168, 1536, 32): (2, 6, 4, 8, 4), + (7168, 1536, 64): (2, 4, 4, 8, 4), + # MLA output gate (helix8 shard) + (1536, 7168, 8): (12, 5, 4, 4, 4), + (1536, 7168, 16): (7, 8, 4, 8, 4), + (1536, 7168, 32): (8, 7, 4, 8, 4), + (1536, 7168, 64): (12, 5, 4, 8, 2), +} +_SM107_MID_TUNED = { + (49152, 7168, 64): (1, 4), + (49152, 7168, 96): (1, 4), + (49152, 7168, 128): (1, 4), + (7168, 12288, 64): (3, 4), + (7168, 12288, 96): (3, 4), + (7168, 12288, 128): (3, 4), + (12288, 7168, 64): (2, 4), + (12288, 7168, 96): (2, 4), + (12288, 7168, 128): (2, 4), + (7168, 6144, 64): (2, 4), + (7168, 6144, 96): (2, 4), + (7168, 6144, 128): (2, 4), + (18432, 1536, 64): (1, 4), + (18432, 1536, 96): (1, 4), + (18432, 1536, 128): (1, 4), + (18432, 1536, 160): (1, 4), + (3584, 7168, 64): (5, 4), + (3584, 7168, 96): (5, 4), + (3584, 7168, 128): (4, 4), + (3584, 7168, 160): (4, 4), + (7168, 3584, 64): (2, 4), + (7168, 3584, 96): (2, 4), + (7168, 3584, 128): (2, 4), + (3072, 7168, 64): (6, 4), + (3072, 7168, 96): (6, 4), + (3072, 7168, 128): (5, 4), + (3072, 7168, 160): (6, 4), + (2176, 7168, 64): (7, 4), + (2176, 7168, 96): (7, 4), + (2176, 7168, 128): (7, 4), + (2176, 7168, 160): (8, 4), + (2176, 7168, 192): (8, 4), + (768, 7168, 64): (12, 4), + (768, 7168, 96): (8, 4), + (768, 7168, 128): (8, 4), + (768, 7168, 160): (10, 4), + (768, 7168, 192): (8, 4), + (768, 7168, 224): (7, 4), + (768, 7168, 256): (8, 3), + (7168, 768, 64): (3, 3), + (7168, 768, 96): (1, 4), + (7168, 768, 128): (1, 4), + (6144, 7168, 64): (3, 4), + (6144, 7168, 96): (3, 4), + (6144, 7168, 128): (3, 4), + (6144, 7168, 160): (3, 4), + (7168, 1536, 64): (4, 3), + (7168, 1536, 96): (1, 4), + (7168, 1536, 128): (1, 4), + (1536, 7168, 64): (8, 4), + (1536, 7168, 96): (7, 4), + (1536, 7168, 128): (8, 4), + (1536, 7168, 160): (10, 3), + (1536, 7168, 192): (10, 3), + (1536, 7168, 224): (10, 3), +} +#: Mid-M buckets without an entry above (beyond the gate) reuse the 128-token tactic. +_SM107_MID_DEFAULT = { + (49152, 7168): (1, 4), + (7168, 12288): (3, 4), + (12288, 7168): (2, 4), + (7168, 6144): (2, 4), + (18432, 1536): (1, 4), + (3584, 7168): (4, 4), + (7168, 3584): (2, 4), + (3072, 7168): (5, 4), + (2176, 7168): (7, 4), + (768, 7168): (8, 4), + (7168, 768): (1, 4), + (6144, 7168): (3, 4), + (7168, 1536): (1, 4), + (1536, 7168): (8, 4), +} +#: Largest M per (N, K) up to which this kernel (low-M, or mid-M above the cut-over) runs +#: ``_Fp8BlockScaleWeightReadLinear.forward`` faster than the SM107 production path +#: (``fp8_quantize_1x128_cutedsl_ue8m0`` + autotuned ``cute_dsl_mxfp8_gemm_rubin``), CUDA graph +#: per call, cold weights. dep16-only weights: at least 5% faster at every measured M in +#: {1, 8, 16, 32, 48, 64, 96, 128} (back-to-back launches with PDL; shapes close to their gate +#: re-measured on a second GPU). Weights of the helix decode (head shards, MLA q_a|kv_a / q_b, +#: MoE latent): at every measured M in {8, 16, 32, 48, 64, 96, 128, 144, 192, 224, 256} up to +#: the gate at least 5% faster with PDL and not slower on the mean of the times with and +#: without PDL, against the mean of two runs of the CuTe path (each autotuned on its own; its +#: per-run time differs by up to ~1 us for these weights); a gate above 64 ends at its mid-M +#: token bucket. Other shapes stop at 32 tokens, where every K3 shape wins. +_SM107_GATED_MAX_M = { + (49152, 7168): 32, + (7168, 12288): 96, + (12288, 7168): 96, + (7168, 6144): 128, + (18432, 1536): 160, + (3584, 7168): 160, + (7168, 3584): 96, + (3072, 7168): 160, + (2176, 7168): 192, + (768, 7168): 256, + (7168, 768): 128, + (6144, 7168): 160, + (7168, 1536): 96, + (1536, 7168): 224, +} +_SM107_DEFAULT_GATED_MAX_M = 32 +#: Low-M kernel up to 32 tokens where the mid-M path (separate quantization, no per-CTA +#: re-quantization of the activation tile) is faster for the 33..64-token bucket. +_SM107_LOW_M_CUTOVER = { + (49152, 7168): 32, + (7168, 12288): 32, + (12288, 7168): 32, + (7168, 6144): 32, + (18432, 1536): 32, + (7168, 3584): 32, + (3072, 7168): 32, + (2176, 7168): 32, + (6144, 7168): 32, +} +#: Largest M per (N, K) for which this kernel on the BF16 activation (quantized inside the +#: kernel) beats the CuTe MXFP8 GEMM alone on an activation its producer already quantized +#: (``_Fp8BlockScaleWeightReadLinear.forward_prequantized``, the KDA MTP verify o_proj), by the +#: helix gate rule above. Other shapes always take the prequantized operand. +_SM107_BF16_OVER_PREQUANTIZED_MAX_M = { + (7168, 1536): 32, + (7168, 768): 32, +} + + +@dataclasses.dataclass(frozen=True) +class _ShapeTables: + """Per-architecture tactic and gating tables (see :func:`_tables`).""" + + #: (N, K, mma_n) -> (split_k, x_stages, b_stages, quant_warps, quant_groups), low-M path. + tuned: dict + #: (N, K, mma_n) -> (split_k, b_stages), mid-M path; ``mid_default`` per (N, K) otherwise. + mid_tuned: dict + mid_default: dict + #: (N, K) -> largest M :func:`can_run` accepts; ``default_gated_max_m`` for other shapes. + gated_max_m: dict + default_gated_max_m: int + #: (N, K) -> largest M of the low-M kernel (mid-M path above); LOW_M_MAX otherwise. + low_m_cutover: dict + + +_SM100_TABLES = _ShapeTables( + _TUNED, _MID_TUNED, _MID_DEFAULT, _GATED_MAX_M, _DEFAULT_GATED_MAX_M, _LOW_M_CUTOVER +) +_SM107_TABLES = _ShapeTables( + _SM107_TUNED, + _SM107_MID_TUNED, + _SM107_MID_DEFAULT, + _SM107_GATED_MAX_M, + _SM107_DEFAULT_GATED_MAX_M, + _SM107_LOW_M_CUTOVER, +) + + +@functools.cache +def _device_capability(device_index: int) -> tuple: + return torch.cuda.get_device_capability(device_index) + + +def _tables() -> _ShapeTables: + """Tables of the current device: Rubin (SM107) or GB300 (SM100/SM103 and the rest).""" + if _device_capability(torch.cuda.current_device()) == (10, 7): + return _SM107_TABLES + return _SM100_TABLES + + +def max_m_for(n: int, k: int) -> int: + """Largest M for which :func:`can_run` selects this kernel for an ``[N, K]`` weight.""" + tables = _tables() + return min(MAX_M, tables.gated_max_m.get((_tile_n(n), k), tables.default_gated_max_m)) + + +def low_m_max_for(n: int, k: int) -> int: + """Largest M served by the low-M kernel for an ``[N, K]`` weight (mid-M path above).""" + return _tables().low_m_cutover.get((_tile_n(n), k), LOW_M_MAX) + + +#: Rows above the table's cut-over for every shape (the KDA verify then hands o_proj BF16 rows with its +#: fused forget gate and PDL instead of the MXFP8 operand); 0 keeps the table. +_BF16_OVER_PREQUANTIZED_MAX_M_ENV = 0 + + +def bf16_over_prequantized_max_m(n: int, k: int) -> int: + """Largest M up to which a producer that can emit the SM107 MXFP8 operand of an ``[N, K]`` + weight should rather hand its BF16 output to this kernel (0: never; always 0 off SM107).""" + if _device_capability(torch.cuda.current_device()) != (10, 7): + return 0 + table = _SM107_BF16_OVER_PREQUANTIZED_MAX_M.get((_tile_n(n), k), 0) + return min(max_m_for(n, k), max(table, _BF16_OVER_PREQUANTIZED_MAX_M_ENV)) + + +@functools.cache +def default_mid_tactic(m: int, n: int, k: int) -> K3MidMTactic: + """Deterministic mid-M tactic: tuned table first, heuristic otherwise.""" + mma_n = mid_mma_n_for(m) + tables = _tables() + tuned = tables.mid_tuned.get((_tile_n(n), k, mma_n)) or tables.mid_default.get((_tile_n(n), k)) + split, b_stages = tuned if tuned is not None else (_heuristic_split(n, k), 4) + base = K3MidMTactic(mma_n=mma_n, split_k=split, a_stages=1, b_stages=b_stages) + tactic = dataclasses.replace(base, a_stages=max_mid_a_stages(n, k, base)) + # A per-shape default (the SM107 tables tune each shape's buckets only up to its gate) may + # not fit shared memory with the wider token tiles: drop activation stages, then split-K + # ranks. + while mid_smem_bytes(n, k, tactic) > _MID_SMEM_CAPACITY and ( + tactic.b_stages > 2 or tactic.split_k > 1 + ): + if tactic.b_stages > 2: + base = dataclasses.replace(base, b_stages=base.b_stages - 1) + else: + base = dataclasses.replace(base, split_k=base.split_k - 1) + tactic = dataclasses.replace(base, a_stages=max_mid_a_stages(n, k, base)) + return tactic + + +@functools.cache +def split_tactic(m: int, n: int, k: int, split_k: int, mid: bool) -> K3LowMTactic | K3MidMTactic: + """The default tactic for M tokens with ``split_k`` cluster ranks instead of the tuned count: + the low-M kernel's (its ring depths rounded down to whole quantization groups) or, with + ``mid``, the mid-M kernel's (4 activation stages, fewer if they do not fit).""" + if mid: + for b_stages in (4, 3, 2): + base = K3MidMTactic( + mma_n=mid_mma_n_for(m), split_k=split_k, a_stages=1, b_stages=b_stages + ) + tactic = dataclasses.replace(base, a_stages=max_mid_a_stages(n, k, base)) + try: + validate_mid_tactic(n, k, tactic) + except ValueError: + continue + return tactic + raise ValueError(f"no mid-M tactic with split_k={split_k} for M={m}, N={n}, K={k}") + base = default_tactic(m, n, k) + kt_max = _ceil_div(k // CTA_K, split_k) + groups = base.quant_groups + + def whole_groups(stages): + return max(groups, min(stages, kt_max) // groups * groups) + + tactic = dataclasses.replace( + base, + split_k=split_k, + a_stages=1, + x_stages=whole_groups(base.x_stages), + b_stages=max(2, whole_groups(base.b_stages)), + ) + tactic = dataclasses.replace(tactic, a_stages=max_a_stages(n, k, tactic)) + validate_tactic(n, k, tactic) + return tactic + + +@functools.cache +def _compiled_quant(k: int, sf_blocks: int, use_pdl: bool): + kernel = K3MxFp8QuantKernel(k, sf_blocks, use_pdl) + dummy = 1 << 20 + x_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + xq_ptr = make_ptr(cutlass.Float8E4M3FN, dummy, cute.AddressSpace.gmem, assumed_align=16) + xs_ptr = make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile(kernel, x_ptr, xq_ptr, xs_ptr, Int32(1), stream) + + +@functools.cache +def _compiled_mid_kernel( + n: int, k: int, tactic: K3MidMTactic, use_pdl: bool, epilogue: str = EPI_BF16 +): + kernel = K3Fp8wMidMGemmKernel(n, k, tactic, use_pdl, epilogue) + dummy = 1 << 20 + ptrs = [ + make_ptr(dtype, dummy, cute.AddressSpace.gmem, assumed_align=16) + for dtype in ( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + cutlass.Uint8, + cutlass.Uint8, + cutlass.BFloat16, + ) + + (cutlass.Uint8,) * 3 + ] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile(kernel, *ptrs, Int32(1), stream) + + +def _gmem_ptr(dtype, tensor: torch.Tensor): + return make_ptr(dtype, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def _validate_epilogue(n: int, epilogue: str) -> None: + if epilogue not in _EPILOGUES: + raise ValueError(f"epilogue must be one of {_EPILOGUES}, got {epilogue!r}") + if epilogue != EPI_BF16 and n % MMA_M: + raise ValueError(f"the {epilogue} epilogue needs N={n} to be a multiple of {MMA_M}") + + +def _run_mid( + x, w, w_scale, out, tactic: K3MidMTactic, use_pdl: bool, epilogue=EPI_BF16, extra=None +) -> torch.Tensor: + m, k = x.shape + n = w.shape[0] + sf_blocks = _sf_blocks(tactic.mma_n) + xq = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device) + xs = torch.empty((k // CTA_K) * sf_blocks * 512, dtype=torch.uint8, device=x.device) + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + _compiled_quant(k, sf_blocks, use_pdl)( + _gmem_ptr(cutlass.BFloat16, x), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + Int32(m), + stream, + ) + _compiled_mid_kernel(n, k, tactic, use_pdl, epilogue)( + _gmem_ptr(cutlass.Float8E4M3FN, w), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + _gmem_ptr(cutlass.Uint8, w_scale), + _gmem_ptr(cutlass.BFloat16, out), + *(_gmem_ptr(cutlass.Uint8, t) for t in (extra or (out, out, out))), + Int32(m), + stream, + ) + return out + + +def mid_prequant_plan(m: int, n: int, k: int) -> Optional[Tuple[K3MidMTactic, int]]: + """c240: ``(tactic, sf_blocks)`` of the mid-M path :func:`run_k3_fp8w_lowm_gemm` runs for ``m`` + rows of an ``[N, K]`` weight (its default tactic), or None where it runs the low-M kernel. + + A producer that stores the MXFP8 form of those rows exactly as :class:`K3MxFp8QuantKernel` + does (E4M3 ``[M, K]``; UE8M0 1x32 exponents in ``[K / 128][sf_blocks][512 B]`` tcgen05 atoms, + byte ``(r % 32) * 16 + (r % 128 // 32) * 4 + k32 % 4``) can hand them to + :func:`run_mid_prequantized` instead: bitwise the same output. + """ + if not low_m_max_for(n, k) < m <= MAX_M: + return None + tactic = default_mid_tactic(m, n, k) + if m > tactic.mma_n: + return None + return tactic, _sf_blocks(tactic.mma_n) + + +def run_mid_prequantized( + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + tactic: K3MidMTactic, + use_pdl: bool = True, +) -> torch.Tensor: + """c240: :func:`_run_mid` without its quantization kernel: the mid-M GEMM (BF16 epilogue into + ``out [M, N]``) on the MXFP8 operand ``xq`` (E4M3 ``[M, K]``) / ``xs`` (UE8M0, + ``(K / 128) * sf_blocks * 512`` bytes, see :func:`mid_prequant_plan`). With ``use_pdl`` it + launches as a programmatic dependent of the producer and streams its weights before its + ``griddepcontrol.wait``.""" + m, k = xq.shape + n = w.shape[0] + stream = cuda.CUstream(torch.cuda.current_stream(xq.device).cuda_stream) + _compiled_mid_kernel(n, k, tactic, use_pdl, EPI_BF16)( + _gmem_ptr(cutlass.Float8E4M3FN, w), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + _gmem_ptr(cutlass.Uint8, w_scale), + _gmem_ptr(cutlass.BFloat16, out), + *(_gmem_ptr(cutlass.Uint8, t) for t in (out, out, out)), + Int32(m), + stream, + ) + return out + + +def quantize_mid_operand( + x: torch.Tensor, tactic: K3MidMTactic, use_pdl: bool = True +) -> Tuple[torch.Tensor, torch.Tensor]: + """c240: the MXFP8 operand of the mid-M GEMM (``tactic``) for the BF16 rows ``x [M, K]``, as + :func:`run_k3_fp8w_lowm_gemm` quantizes them (:class:`K3MxFp8QuantKernel`): E4M3 ``[M, K]`` and + UE8M0 scales ``[(K / 128) * sf_blocks * 512]`` in that GEMM's layout, for + :func:`run_mid_prequantized` (``_ar_push``). With one scale block (``tactic.mma_n <= 128``) + the scales are also the R128c4 layout of the CuTe MXFP8 GEMM.""" + m, k = x.shape + sf_blocks = _sf_blocks(tactic.mma_n) + xq = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device) + xs = torch.empty((k // CTA_K) * sf_blocks * 512, dtype=torch.uint8, device=x.device) + _compiled_quant(k, sf_blocks, use_pdl)( + _gmem_ptr(cutlass.BFloat16, x), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + Int32(m), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + return xq, xs + + +def mid_ar_push_tactic(m: int, n: int, k: int) -> Optional[K3MidMTactic]: + """c240: the mid-M tactic :func:`run_k3_fp8w_lowm_gemm` runs for ``m`` rows of an ``[N, K]`` + weight (:func:`mid_prequant_plan`), which the mid-M all-reduce push epilogue runs with the + same bits; None where the low-M kernel serves ``m`` (see :func:`ar_push_tactic`) or above + the mid-M range.""" + plan = mid_prequant_plan(m, n, k) + if plan is None or n % 8: + return None + return plan[0] + + +@functools.cache +def _compiled_mid_ar_push_kernel( + n: int, k: int, tactic: K3MidMTactic, use_pdl: bool, ar_push: int, ar_world: int +): + kernel = K3Fp8wMidMGemmKernel(n, k, tactic, use_pdl, ar_push=ar_push, ar_world=ar_world) + dummy = 1 << 20 + ptrs = [ + make_ptr(dtype, dummy, cute.AddressSpace.gmem, assumed_align=16) + for dtype in (cutlass.Float8E4M3FN, cutlass.Float8E4M3FN, cutlass.Uint8, cutlass.Uint8) + ] + addr = Int64(dummy) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile(kernel.call_ar_push, *ptrs, Int32(1), addr, addr, addr, Int32(0), stream) + + +def run_mid_prequantized_ar_push( + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + tactic: K3MidMTactic, + use_pdl: bool = True, +) -> None: + """c240: :func:`run_mid_prequantized` with the all-reduce push epilogue: ``bf16(xq @ W^T)`` + goes into the Lamport slots of the next all-reduce on ``ws`` instead of an output tensor; + :func:`k3_lamport_ar.run_lamport_finish` with the same ``one_shot`` (or a consumer that + finishes it) completes it, bitwise the GEMM followed by the MNNVL all-reduce of its + output.""" + m, k = xq.shape + n = w.shape[0] + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + need = k3_lamport_ar.required_bytes(m, n, ws.world, one_shot) + if need > ws.bytes_per_buffer: + raise ValueError(f"Lamport buffer of {ws.bytes_per_buffer} B < {need} B") + mode = AR_PUSH_ONE_SHOT if one_shot else AR_PUSH_TWO_SHOT + compiled = _compiled_mid_ar_push_kernel(n, k, tactic, use_pdl, mode, ws.world) + stream = cuda.CUstream(torch.cuda.current_stream(xq.device).cuda_stream) + compiled( + _gmem_ptr(cutlass.Float8E4M3FN, w), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + _gmem_ptr(cutlass.Uint8, w_scale), + Int32(m), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(ws.rank), + stream, + ) + + +def pack_ue8m0_block_scale(w_scale: torch.Tensor, n: int, k: int) -> torch.Tensor: + """Pack power-of-two ``[ceil(N/128), K/128]`` FP32 block scales into DeepGEMM's layout. + + Produces the same tensor as ``transform_sf_into_required_layout(w_scale, mn=N, k=K, + recipe=(1, 128, 128), is_sfa=False)``: UE8M0 bytes replicated per row, four K blocks + per int32 word, MN-major ``int32 [N, ceil(K/512)]`` with strides ``(1, N)``. + """ + ue8m0 = ((w_scale.contiguous().view(torch.int32) >> 23) & 0xFF).to(torch.uint8) + return _pack_ue8m0_exponents(ue8m0, n, k) + + +def _pack_ue8m0_exponents(ue8m0: torch.Tensor, n: int, k: int) -> torch.Tensor: + """UE8M0 block exponents ``uint8 [ceil(N/128), K/128]`` -> DeepGEMM-packed int32 scale.""" + kb = k // CTA_K + rows = ue8m0.repeat_interleave(MMA_M, dim=0)[:n] + padded = torch.zeros((n, _ceil_div(kb, 4) * 4), dtype=torch.uint8, device=ue8m0.device) + padded[:, :kb] = rows + words = padded.view(torch.int32) + packed = torch.empty((words.shape[1], n), dtype=torch.int32, device=ue8m0.device).t() + packed.copy_(words) + return packed + + +def pack_block_scale_from_mxfp8_layout(sf: torch.Tensor, n: int, k: int) -> torch.Tensor: + """DeepGEMM-packed UE8M0 scale of a weight from its Rubin MXFP8 GEMM scale. + + ``sf`` is the R128c4 K32 UE8M0 tensor ``transform_k128_scales_to_cutedsl_mxfp8_layout`` + builds from 128x128 block scales (each block exponent replicated over its 128 rows and 4 + K32 slots; byte offset ``((r // 128) * (K / 128) + c // 4) * 512 + (r % 32) * 16 + + ((r % 128) // 32) * 4 + c % 4`` for row ``r``, K32 column ``c``). The first byte of every + 512-byte atom is the exponent of block ``(r // 128, c // 4)``; the result equals + :func:`pack_ue8m0_block_scale` of the original FP32 block scales. + """ + if k % CTA_K: + raise ValueError(f"K={k} must be a multiple of {CTA_K}") + nb, kb = _ceil_div(n, MMA_M), k // CTA_K + flat = sf.reshape(-1) + if sf.dtype != torch.uint8 or flat.numel() != nb * kb * 512: + raise ValueError(f"sf must be the uint8 R128c4 scale of an [{n}, {k}] weight") + exponents = flat.view(nb, kb, 512)[:, :, 0].contiguous() + return _pack_ue8m0_exponents(exponents, n, k) + + +def prepare_k3_fp8w_lowm_weight( + w_fp8: torch.Tensor, w_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """One-time weight preparation for standalone use. + + The kernel consumes exactly what ``_Fp8BlockScaleWeightReadLinear`` (and the other + ``fp8_swap_ab_gemm`` call sites) hold on SM100/SM103: the UE8M0-resmoothed FP8 weight + and the DeepGEMM-packed scale, so modules that already went through + ``resmooth_to_fp8_e8m0`` + ``transform_sf_into_required_layout`` need no extra + preparation or scale copy. + + Args: + w_fp8: ``[N, K]`` ``float8_e4m3fn`` weight (N a multiple of :data:`N_ALIGN`). + w_scale: ``[ceil(N/128), K/128]`` FP32 block scales. + + Returns: + ``(weight, packed_scale)``. Non-power-of-two scales are re-smoothed to UE8M0 + first (on a copy), which re-rounds the affected weight blocks. + """ + if w_fp8.dtype != torch.float8_e4m3fn or w_fp8.dim() != 2: + raise ValueError("w_fp8 must be a 2D float8_e4m3fn tensor") + n, k = w_fp8.shape + if n % N_ALIGN or k % CTA_K: + raise ValueError(f"N={n} must be a multiple of {N_ALIGN} and K={k} of {CTA_K}") + if tuple(w_scale.shape) != (_ceil_div(n, MMA_M), k // CTA_K): + raise ValueError(f"w_scale must have shape {(_ceil_div(n, MMA_M), k // CTA_K)}") + w_fp8 = w_fp8.contiguous() + w_scale = w_scale.to(device=w_fp8.device, dtype=torch.float32).contiguous() + bits = w_scale.view(torch.int32) + if bool(((bits & 0x7FFFFF) != 0).any()) or bool((w_scale <= 0).any()): + from tensorrt_llm.quantization.utils.fp8_utils import resmooth_to_fp8_e8m0 + + w_fp8, w_scale = resmooth_to_fp8_e8m0(w_fp8.clone(), w_scale.clone()) + return w_fp8, pack_ue8m0_block_scale(w_scale, n, k) + + +def run_k3_fp8w_lowm_gemm( + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + *, + tactic: K3LowMTactic | K3MidMTactic | None = None, + use_pdl: bool = True, +) -> torch.Tensor: + """Launch into ``out`` (``bf16 [M, N]``) on the current stream. + + ``M <= low_m_max_for(N, K)`` runs the low-M kernel (activation quantized inside the + GEMM); larger M runs a small MXFP8 quantization kernel followed by the mid-M GEMM, chained + with PDL so the GEMM's weight stream overlaps the quantization. + """ + _launch(x, w, w_scale, out, None, EPI_BF16, tactic, use_pdl) + return out + + +def run_k3_fp8w_lowm_gemm_fused( + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + epilogue: str, + out0: torch.Tensor, + out1: torch.Tensor, + aux0: torch.Tensor, + aux1: torch.Tensor, + *, + tactic: K3LowMTactic | K3MidMTactic | None = None, + use_pdl: bool = True, +) -> None: + """The GEMM of :func:`run_k3_fp8w_lowm_gemm` (same kernels, tactics and numerics) with a + fused epilogue on the BF16 result ``y``: + + * :data:`EPI_NVFP4`: ``out0`` / ``out1`` = ``fp4_quantize(y, aux0, 16)`` codes ``uint8 + [M, N / 2]`` / linear E4M3 scales ``uint8 [M, N / 16]``; ``aux0`` is the FP32 global + scale (``aux1`` unused). + + N must be a multiple of 128 and the low-M tactic must stage the output tile in its + quantized-activation ring (``b_stages >= 2``, true for every default tactic). + """ + if epilogue == EPI_BF16: + raise ValueError("use run_k3_fp8w_lowm_gemm for the BF16 output") + _launch(x, w, w_scale, out0, (out1, aux0, aux1), epilogue, tactic, use_pdl) + + +def _launch(x, w, w_scale, out, extra, epilogue, tactic, use_pdl) -> None: + m, k = x.shape + n = w.shape[0] + if not 1 <= m <= MAX_M: + raise ValueError(f"M={m} must be in [1, {MAX_M}]") + if tactic is None: + if m <= low_m_max_for(n, k): + tactic = default_tactic(m, n, k) + else: + tactic = default_mid_tactic(m, n, k) + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + if isinstance(tactic, K3MidMTactic): + _run_mid(x, w, w_scale, out, tactic, use_pdl, epilogue, extra) + return + compiled = _compiled_kernel(n, k, tactic, use_pdl, epilogue) + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + compiled( + make_ptr(cutlass.Float8E4M3FN, w.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, w_scale.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, out.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + *(_gmem_ptr(cutlass.Uint8, t) for t in (extra or (out, out, out))), + Int32(m), + stream, + ) + + +def run_k3_fp8w_lowm_gemm_gated( + x: torch.Tensor, + gate: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + *, + tactic: K3LowMTactic | None = None, + use_pdl: bool = True, +) -> torch.Tensor: + """Launch ``out = bf16(x * gate) @ dequant(w)^T`` into ``out`` on the current stream. + + Bitwise equal to :func:`run_k3_fp8w_lowm_gemm` on ``x * gate`` (PyTorch BF16 multiply) + with the same tactic. Low-M kernel only (``M <= low_m_max_for(N, K)``). ``gate`` is BF16 + ``[M, K]`` with unit column stride and a row stride that is a multiple of 8 elements; + like ``x`` it is read only after the kernel's ``griddepcontrol.wait``. + """ + m, k = x.shape + n = w.shape[0] + if not 1 <= m <= low_m_max_for(n, k): + raise ValueError(f"M={m} must be in [1, {low_m_max_for(n, k)}] for the gated GEMM") + if tactic is None: + tactic = gated_tactic(m, n, k) + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + compiled = _compiled_gated_kernel(n, k, tactic, use_pdl) + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + compiled( + make_ptr(cutlass.Float8E4M3FN, w.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, gate.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int32(gate.stride(0)), + make_ptr(cutlass.Uint8, w_scale.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, out.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int32(m), + stream, + ) + return out + + +@functools.cache +def _compiled_ar_push_kernel( + n: int, k: int, tactic: K3LowMTactic, use_pdl: bool, ar_push: int, ar_world: int, gated: bool +): + kernel = K3Fp8wLowMGemmKernel( + n, k, tactic, use_pdl, gated=gated, ar_push=ar_push, ar_world=ar_world + ) + dummy = 1 << 20 + w_ptr = make_ptr(cutlass.Float8E4M3FN, dummy, cute.AddressSpace.gmem, assumed_align=16) + x_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + ws_ptr = make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + addr = Int64(dummy) + if gated: + g_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + return cute.compile( + kernel.call_ar_push_gated, + w_ptr, + x_ptr, + g_ptr, + Int32(k), + ws_ptr, + Int32(1), + addr, + addr, + addr, + Int32(0), + stream, + ) + return cute.compile( + kernel.call_ar_push, w_ptr, x_ptr, ws_ptr, Int32(1), addr, addr, addr, Int32(0), stream + ) + + +def ar_push_tactic(m: int, n: int, k: int, gated: bool = False) -> Optional[K3LowMTactic]: + """The tactic the plain (or gated) low-M GEMM uses for this shape, if the all-reduce push + epilogue can run it (low-M kernel; ``split_k == 1``, or on SM107 also split-K); else None. + + The SM107 tables split K for most Kimi-K3 decode projections (212 SMs), so on Rubin the + split-K owners push their rows (:meth:`K3Fp8wLowMGemmKernel._push_owned_rows`) and the + push keeps the tactic, and with it the bits, of the plain GEMM. Whether a pushed split-K + GEMM beats the GEMM plus the C++ all-reduce at each shape and token count still has to be + measured on Rubin; SM100/SM103 keep the ``split_k == 1`` rule they were measured with. + """ + if not 1 <= m <= low_m_max_for(n, k) or n % 8: + return None + tactic = gated_tactic(m, n, k) if gated else default_tactic(m, n, k) + if tactic.split_k == 1: + return tactic + if _device_capability(torch.cuda.current_device()) != (10, 7): + return None + kt_max = _ceil_div(k // CTA_K, tactic.split_k) + if min(tactic.b_stages, kt_max) * tactic.mma_n * CTA_K < tactic.mma_n * MMA_M * 2: + return None + return tactic + + +def run_k3_fp8w_lowm_gemm_ar_push( + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + *, + gate: Optional[torch.Tensor] = None, + tactic: K3LowMTactic | K3MidMTactic | None = None, + use_pdl: bool = True, +) -> None: + """Low-M GEMM with the all-reduce push epilogue, on the current stream. + + Stores ``bf16(x @ dequant(w)^T)`` (``bf16(x * gate) @ ...`` with ``gate``) into the + Lamport slots of the next all-reduce on ``ws`` instead of an output tensor; + :func:`k3_lamport_ar.run_lamport_finish` with the same ``one_shot`` completes it. Same + tactic and rounding as :func:`run_k3_fp8w_lowm_gemm` (:func:`run_k3_fp8w_lowm_gemm_gated`), + so the pair is bitwise that GEMM followed by the MNNVL all-reduce of its output. + + c240: above the low-M range (no gate) the mid-M path runs, its quantization kernel then the + mid-M GEMM with the push epilogue (:func:`mid_ar_push_tactic`). + """ + m, k = x.shape + n = w.shape[0] + if tactic is None: + tactic = ar_push_tactic(m, n, k, gated=gate is not None) + if tactic is None and gate is None: + tactic = mid_ar_push_tactic(m, n, k) + if tactic is None: + raise ValueError(f"no all-reduce push tactic for M={m}, N={n}, K={k}") + if isinstance(tactic, K3MidMTactic): + if gate is not None: + raise ValueError("the mid-M all-reduce push has no gated variant") + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + xq, xs = quantize_mid_operand(x, tactic, use_pdl) + run_mid_prequantized_ar_push(xq, xs, w, w_scale, ws, one_shot, tactic, use_pdl) + return + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + need = k3_lamport_ar.required_bytes(m, n, ws.world, one_shot) + if need > ws.bytes_per_buffer: + raise ValueError(f"Lamport buffer of {ws.bytes_per_buffer} B < {need} B") + mode = AR_PUSH_ONE_SHOT if one_shot else AR_PUSH_TWO_SHOT + compiled = _compiled_ar_push_kernel(n, k, tactic, use_pdl, mode, ws.world, gate is not None) + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + ptrs = ( + make_ptr(cutlass.Float8E4M3FN, w.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + ) + ws_ptr = make_ptr(cutlass.Uint8, w_scale.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + tail = ( + Int32(m), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(ws.rank), + stream, + ) + if gate is not None: + g_ptr = make_ptr( + cutlass.BFloat16, gate.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + compiled(*ptrs, g_ptr, Int32(gate.stride(0)), ws_ptr, *tail) + else: + compiled(*ptrs, ws_ptr, *tail) + + +@functools.cache +def _compiled_split_kernel( + n: int, + k: int, + tactic: K3LowMTactic, + use_pdl: bool, + split_src: Tuple[int, int], + ar_push: int, + ar_world: int, + push_addend: bool = False, +): + kernel = K3Fp8wLowMGemmKernel( + n, + k, + tactic, + use_pdl, + ar_push=ar_push, + ar_world=ar_world, + split_src=split_src, + push_addend=push_addend, + ) + dummy = 1 << 20 + w_ptr = make_ptr(cutlass.Float8E4M3FN, dummy, cute.AddressSpace.gmem, assumed_align=16) + x_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + ws_ptr = make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + addr = Int64(dummy) + add_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + return cute.compile( + kernel.call_split_ar_push, + w_ptr, + x_ptr, + ws_ptr, + add_ptr, + Int32(1), + addr, + addr, + addr, + Int32(0), + stream, + ) + + +def split_source_tactic(m: int, n: int, k: int) -> Optional[K3LowMTactic]: + """The tactic of :func:`run_k3_fp8w_lowm_gemm_split_ar_push` for ``M`` tokens of an + ``[N, K]`` weight: the low-M kernel's, where its all-reduce push epilogue can run it + (:func:`ar_push_tactic`); else None.""" + return ar_push_tactic(m, n, k) + + +def _check_split_inputs(x, w, w_scale) -> Tuple[int, int, int]: + """``(m, n, k)`` of valid strided operands (see :func:`run_k3_fp8w_lowm_gemm_split_ar_push`); + raises ValueError otherwise.""" + for t, name, dtype in ((x, "x", torch.bfloat16), (w, "w", torch.float8_e4m3fn)): + if t.dtype != dtype or t.dim() != 2 or t.stride(1) != 1 or t.data_ptr() % 16: + raise ValueError(f"{name} must be a 16-byte aligned 2D {dtype} tensor, unit col stride") + m, k = x.shape + n = w.shape[0] + if tuple(w.shape) != (n, k): + raise ValueError("operand shapes do not match") + if k % CTA_K or n % N_ALIGN: + raise ValueError(f"K={k} must be a multiple of {CTA_K}, N={n} of {N_ALIGN}") + if x.stride(0) % 8 or w.stride(0) % 16: + raise ValueError("x / w row strides must be 16-byte multiples") + words = _ceil_div(k // CTA_K, 4) + if ( + w_scale.dtype != torch.int32 + or tuple(w_scale.shape) != (n, words) + or w_scale.stride(0) != 1 + or (words > 1 and w_scale.stride(1) != n) + ): + raise ValueError(f"w_scale must be DeepGEMM-packed UE8M0 int32 [{n}, {words}]") + return m, n, k + + +def run_k3_fp8w_lowm_gemm_split_ar_push( + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + *, + addend: Optional[torch.Tensor] = None, + tactic: K3LowMTactic | None = None, + use_pdl: bool = True, +) -> None: + """``x @ dequant(w, w_scale)^T`` with the all-reduce push epilogue of + :func:`run_k3_fp8w_lowm_gemm_ar_push`, low-M kernel only: the BF16 ``[M, N]`` result goes + into the Lamport slots of the next all-reduce on ``ws`` (``one_shot`` or two-shot transport). + + ``x`` (BF16 ``[M, K]``) and ``w`` (E4M3 ``[N, K]``) may be column slices of wider tensors + (unit column stride, 16-byte aligned rows), read in place; ``w_scale`` is the DeepGEMM-packed + UE8M0 scale of the ``[N, K]`` weight. Bitwise equal to :func:`run_k3_fp8w_lowm_gemm_ar_push` + on contiguous copies with the same tactic. With ``addend`` (a contiguous BF16 ``[M, N]`` + tensor, read after the kernel's ``griddepcontrol.wait``) the pushed values are + ``result + addend`` as the PyTorch BF16 add computes it.""" + m, n, k = _check_split_inputs(x, w, w_scale) + if addend is not None and ( + addend.dtype != torch.bfloat16 + or tuple(addend.shape) != (m, n) + or not addend.is_contiguous() + or addend.data_ptr() % 16 + ): + raise ValueError("addend must be a contiguous, 16-byte aligned bf16 [M, N] tensor") + if tactic is None: + tactic = split_source_tactic(m, n, k) + if tactic is None: + raise ValueError(f"no all-reduce push tactic for M={m}, N={n}, K={k}") + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + need = k3_lamport_ar.required_bytes(m, n, ws.world, one_shot) + if need > ws.bytes_per_buffer: + raise ValueError(f"Lamport buffer of {ws.bytes_per_buffer} B < {need} B") + mode = AR_PUSH_ONE_SHOT if one_shot else AR_PUSH_TWO_SHOT + split_src = (w.stride(0), x.stride(0)) + compiled = _compiled_split_kernel( + n, k, tactic, use_pdl, split_src, mode, ws.world, addend is not None + ) + compiled( + _gmem_ptr(cutlass.Float8E4M3FN, w), + _gmem_ptr(cutlass.BFloat16, x), + _gmem_ptr(cutlass.Uint8, w_scale), + _gmem_ptr(cutlass.BFloat16, x if addend is None else addend), + Int32(m), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(ws.rank), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + + +def pack_ue8m0_exponents(ue8m0: torch.Tensor, n: int, k: int) -> torch.Tensor: + """UE8M0 block exponents ``uint8 [ceil(N/128), K/128]`` -> the DeepGEMM-packed int32 scale + ``[N, ceil(K/512)]`` (strides ``(1, N)``) the low-M kernels read.""" + return _pack_ue8m0_exponents(ue8m0, n, k) + + +def block_exponents_of(scale: torch.Tensor, n: int, k: int) -> torch.Tensor: + """The UE8M0 exponents ``uint8 [ceil(N/128), K/128]`` of an ``[N, K]`` weight's 128x128 + blocks, from either layout ``_Fp8BlockScaleWeightReadLinear`` holds: the Rubin MXFP8 R128c4 + ``uint8`` scale or the DeepGEMM-packed ``int32 [N, ceil(K/512)]`` scale.""" + nb, kb = _ceil_div(n, MMA_M), k // CTA_K + if k % CTA_K: + raise ValueError(f"K={k} must be a multiple of {CTA_K}") + if scale.dtype == torch.uint8: + flat = scale.reshape(-1) + if flat.numel() != nb * kb * 512: + raise ValueError(f"not the R128c4 scale of an [{n}, {k}] weight") + return flat.view(nb, kb, 512)[:, :, 0].contiguous() + words = _ceil_div(kb, 4) + if scale.dtype != torch.int32 or tuple(scale.shape) != (n, words): + raise ValueError(f"not the packed int32 scale of an [{n}, {k}] weight") + rows = scale[::MMA_M].contiguous() # [nb, words], one row per 128-row block + return rows.view(torch.uint8).view(nb, words * 4)[:, :kb].contiguous() + + +def k3_fp8w_lowm_gemm_sigmoid( + x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + """``torch.sigmoid(k3_fp8w_lowm_gemm(x, w, w_scale))`` in one launch, bit for bit. + + Low-M kernel only (``M <= low_m_max_for(N, K)``, same tactic as the plain GEMM); the + epilogue applies :func:`sigmoid_bf16` to each BF16 output. ``out`` (contiguous BF16 + ``[M, N]``) is optional. + """ + _check_inputs(x, w, w_scale) + m, k = x.shape + n = w.shape[0] + if not 1 <= m <= low_m_max_for(n, k): + raise ValueError(f"M={m} must be in [1, {low_m_max_for(n, k)}] for the sigmoid GEMM") + if out is None: + out = torch.empty((m, n), dtype=torch.bfloat16, device=x.device) + elif out.dtype != torch.bfloat16 or tuple(out.shape) != (m, n) or not out.is_contiguous(): + raise ValueError("out must be a contiguous bf16 [M, N] tensor") + compiled = _compiled_kernel(n, k, default_tactic(m, n, k), True, EPI_BF16, sigmoid=True) + compiled( + make_ptr(cutlass.Float8E4M3FN, w.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, w_scale.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, out.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + *(_gmem_ptr(cutlass.Uint8, out) for _ in range(3)), + Int32(m), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + return out + + +def _check_inputs(x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor) -> None: + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous(): + raise ValueError("x must be a contiguous 2D bfloat16 tensor") + if w.dtype != torch.float8_e4m3fn or w.dim() != 2 or not w.is_contiguous(): + raise ValueError("w must be a contiguous 2D float8_e4m3fn tensor") + n, k = w.shape + if x.shape[1] != k: + raise ValueError(f"x has K={x.shape[1]}, weight has K={k}") + if n % N_ALIGN or k % CTA_K: + raise ValueError(f"N={n} must be a multiple of {N_ALIGN} and K={k} of {CTA_K}") + words = _ceil_div(k // CTA_K, 4) + if ( + w_scale.dtype != torch.int32 + or tuple(w_scale.shape) != (n, words) + or w_scale.stride(0) != 1 + or (words > 1 and w_scale.stride(1) != n) + ): + raise ValueError( + f"w_scale must be DeepGEMM-packed UE8M0 int32 [{n}, {words}] with strides (1, {n})" + ) + if x.data_ptr() % 16 or w.data_ptr() % 16: + raise ValueError("x and w must be 16-byte aligned") + + +def _check_gate(x: torch.Tensor, gate: torch.Tensor) -> None: + if gate.dtype != torch.bfloat16 or gate.dim() != 2 or tuple(gate.shape) != tuple(x.shape): + raise ValueError("gate must be a 2D bfloat16 tensor shaped like x") + if gate.stride(1) != 1 or gate.stride(0) % 8 or gate.stride(0) < gate.shape[1]: + raise ValueError("gate needs unit column stride and a row stride that is a multiple of 8") + if gate.data_ptr() % 16: + raise ValueError("gate must be 16-byte aligned") + + +@functools.cache +def _is_supported_device(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3), (10, 7)) + + +def can_run( + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + allow_partial_n_tile: bool = False, +) -> bool: + """Whether :func:`k3_fp8w_lowm_gemm` accepts these inputs (non-raising check). + + Call sites use it to pick this kernel over ``fp8_swap_ab_gemm``: only a 2D CUDA input + with ``1 <= M <= max_m_for(N, K)`` on SM100/SM103/SM107 qualifies, plus everything + :func:`_check_inputs` validates. A weight whose N is not a multiple of 128 (a partial + last row tile) qualifies only with ``allow_partial_n_tile``. + """ + if x.dim() != 2 or w.dim() != 2 or not x.is_cuda: + return False + if w.shape[0] % MMA_M and not allow_partial_n_tile: + return False + if not 1 <= x.shape[0] <= max_m_for(w.shape[0], w.shape[1]): + return False + if not _is_supported_device(x.device.index): + return False + try: + _check_inputs(x, w, w_scale) + except ValueError: + return False + return True + + +def can_run_gated( + x: torch.Tensor, + gate: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + allow_partial_n_tile: bool = False, +) -> bool: + """Whether :func:`k3_fp8w_lowm_gemm_gated` accepts these inputs (non-raising check): + :func:`can_run` with M served by the low-M kernel, a gated tactic that fits in shared + memory, and a valid ``gate``.""" + if not can_run(x, w, w_scale, allow_partial_n_tile): + return False + if x.shape[0] > low_m_max_for(w.shape[0], w.shape[1]): + return False + if not _gated_tactic_fits(x.shape[0], w.shape[0], w.shape[1]): + return False + try: + _check_gate(x, gate) + except ValueError: + return False + return True + + +def k3_fp8w_lowm_gemm_gated( + x: torch.Tensor, + gate: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """``bf16 [M, N] = bf16(x * gate) @ dequant(w, w_scale)^T`` on the low-M kernel. + + The output-gated projection ``proj(x * gate)`` in one launch, bitwise equal to the + multiply followed by :func:`k3_fp8w_lowm_gemm`. ``out`` (contiguous BF16 ``[M, N]``) is + optional. + """ + _check_inputs(x, w, w_scale) + _check_gate(x, gate) + if out is None: + out = torch.empty((x.shape[0], w.shape[0]), dtype=torch.bfloat16, device=x.device) + elif out.dtype != torch.bfloat16 or tuple(out.shape) != (x.shape[0], w.shape[0]): + raise ValueError("out must be bf16 [M, N]") + elif not out.is_contiguous(): + raise ValueError("out must be contiguous") + return run_k3_fp8w_lowm_gemm_gated(x, gate, w, w_scale, out) + + +def k3_fp8w_lowm_gemm(x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor) -> torch.Tensor: + """``bf16 [M, N] = x [M, K] @ dequant(w [N, K], w_scale)^T`` for ``1 <= M <= MAX_M``. + + Drop-in for ``fp8_swap_ab_gemm(x, w, w_scale, disable_ue8m0_cast=True)``: ``w`` is the + UE8M0-resmoothed FP8 weight and ``w_scale`` the DeepGEMM-packed UE8M0 scale. + """ + _check_inputs(x, w, w_scale) + out = torch.empty((x.shape[0], w.shape[0]), dtype=torch.bfloat16, device=x.device) + return run_k3_fp8w_lowm_gemm(x, w, w_scale, out) + + +def k3_fp8w_lowm_gemm_out( + x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + """Variant of :func:`k3_fp8w_lowm_gemm` writing into a preallocated ``out``.""" + _check_inputs(x, w, w_scale) + if out.dtype != torch.bfloat16 or tuple(out.shape) != (x.shape[0], w.shape[0]): + raise ValueError("out must be bf16 [M, N]") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + return run_k3_fp8w_lowm_gemm(x, w, w_scale, out) + + +__all__ = [ + "EPI_BF16", + "EPI_NVFP4", + "K3LowMTactic", + "K3MidMTactic", + "LOW_M_MAX", + "MAX_M", + "N_ALIGN", + "bf16_over_prequantized_max_m", + "block_exponents_of", + "can_run", + "can_run_gated", + "default_tactic", + "gated_tactic", + "k3_fp8w_lowm_gemm", + "k3_fp8w_lowm_gemm_gated", + "k3_fp8w_lowm_gemm_out", + "k3_fp8w_lowm_gemm_sigmoid", + "low_m_max_for", + "max_m_for", + "pack_block_scale_from_mxfp8_layout", + "pack_ue8m0_block_scale", + "pack_ue8m0_exponents", + "prepare_k3_fp8w_lowm_weight", + "run_k3_fp8w_lowm_gemm", + "run_k3_fp8w_lowm_gemm_fused", + "run_k3_fp8w_lowm_gemm_gated", + "run_k3_fp8w_lowm_gemm_split_ar_push", + "sigmoid_bf16", + "split_source_tactic", + "split_tactic", + "validate_tactic", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fp8w_lowm_mxfp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fp8w_lowm_mxfp8.py new file mode 100644 index 000000000000..2d0bb9488e5a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fp8w_lowm_mxfp8.py @@ -0,0 +1,1571 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M FP8-weight GEMM fed by an MXFP8 activation made upstream (SM100/SM103). + +``out[M, N] = x[M, K] @ (W[N, K] * S)^T`` like :mod:`k3_fp8w_lowm_gemm`'s low-M kernel, but the +producer of ``x`` hands over its MXFP8 form instead of BF16: + +* ``xq``: E4M3 ``[M, K]`` row-major, and +* ``xs``: UE8M0 ``[M, K / 32]`` row-major, one exponent per 1x32 group, + +computed from the BF16 ``x`` exactly like the low-M kernel's in-kernel quantization +(:func:`mxfp8_quantize_reference` documents the arithmetic; producers call the same PTX helpers +``_abs_max_bf16x32`` / ``_ue8m0_of_amax`` / ``_bf16x2_inv_scale`` / ``_quant_bf16x8_to_e4m3``). +With the low-M kernel's tiling (same ``mma_n`` and split-K), K-tile order, scale folding +(``SFB = e_act + e_w - 127``, ``SFA = 1``) and split-K reduction order, the output is bitwise the +low-M kernel's on the BF16 ``x``. + +Why: at decode batch sizes the low-M kernel's tail after its ``griddepcontrol.wait`` is +dominated by the activation path (BF16 tiles arrive by TMA one per ~0.1-0.2 us per quantization +group and take ~0.25 us each to quantize), not by the weights, which a programmatic head start +already streams into the ring. Here the E4M3 tiles go by TMA (128B swizzle) straight into the +UMMA B ring (every K tile of the CTA in flight at once) and one warp folds the weight block +exponents into the activation exponents. + +Warp roles: warp 0 streams the weight ring (and releases the dependent grid at once, before its +own wait: it only reads weights), warp 1 issues the MMAs, warp 2 loads the E4M3 tiles after +``griddepcontrol.wait``, warp 3 loads and folds the scales after it, warps 4-7 run the epilogue +(TMEM -> registers in 32-column blocks -> split-K partials pushed with ``st.async`` into the +owner's dead weight ring once every CTA of the cluster is done with its MMAs -> token-major +staging -> 16-byte BF16 stores). + +Tensor-memory weight staging (``K3LowMMxTactic.tmem_tiles`` = ``t > 0``, :func:`staged_mx_tactic`): +the MMA warp copies the first ``t`` weight K tiles of its slice from the ring into tensor memory +(``tcgen05.cp`` 128x256b, four per tile) as they land and releases their ring stages at once, so +the weight warp refills them with the next tiles before ``griddepcontrol.wait``. With ``t`` = +K tiles minus ring stages the CTA's whole weight slice is on chip when the activation lands. The +MMAs of those tiles read A from tensor memory (``kind::mxf8f6f4.block_scale`` with A in TMEM), +in the same K-tile and K-block order as the ring MMAs, so the output is bitwise the unstaged +kernel's. The CTA then allocates 256 or 512 tensor-memory columns instead of 128. +""" + +from __future__ import annotations + +import dataclasses +import functools + +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 +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass import Int32, Int64 +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +from . import k3_fp8w_lowm_gemm as lowm +from . import k3_lamport_ar +from .k3_fp8w_lowm_gemm import ( + _EVICT_FIRST, + _SF_TMEM_COLS, + _TMEM_COLS, + CTA_K, + MMA_M, + N_ALIGN, + SF_VEC, + SUPPORTED_MMA_N, + _abs_max_bf16x32, + _add_bf16x2, + _asm, + _bf16x2_inv_scale, + _ceil_div, + _f32_to_bf16_bits, + _griddepcontrol_launch_dependents, + _griddepcontrol_wait, + _iv, + _ldg_v4_b32, + _lds_v4_b32, + _lds_v4_f32, + _mbarrier_arrive_remote, + _quant_bf16x8_to_e4m3, + _s2t_partition_sf, + _st_async_v4_f32, + _stg_v4_b32, + _sts_u16, + _sts_v4_b32, + _ue8m0_of_amax, +) + +#: Largest M of this kernel (the low-M kernel's range). +MAX_M = lowm.LOW_M_MAX + +_TMA_WARP = 0 +_MMA_WARP = 1 +_X_WARP = 2 +_SF_WARP = 3 +_EPI_WARP0 = 4 +_THREADS = 256 +_EPI_BARRIER_ID = 3 + +__all__ = [ + "MAX_M", + "K3LowMMxTactic", + "can_run_mxfp8", + "default_mx_tactic", + "can_stage_weights", + "ring_holds_weights", + "weights_on_chip", + "staged_matches_mid_m", + "staged_mx_tactic", + "mxfp8_quantize_reference", + "run_k3_fp8w_lowm_gemm_mxfp8", + "run_k3_fp8w_lowm_gemm_mxfp8_ag_push", + "run_k3_fp8w_lowm_gemm_mxfp8_ar_push", + "run_k3_fp8w_lowm_gemm_mxfp8_slice_ar_push", +] + + +#: First tensor-memory column of the staged weight tiles (the accumulator and the scale-factor +#: slots use the columns below it). +_A_TMEM_COL0 = _TMEM_COLS +#: Tensor-memory columns of one staged 128 x 128 E4M3 weight tile (four codes per column). +_A_TMEM_TILE_COLS = CTA_K // 4 +#: Largest tensor-memory allocation of one CTA. +_TMEM_MAX_COLS = 512 +#: Most weight K tiles that fit in tensor memory next to the accumulator and scale factors. +MAX_TMEM_TILES = (_TMEM_MAX_COLS - _A_TMEM_COL0) // _A_TMEM_TILE_COLS +#: Upper 32 bits of the shared-memory matrix descriptor of a 128B-swizzled K-major weight stage +#: (stride between 8-row groups 1024 B, descriptor version 1, 128-byte swizzle), the source of +#: ``tcgen05.cp``; the lower word is the 16-byte start address with a leading offset of one. +_A_DESC_HI = 0x40004040 +_A_DESC_LO_LBO = 0x10000 + + +@dataclasses.dataclass(frozen=True) +class K3LowMMxTactic: + """``mma_n`` padded tokens, cluster ``split_k``, weight ring ``a_stages`` and activation + ring ``b_stages`` (``b_stages`` = the K tiles of a CTA keeps every activation load in + flight); ``tmem_tiles``: the first weight K tiles of the slice staged in tensor memory + (0: none).""" + + mma_n: int + split_k: int + a_stages: int + b_stages: int + tmem_tiles: int = 0 + + +def _mailbox_bytes(t: K3LowMMxTactic) -> int: + if t.split_k == 1: + return 0 + return (t.split_k - 1) * _ceil_div(t.mma_n // 4, t.split_k) * MMA_M * 16 + + +def smem_bytes(n: int, k: int, t: K3LowMMxTactic) -> int: + """Conservative host-side mirror of the kernel's shared-memory footprint.""" + kt_max = _ceil_div(k // CTA_K, t.split_k) + a_st = min(t.a_stages, kt_max) + b_st = min(t.b_stages, kt_max) + total = 8 * (2 * (a_st + b_st) + 5) + 4 + total = lowm._align_up(total, 1024) + max(a_st * MMA_M * CTA_K, _mailbox_bytes(t)) + total = lowm._align_up(total, 1024) + max(b_st * t.mma_n * CTA_K, t.mma_n * MMA_M * 2) + total = lowm._align_up(total, 128) + 512 # SFA (one stage, constant 1.0) + total = lowm._align_up(total, 128) + b_st * 512 + return total + 1024 + + +def validate_tactic(n: int, k: int, t: K3LowMMxTactic) -> None: + if t.mma_n not in SUPPORTED_MMA_N: + raise ValueError(f"mma_n must be one of {SUPPORTED_MMA_N}, got {t.mma_n}") + if n % N_ALIGN or k % CTA_K: + raise ValueError(f"N={n} must be a multiple of {N_ALIGN} and K={k} of {CTA_K}") + if not 1 <= t.split_k <= min(16, k // CTA_K): + raise ValueError(f"split_k={t.split_k} invalid for K={k}") + if t.a_stages < 1 or t.b_stages < 1: + raise ValueError("a_stages and b_stages must be >= 1") + # Every rank of the cluster owns at least k // CTA_K // split_k K tiles; the staged tiles + # must be K tiles of every rank. + if not 0 <= t.tmem_tiles <= min(MAX_TMEM_TILES, (k // CTA_K) // t.split_k): + raise ValueError(f"tmem_tiles={t.tmem_tiles} invalid for K={k}, split_k={t.split_k}") + need = smem_bytes(n, k, t) + if need > lowm._SMEM_CAPACITY: + raise ValueError(f"tactic {t} needs {need} B of SMEM > {lowm._SMEM_CAPACITY} B") + + +def tmem_alloc_cols(t: K3LowMMxTactic) -> int: + """Tensor-memory columns a CTA of tactic ``t`` allocates (a power of two).""" + need = _A_TMEM_COL0 + t.tmem_tiles * _A_TMEM_TILE_COLS + cols = _TMEM_COLS + while cols < need: + cols *= 2 + return cols + + +#: Weight-ring depth :func:`default_mx_tactic` keeps before it shortens the activation ring. +_MIN_A_STAGES = 6 + + +@functools.cache +def default_mx_tactic(m: int, n: int, k: int) -> K3LowMMxTactic: + """The low-M kernel's ``mma_n`` and split-K for ``(M, N, K)`` (same bits) and the deepest + weight ring that fits next to the activation ring. The activation ring holds every K tile + of the CTA unless that leaves fewer than ``_MIN_A_STAGES`` weight stages (long K slices at + ``mma_n >= 32``); then it is the deepest ring that does.""" + base = lowm.default_tactic(m, n, k) + kt_max = _ceil_div(k // CTA_K, base.split_k) + best = None + for b_stages in range(kt_max, 0, -1): + t = K3LowMMxTactic(mma_n=base.mma_n, split_k=base.split_k, a_stages=1, b_stages=b_stages) + fits = [ + a + for a in range(1, kt_max + 1) + if smem_bytes(n, k, dataclasses.replace(t, a_stages=a)) <= lowm._SMEM_CAPACITY + ] + if fits: + t = dataclasses.replace(t, a_stages=max(fits)) + best = best or t + if t.a_stages >= min(_MIN_A_STAGES, kt_max): + return t + if best is None: + raise ValueError(f"no MXFP8 tactic fits for M={m}, N={n}, K={k}") + return best + + +#: Shared memory of a staged CTA (:func:`staged_mx_tactic`). With the predecessor's CTAs at the +#: largest carve-out (``k3_attn_res`` ``entry_trigger``), a CTA of up to 225 KB lands on an SM +#: that runs a 1 KB CTA (measured on Rubin; 226 KB does not); 220 KB leaves room for an attn_res +#: CTA's own few KB. +STAGED_SMEM_BYTES = 220 * 1024 + + +@functools.cache +def staged_mx_tactic(m: int, n: int, k: int) -> K3LowMMxTactic: + """:func:`default_mx_tactic`'s ``mma_n``, split-K and activation ring (so the same bits) + with the weight ring cut to :data:`STAGED_SMEM_BYTES` of shared memory and as many weight K + tiles staged in tensor memory as the ring cannot hold: the CTA's whole weight slice is on + chip before ``griddepcontrol.wait`` where tensor memory allows.""" + base = default_mx_tactic(m, n, k) + kt_max = _ceil_div(k // CTA_K, base.split_k) + budget = min(STAGED_SMEM_BYTES, lowm._SMEM_CAPACITY) + a_stages = base.a_stages + while a_stages > 1 and smem_bytes(n, k, dataclasses.replace(base, a_stages=a_stages)) > budget: + a_stages -= 1 + tiles = min(max(kt_max - a_stages, 0), MAX_TMEM_TILES, (k // CTA_K) // base.split_k) + tactic = dataclasses.replace(base, a_stages=a_stages, tmem_tiles=tiles) + validate_tactic(n, k, tactic) + return tactic + + +@functools.cache +def can_stage_weights(m: int, n: int, k: int) -> bool: + """Whether :func:`staged_mx_tactic` has a tactic with staged tiles for ``(M, N, K)``.""" + try: + return staged_mx_tactic(m, n, k).tmem_tiles > 0 + except ValueError: + return False + + +@functools.cache +def ring_holds_weights(m: int, n: int, k: int) -> bool: + """Whether :func:`staged_mx_tactic` keeps the CTA's whole weight K slice in its shared-memory + ring (at most :data:`STAGED_SMEM_BYTES`), with no tiles in tensor memory (e.g. the KDA QKVG + projection at TP16: 8 K tiles per CTA of the 7-way split). The staged tactic only sizes the + ring: the same bits.""" + try: + tactic = staged_mx_tactic(m, n, k) + except ValueError: + return False + return tactic.tmem_tiles == 0 and tactic.a_stages >= _ceil_div(k // CTA_K, tactic.split_k) + + +def weights_on_chip(m: int, n: int, k: int) -> bool: + """Whether :func:`staged_mx_tactic` has the CTA's weight slice on chip before + ``griddepcontrol.wait`` (:func:`can_stage_weights` or :func:`ring_holds_weights`), so + that the GEMM may run as the entry-released dependent of a ``k3_attn_res`` + (``run_k3_fp8w_lowm_gemm_mxfp8(..., stage_weights=True)``).""" + return can_stage_weights(m, n, k) or ring_holds_weights(m, n, k) + + +@functools.cache +def staged_matches_mid_m(m: int, n: int, k: int) -> bool: + """Whether the staged kernel on ``M`` rows above the low-M cut-over (``low_m_max_for(N, K) + < M <= MAX_M``, whole 128-row weight tiles) computes the mid-M GEMM's bits: both run the + same split-K over the same K-tile order, the mid-M kernel on the same MXFP8 quantization of + the activation (checked bitwise in the tests for the KDA QKVG weight).""" + if not lowm.low_m_max_for(n, k) < m <= MAX_M or n % MMA_M: + return False + try: + mid = lowm.default_mid_tactic(m, n, k) + except ValueError: + return False + return can_stage_weights(m, n, k) and mid.split_k == staged_mx_tactic(m, n, k).split_k + + +@dsl_user_op +def _ldg_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fold_scales_x4(act4, w_exp, *, loc=None, ip=None): + """Four bytes ``clamp(e_act_i + w_exp - 127, 0, 254)`` of a packed word of four UE8M0 + activation exponents (``_sts_scale`` of the low-M kernel, four at a time).""" + return Int32( + _asm( + T.i32(), + [_iv(act4, loc=loc, ip=ip), _iv(w_exp, loc=loc, ip=ip)], + "{\n" + ".reg .s32 a, v, r;\n" + "mov.b32 r, 0;\n" + + "".join( + f"bfe.u32 a, $1, {8 * i}, 8;\n" + "add.s32 v, a, $2;\n" + "sub.s32 v, v, 127;\n" + "max.s32 v, v, 0;\n" + "min.s32 v, v, 254;\n" + f"bfi.b32 r, v, r, {8 * i}, 8;\n" + for i in range(4) + ) + + "mov.b32 $0, r;\n}", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sts_b32(addr, v, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(v, loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _tmem_cp_128x256b(taddr, smem_addr, *, loc=None, ip=None): + """``tcgen05.cp`` of 128 rows x 32 bytes (one K block of E4M3 codes) of the 128B-swizzled + K-major tile at shared address ``smem_addr`` into eight tensor-memory columns at + ``taddr`` (row r -> lane r).""" + _asm( + None, + [_iv(taddr, loc=loc, ip=ip), _iv(smem_addr, loc=loc, ip=ip)], + "{\n" + ".reg .b32 lo, hi;\n" + ".reg .b64 desc;\n" + "shr.u32 lo, $1, 4;\n" + "and.b32 lo, lo, 0x3FFF;\n" + f"or.b32 lo, lo, {_A_DESC_LO_LBO};\n" + f"mov.b32 hi, {_A_DESC_HI};\n" + "mov.b64 desc, {lo, hi};\n" + "tcgen05.cp.cta_group::1.128x256b [$0], desc;\n" + "}", + "r,r", + loc=loc, + ip=ip, + ) + + +class K3Fp8wLowMMxGemmKernel: + """Swap-AB MXFP8 GEMM on an upstream-quantized activation (``M <= 64``), cluster split-K. + + ``ar_push`` (``AR_PUSH_TWO_SHOT`` / ``AR_PUSH_ONE_SHOT`` of :mod:`k3_fp8w_lowm_gemm`, + ``ar_world`` ranks): instead of the BF16 store, every 16-byte output vector goes into the + Lamport slots of the next all-reduce on a ``k3_lamport_ar`` workspace (``push_slot`` / + ``push_v4``, as the low-M kernel's push epilogue). The dependent grid is then released from + the activation warp after its ``griddepcontrol.wait`` (the low-M push variant's order), since + the consumer reads the workspace's rotation state as soon as it is resident. + """ + + def __init__( + self, + n: int, + k: int, + tactic: K3LowMMxTactic, + use_pdl: bool, + ar_push: int = 0, + ar_world: int = 0, + strides: tuple[int, int, int] | None = None, + push_addend: bool = False, + ag_hidden: int = 0, + ): + validate_tactic(n, k, tactic) + if ar_push and ( + ar_push not in (lowm.AR_PUSH_TWO_SHOT, lowm.AR_PUSH_ONE_SHOT) or ar_world < 2 + ): + raise ValueError(f"invalid all-reduce push {ar_push} over {ar_world} ranks") + # ag_hidden > 0: the output is these ranks' columns of the call's all-gather rows of width + # ag_hidden (k3_lamport_ar "ag": after the call's M x ar_world slots, one-shot transport), + # starting at the column the launch passes in place of the rank. + if ag_hidden and (ar_push != lowm.AR_PUSH_ONE_SHOT or push_addend or ag_hidden < n): + raise ValueError("the all-gather push is a one-shot push without addend") + self.ag_hidden = ag_hidden + if push_addend and not ar_push: + raise ValueError("push_addend goes with the all-reduce push") + self.ar_push = ar_push + self.ar_world = ar_world + # Row strides (elements) of the weight, the E4M3 activation and its UE8M0 scales: a + # column slice of wider operands (the row-parallel MoE tail reads its latent_up columns + # and its slice of the latent all-reduce's MXFP8 output in place). + self.w_ld, self.x_ld, self.xs_ld = strides if strides is not None else (k, k, k // SF_VEC) + if ( + self.w_ld < k + or self.x_ld < k + or self.xs_ld < k // SF_VEC + or self.w_ld % 16 + or self.x_ld % 16 + ): + raise ValueError(f"invalid row strides {strides} for K={k}") + # Every pushed BF16 value y becomes bf16(y + a), a from a BF16 [M, N] addend read after + # griddepcontrol.wait (the PyTorch BF16 add; k3_fp8w_lowm_gemm's ``push_addend``). + self.push_addend = push_addend + self.n = n + self.k = k + self.mma_n = tactic.mma_n + self.split_k = tactic.split_k + self.kt_total = k // CTA_K + self.kt_max = _ceil_div(self.kt_total, self.split_k) + self.a_stages = min(tactic.a_stages, self.kt_max) + self.b_stages = min(tactic.b_stages, self.kt_max) + self.n_tiles = _ceil_div(n, MMA_M) + self.use_pdl = use_pdl + self.mma_tiler = (MMA_M, self.mma_n, CTA_K) + self.a_stage_bytes = MMA_M * CTA_K + self.b_stage_bytes = self.mma_n * CTA_K + self.chunks = self.mma_n // 4 + self.chunks_cap = _ceil_div(self.chunks, self.split_k) + self.a_region_bytes = max(self.a_stages * self.a_stage_bytes, _mailbox_bytes(tactic)) + self.out_stage_bytes = self.mma_n * MMA_M * 2 + self.groups = k // SF_VEC + self.blk = min(32, self.mma_n) + self.tmem_tiles = tactic.tmem_tiles + self.tmem_cols = tmem_alloc_cols(tactic) + + @cute.jit + def __call__( + self, + w_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ): + self._launch( + w_ptr, + xq_ptr, + xs_ptr, + ws_ptr, + out_ptr, + m, + stream, + Int64(0), + Int64(0), + Int64(0), + Int32(0), + ) + + @cute.jit + def call_ar_push( + self, + w_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + m: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + stream: cuda.CUstream, + ): + """All-reduce push launch (``ar_push``): the output goes into the Lamport slots of the + next call on the workspace (``k3_lamport_ar.LamportWorkspace`` addresses: peer pointer + table, multicast base, state flags; this is rank ``rank``). Nothing is stored through + ``out_ptr``.""" + self._launch( + w_ptr, xq_ptr, xs_ptr, ws_ptr, out_ptr, m, stream, peer_ptrs, mc_base, flags, rank + ) + + @cute.jit + def _launch( + self, + w_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + out_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ar_peers: Int64, + ar_mc: Int64, + ar_flags: Int64, + ar_rank: Int32, + ): + n, k = self.n, self.k + w_ld, x_ld = self.w_ld, self.x_ld + mW = cute.make_tensor(w_ptr, cute.make_layout((n, k, 1), stride=(w_ld, 1, n * w_ld))) + mXq = cute.make_tensor(xq_ptr, cute.make_layout((m, k, 1), stride=(x_ld, 1, m * x_ld))) + mWS = cute.make_tensor( + ws_ptr, + cute.make_layout( + (_ceil_div(n, MMA_M), (4, _ceil_div(k // CTA_K, 4))), + stride=(MMA_M * 4, (1, 4 * n)), + ), + ) + mOut = cute.make_tensor(out_ptr, cute.make_layout((m, n), stride=(n, 1))) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float8E8M0FNU, + SF_VEC, + tcgen05.CtaGroup.ONE, + (MMA_M, self.mma_n), + ) + a_layout = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.a_stages + ) + # The same MMA with A in tensor memory for the staged weight tiles (unused without them). + tiled_mma_ts = tiled_mma + a_tmem_layout = a_layout + if cutlass.const_expr(self.tmem_tiles > 0): + tiled_mma_ts = sm100_utils.make_blockscaled_trivial_tiled_mma( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float8E8M0FNU, + SF_VEC, + tcgen05.CtaGroup.ONE, + (MMA_M, self.mma_n), + tcgen05.OperandSource.TMEM, + ) + a_tmem_layout = sm100_utils.make_smem_layout_a( + tiled_mma_ts, self.mma_tiler, cutlass.Float8E4M3FN, 1 + ) + b_layout = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.b_stages + ) + sfa_layout = blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler, SF_VEC, 1) + sfb_layout = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, self.mma_tiler, SF_VEC, self.b_stages + ) + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) + ) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mW, + cute.slice_(a_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mXq, + cute.slice_(b_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + + @cute.struct + class SharedStorage: + a_full: cute.struct.MemRange[cutlass.Int64, self.a_stages] + a_empty: cute.struct.MemRange[cutlass.Int64, self.a_stages] + b_full: cute.struct.MemRange[cutlass.Int64, self.b_stages] + b_empty: cute.struct.MemRange[cutlass.Int64, self.b_stages] + acc_full: cutlass.Int64 + tmem_ready: cutlass.Int64 + tmem_free: cutlass.Int64 + mbox_full: cutlass.Int64 + peers_done: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # Weight ring; the split-K mailbox once every CTA of the cluster is done with it. + sA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E4M3FN, self.a_region_bytes], 1024 + ] + # Activation ring; stages the BF16 output tile token-major after the mainloop. + sB: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float8E4M3FN, max(cute.cosize(b_layout.outer), self.out_stage_bytes) + ], + 1024, + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfa_layout)], 128 + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfb_layout)], 128 + ] + + self.shared_storage = SharedStorage + + self.kernel( + tiled_mma, + tiled_mma_ts, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + xs_ptr, + mWS, + mOut, + m, + a_layout, + a_tmem_layout, + b_layout, + sfa_layout, + sfb_layout, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + ).launch( + grid=(self.n_tiles * self.split_k, 1, 1), + block=(_THREADS, 1, 1), + cluster=(self.split_k, 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_ts: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB: cute.Tensor, + xs_ptr: cute.Pointer, + mWS: cute.Tensor, + mOut: cute.Tensor, + m: Int32, + a_layout: cute.ComposedLayout, + a_tmem_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + ar_peers: Int64, + ar_mc: Int64, + ar_flags: Int64, + ar_rank: Int32, + ): + S = self.split_k + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + rank = Int32(0) + if cutlass.const_expr(S > 1): + rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_tile = bidx // S + + kt_base = self.kt_total // S + kt_rem = self.kt_total % S + kt_count = Int32(kt_base) + kt_start = rank * kt_base + if cutlass.const_expr(kt_rem > 0): + if rank < kt_rem: + kt_count = kt_count + 1 + kt_start = kt_start + rank + else: + kt_start = kt_start + kt_rem + + storage = utils.SmemAllocator().allocate(self.shared_storage) + a_full = storage.a_full.data_ptr() + a_empty = storage.a_empty.data_ptr() + b_full = storage.b_full.data_ptr() + b_empty = storage.b_empty.data_ptr() + + if warp_idx == _TMA_WARP: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(self.a_stages): + cute.arch.mbarrier_init(a_full + i, 1) + cute.arch.mbarrier_init(a_empty + i, 1) + for i in cutlass.range_constexpr(self.b_stages): + # the E4M3 tile (TMA, expect_tx) and its folded scales (warp 3) + cute.arch.mbarrier_init(b_full + i, 2) + cute.arch.mbarrier_init(b_empty + i, 1) + cute.arch.mbarrier_init(storage.acc_full, 1) + cute.arch.mbarrier_init(storage.tmem_ready, 1) + cute.arch.mbarrier_init(storage.tmem_free, 4) + cute.arch.mbarrier_init(storage.mbox_full, 1) + cute.arch.mbarrier_init(storage.peers_done, max(S - 1, 1)) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + if cutlass.const_expr(S > 1): + cute.arch.cluster_arrive_relaxed() + + if warp_idx == _TMA_WARP: + # Weight producer: never waits on the previous grid (weights are static), so the + # dependent grid may launch at once (with the push: after the activation warp's wait). + if cutlass.const_expr(self.use_pdl and not self.ar_push): + _griddepcontrol_launch_dependents() + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + gA = cute.local_tile( + mA, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + tCgA = tiled_mma.get_slice(0).partition_A(gA) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + for kt in cutlass.range(kt_count, unroll=1): + stage = kt % self.a_stages + if kt >= self.a_stages: + cute.arch.mbarrier_wait(a_empty + stage, ((kt // self.a_stages) & 1) ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(a_full + stage, self.a_stage_bytes) + cute.copy( + tma_atom_a, + tAgA[(None, n_tile, kt_start + kt, 0)], + tAsA[(None, stage)], + tma_bar_ptr=a_full + stage, + cache_policy=Int64(_EVICT_FIRST), + ) + elif warp_idx == _X_WARP: + # Activation producer: the E4M3 tiles of every K tile (128B swizzle = UMMA B layout). + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + if cutlass.const_expr(self.ar_push != 0): + _griddepcontrol_launch_dependents() + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + gB = cute.local_tile( + mB, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + tCgB = tiled_mma.get_slice(0).partition_B(gB) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + for kt in cutlass.range(kt_count, unroll=1): + stage = kt % self.b_stages + if kt >= self.b_stages: + cute.arch.mbarrier_wait(b_empty + stage, ((kt // self.b_stages) & 1) ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(b_full + stage, self.b_stage_bytes) + cute.copy( + tma_atom_b, + tBgB[(None, 0, kt_start + kt, 0)], + tBsB[(None, stage)], + tma_bar_ptr=b_full + stage, + ) + elif warp_idx == _SF_WARP: + self.scale_warp(storage, xs_ptr, mWS, m, n_tile, kt_start, kt_count) + elif warp_idx == _MMA_WARP: + self.mma_warp( + tiled_mma, tiled_mma_ts, storage, kt_count, a_layout, a_tmem_layout, b_layout, + sfa_layout, sfb_layout, + ) # fmt: skip + elif warp_idx >= _EPI_WARP0: + self.epilogue( + tiled_mma, storage, mOut, m, tidx - 32 * _EPI_WARP0, n_tile, rank, + ar_peers, ar_mc, ar_flags, ar_rank, + ) # fmt: skip + + @cute.jit + def scale_warp( + self, + storage, + xs_ptr: cute.Pointer, + mWS: cute.Tensor, + m: Int32, + n_tile: Int32, + kt_start: Int32, + kt_count: Int32, + ): + """SFB of every K tile: lane r (and r + 32) writes token row r's four folded exponents + ``clamp(e_act + e_w - 127, 0, 254)`` at ``(r % 32) * 16 + (r // 32) * 4``; rows past M + (zero codes) get ``e_act = 127``.""" + lane = cute.arch.lane_idx() + b_full = storage.b_full.data_ptr() + b_empty = storage.b_empty.data_ptr() + sf_base = storage.sSFB.data_ptr().toint() + # Weight block exponents of this CTA's K tiles (static: before the wait). + w_exps = [] + for j in cutlass.range_constexpr(_ceil_div(self.kt_max, 32)): + e = Int32(127) + if lane + 32 * j < kt_count: + e = Int32(mWS[n_tile, kt_start + lane + 32 * j]) + w_exps.append(e) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + rows_per_lane = _ceil_div(self.mma_n, 32) + # Issue every activation exponent load first (one word = the 4 groups of a K tile). + acts = [] + for i in cutlass.range_constexpr(rows_per_lane): + r = lane + 32 * i + row_words = [] + for kt in cutlass.range_constexpr(self.kt_max): + v = Int32(0x7F7F7F7F) + if (r < m) & (kt < kt_count) & (r < self.mma_n): + v = _ldg_b32(xs_ptr.toint() + Int64(r * self.xs_ld + (kt_start + kt) * 4)) + row_words.append(v) + acts.append(row_words) + if cutlass.const_expr(self.b_stages >= self.kt_max): + # Every K tile has its own stage: write all scales, fence once, release all tiles. + for kt in cutlass.range_constexpr(self.kt_max): + if kt < kt_count: + e_w = cute.arch.shuffle_sync(w_exps[kt // 32], kt % 32) + for i in cutlass.range_constexpr(rows_per_lane): + r = lane + 32 * i + if r < self.mma_n: + _sts_b32( + sf_base + kt * 512 + (r & 31) * 16 + (r >> 5) * 4, + _fold_scales_x4(acts[i][kt], e_w), + ) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + for j in cutlass.range_constexpr(_ceil_div(self.kt_max, 32)): + if lane + 32 * j < kt_count: + cute.arch.mbarrier_arrive(b_full + lane + 32 * j) + else: + for kt in cutlass.range_constexpr(self.kt_max): + if kt < kt_count: + stage = kt % self.b_stages + if cutlass.const_expr(kt >= self.b_stages): + cute.arch.mbarrier_wait(b_empty + stage, ((kt // self.b_stages) & 1) ^ 1) + e_w = cute.arch.shuffle_sync(w_exps[kt // 32], kt % 32) + for i in cutlass.range_constexpr(rows_per_lane): + r = lane + 32 * i + if r < self.mma_n: + _sts_b32( + sf_base + stage * 512 + (r & 31) * 16 + (r >> 5) * 4, + _fold_scales_x4(acts[i][kt], e_w), + ) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(b_full + stage) + + @cute.jit + def mma_warp( + self, + tiled_mma: cute.TiledMma, + tiled_mma_ts: cute.TiledMma, + storage, + kt_count: Int32, + a_layout: cute.ComposedLayout, + a_tmem_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + ): + a_full = storage.a_full.data_ptr() + a_empty = storage.a_empty.data_ptr() + b_full = storage.b_full.data_ptr() + b_empty = storage.b_empty.data_ptr() + cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding_buf) + cute.arch.relinquish_tmem_alloc_permit() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(storage.tmem_ready) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding_buf + ) + staged = self.tmem_tiles + if cutlass.const_expr(staged > 0): + # Weight K tiles 0 .. staged - 1 into tensor memory as they land (tile j at columns + # _A_TMEM_COL0 + 32 j, one tcgen05.cp per 32-byte K block); their ring stages are + # released at once, so the weight warp refills them before griddepcontrol.wait. + # tcgen05.commit tracks the copies like MMAs, and the MMAs below that read the tiles + # are issued after them by the same thread. + tmem_base = storage.tmem_holding_buf.load() + sa_base = storage.sA.data_ptr().toint() + for j in cutlass.range_constexpr(staged): + stage = j % self.a_stages + cute.arch.mbarrier_wait(a_full + stage, (j // self.a_stages) & 1) + with cute.arch.elect_one(): + for kb in cutlass.range_constexpr(CTA_K // SF_VEC): + _tmem_cp_128x256b( + tmem_base + _A_TMEM_COL0 + j * _A_TMEM_TILE_COLS + kb * 8, + sa_base + stage * self.a_stage_bytes + kb * SF_VEC, + ) + if cutlass.const_expr(j + self.a_stages < self.kt_max): + tcgen05.commit(a_empty + stage) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tiled_mma.make_fragment_C(acc_shape).layout) + tCtSFA = cute.make_tensor( + cute.recast_ptr(acc_tmem_ptr + self.mma_n, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, self.mma_tiler, SF_VEC, cute.slice_(sfa_layout, (None, None, None, 0)) + ), + ) + s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = _s2t_partition_sf( + storage.sSFA.get_tensor(sfa_layout), tCtSFA + ) + # SFB: TMEM columns [mma_n + 16, +16). + sSFB = storage.sSFB.get_tensor(sfb_layout) + tmem_sfb = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, self.mma_tiler, SF_VEC, cute.slice_(sfb_layout, (None, None, None, 0)) + ) + tCtSFB = cute.make_tensor( + cute.recast_ptr( + acc_tmem_ptr + self.mma_n + _SF_TMEM_COLS, + dtype=cutlass.Float8E8M0FNU, + ), + tmem_sfb, + ) + s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = _s2t_partition_sf(sSFB, tCtSFB) + # The weight scale is folded into SFB, so SFA is a constant 1.0 (UE8M0 127). + lane = cute.arch.lane_idx() + one4 = Int32(0x7F7F7F7F) + _sts_v4_b32(storage.sSFA.data_ptr().toint() + lane * 16, one4, one4, one4, one4) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + cute.copy(s2t_sfa, tCsSFA_s2t[(None, None, None, None, 0)], tCtSFA_s2t) + + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + num_kphases = cute.size(tCrA, mode=[2]) + # A operands of the staged tiles: tile j's fragment at its tensor-memory columns. + tCtA_tiles = [] + if cutlass.const_expr(staged > 0): + tCtA_frag = tiled_mma_ts.make_fragment_A( + cute.make_tensor(acc_tmem_ptr, a_tmem_layout.outer) + )[None, None, None, 0] + for j in cutlass.range_constexpr(staged): + tCtA_tiles.append( + cute.make_tensor( + cute.recast_ptr( + acc_tmem_ptr + _A_TMEM_COL0 + j * _A_TMEM_TILE_COLS, + dtype=tCtA_frag.element_type, + ), + tCtA_frag.layout, + ) + ) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, staged > 0) + tiled_mma_ts.set(tcgen05.Field.ACCUMULATE, False) + # Fully unrolled over the largest K-tile count; the MMA descriptor fields are set outside + # the per-tile predicate (a rank with one K tile fewer skips the last tile's work). The + # staged tiles are K tiles of every rank (``validate_tactic``). + for kt in cutlass.range_constexpr(self.kt_max): + a_stage = kt % self.a_stages + b_stage = kt % self.b_stages + live = kt < kt_count + if live: + if cutlass.const_expr(kt >= staged): + cute.arch.mbarrier_wait(a_full + a_stage, (kt // self.a_stages) & 1) + cute.arch.mbarrier_wait(b_full + b_stage, (kt // self.b_stages) & 1) + cute.copy(s2t_sfb, tCsSFB_s2t[(None, None, None, None, b_stage)], tCtSFB_s2t) + for kp in cutlass.range_constexpr(num_kphases): + if cutlass.const_expr(kt < staged): + tiled_mma_ts.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator) + tiled_mma_ts.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator) + if live: + cute.gemm( + tiled_mma_ts, + tCtAcc, + tCtA_tiles[kt][(None, None, kp)], + tCrB[(None, None, kp, b_stage)], + tCtAcc, + ) + tiled_mma_ts.set(tcgen05.Field.ACCUMULATE, True) + else: + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator) + if live: + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[(None, None, kp, a_stage)], + tCrB[(None, None, kp, b_stage)], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if live: + with cute.arch.elect_one(): + if cutlass.const_expr(kt >= staged and kt + self.a_stages < self.kt_max): + tcgen05.commit(a_empty + a_stage) + if cutlass.const_expr(kt + self.b_stages < self.kt_max): + tcgen05.commit(b_empty + b_stage) + with cute.arch.elect_one(): + tcgen05.commit(storage.acc_full) + + cute.arch.mbarrier_wait(storage.tmem_free, 0) + cute.arch.dealloc_tmem(acc_tmem_ptr, self.tmem_cols) + + @cute.jit + def epilogue( + self, + tiled_mma: cute.TiledMma, + storage, + mOut: cute.Tensor, + m: Int32, + q: Int32, + n_tile: Int32, + rank: Int32, + ar_peers: Int64, + ar_mc: Int64, + ar_flags: Int64, + ar_rank: Int32, + ): + """TMEM -> registers -> token-major SMEM tile -> 16-byte stores (``ar_push``: into the + all-reduce's Lamport slots); with ``split_k > 1`` rank ``r`` reduces every 4-token chunk + ``c`` with ``c % split_k == r`` (own partial, then the peers' in rank order: the low-M + kernel's order).""" + S = self.split_k + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + # The previous call on the workspace has completed: its rotation state is final. + ar_cur, ar_nbytes = Int32(0), Int32(0) + if cutlass.const_expr(self.ar_push != 0): + ar_cur, _, ar_nbytes, _, _, _, _, _ = k3_lamport_ar.lamport_state(ar_flags) + cute.arch.mbarrier_wait(storage.tmem_ready, 0) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tiled_mma.make_fragment_C(acc_shape).layout) + tAcc = tCtAcc[((None, None), 0, 0)] + cAcc = cute.make_identity_tensor((MMA_M, self.mma_n)) + blk = self.blk + tiled_t2r = tcgen05.make_tmem_copy( + cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(blk)), cutlass.Float32), + cute.local_tile(tAcc, (MMA_M, blk), (0, 0)), + ) + thr_t2r = tiled_t2r.get_slice(q) + out_base = storage.sB.data_ptr().toint() + out_gmem = mOut.iterator.toint() + Int64(n_tile * MMA_M * 2) + row_stride = Int64(self.n * 2) + store_cols = Int32(16) + if cutlass.const_expr(self.n % MMA_M != 0): + if n_tile * MMA_M + (q & 15) * 8 >= self.n: + store_cols = Int32(0) + + cute.arch.mbarrier_wait(storage.acc_full, 0) + if cutlass.const_expr(S > 1): + # Every CTA's weight ring is dead once its MMAs completed: it becomes the mailbox. + cute.arch.cluster_wait() + if q == 0: + for p in cutlass.range_constexpr(1, S): + _mbarrier_arrive_remote(storage.peers_done.toint(), (rank + p) % S) + cute.arch.mbarrier_wait(storage.peers_done, 0) + + mbox_base = storage.sA.data_ptr().toint() + mbar_addr = storage.mbox_full.toint() + slot_bytes = self.chunks_cap * MMA_M * 16 + rOwn = cute.make_rmem_tensor((self.chunks_cap * 4,), cutlass.Float32) + nblk = self.mma_n // blk + tTR_cAcc0 = thr_t2r.partition_D(cute.local_tile(cAcc, (MMA_M, blk), (0, 0))) + rBuf = [cute.make_rmem_tensor(tTR_cAcc0.shape, cutlass.Float32) for _ in range(2)] + cute.copy( + tiled_t2r, thr_t2r.partition_S(cute.local_tile(tAcc, (MMA_M, blk), (0, 0))), rBuf[0] + ) + for j in cutlass.range_constexpr(nblk): + cute.arch.fence_view_async_tmem_load() + rAcc = rBuf[j % 2] + if cutlass.const_expr(j + 1 < nblk): + cute.copy( + tiled_t2r, + thr_t2r.partition_S(cute.local_tile(tAcc, (MMA_M, blk), (0, j + 1))), + rBuf[(j + 1) % 2], + ) + for c4 in cutlass.range_constexpr(blk // 4): + c = j * (blk // 4) + c4 + if cutlass.const_expr(S == 1): + if 4 * c < m: + for v in cutlass.range_constexpr(4): + _sts_u16( + out_base + (4 * c + v) * (MMA_M * 2) + q * 2, + _f32_to_bf16_bits(rAcc[4 * c4 + v]), + ) + else: + if rank == c % S: + for v in cutlass.range_constexpr(4): + rOwn[(c // S) * 4 + v] = rAcc[4 * c4 + v] + else: + if 4 * c < m: + slot = rank + if rank > c % S: + slot = rank - 1 + _st_async_v4_f32( + mbox_base + slot * slot_bytes + ((c // S) * MMA_M + q) * 16, + rAcc[4 * c4], + rAcc[4 * c4 + 1], + rAcc[4 * c4 + 2], + rAcc[4 * c4 + 3], + mbar_addr, + c % S, + ) + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(storage.tmem_free) + + if cutlass.const_expr(S == 1): + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + for p in cutlass.range(q, m * store_cols, MMA_M, unroll=1): + t = p >> 4 + col = (p & 15) * 16 + w0, w1, w2, w3 = _lds_v4_b32(out_base + t * (MMA_M * 2) + col) + if cutlass.const_expr(self.ar_push != 0): + self._push( + (w0, w1, w2, w3), t, n_tile * MMA_M + (p & 15) * 8, + ar_peers, ar_mc, ar_cur, ar_nbytes, ar_rank, mOut.iterator.toint(), m, + ) # fmt: skip + else: + _stg_v4_b32(out_gmem + Int64(t) * row_stride + col, w0, w1, w2, w3) + else: + active = (m + 3) >> 2 + my_chunks = Int32(0) + if active > rank: + my_chunks = (active - rank + S - 1) // S + if my_chunks > 0: + if q == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + storage.mbox_full, my_chunks * ((S - 1) * MMA_M * 16) + ) + cute.arch.mbarrier_wait(storage.mbox_full, 0) + for i in cutlass.range_constexpr(self.chunks_cap): + if i < my_chunks: + acc = [rOwn[i * 4 + v] for v in range(4)] + for p in cutlass.range_constexpr(S - 1): + v0, v1, v2, v3 = _lds_v4_f32( + mbox_base + p * slot_bytes + (i * MMA_M + q) * 16 + ) + acc[0] = acc[0] + v0 + acc[1] = acc[1] + v1 + acc[2] = acc[2] + v2 + acc[3] = acc[3] + v3 + for v in cutlass.range_constexpr(4): + _sts_u16( + out_base + (4 * (rank + i * S) + v) * (MMA_M * 2) + q * 2, + _f32_to_bf16_bits(acc[v]), + ) + cute.arch.barrier(barrier_id=_EPI_BARRIER_ID, number_of_threads=MMA_M) + for p in cutlass.range(q, my_chunks * 4 * store_cols, MMA_M, unroll=1): + r = p >> 4 + t = 4 * (rank + (r >> 2) * S) + (r & 3) + if t < m: + col = (p & 15) * 16 + w0, w1, w2, w3 = _lds_v4_b32(out_base + t * (MMA_M * 2) + col) + if cutlass.const_expr(self.ar_push != 0): + self._push( + (w0, w1, w2, w3), t, n_tile * MMA_M + (p & 15) * 8, + ar_peers, ar_mc, ar_cur, ar_nbytes, ar_rank, mOut.iterator.toint(), + m, + ) # fmt: skip + else: + _stg_v4_b32(out_gmem + Int64(t) * row_stride + col, w0, w1, w2, w3) + + @cute.jit + def _push( + self, + w, + t: Int32, + elem: Int32, + peer_ptrs: Int64, + mc_base: Int64, + cur: Int32, + nbytes: Int32, + ar_rank: Int32, + add_base: Int64, + m: Int32, + ): + """One 16-byte output vector (features ``[elem, elem + 8)`` of token ``t``) into its + Lamport slot of the next all-reduce (``-0`` sent as ``+0``); with ``push_addend`` after + the BF16 add of the addend at ``add_base`` (BF16 ``[M, N]``, read after the wait). With + ``ag_hidden`` into columns ``ar_rank + elem`` (the launch's column offset) of the call's + all-gather row ``t`` instead.""" + if cutlass.const_expr(self.ag_hidden > 0): + dst = ( + k3_lamport_ar.ag_row_addr(mc_base, cur, nbytes, m, t, self.ar_world, self.ag_hidden) + + Int64(ar_rank + elem) * 2 + ) + k3_lamport_ar.push_v4(dst, w[0], w[1], w[2], w[3], True) + else: + if cutlass.const_expr(self.push_addend): + a = _ldg_v4_b32(add_base + (Int64(t) * self.n + Int64(elem)) * 2) + w = tuple(_add_bf16x2(w[i], a[i]) for i in range(4)) + one_shot = self.ar_push == lowm.AR_PUSH_ONE_SHOT + dst = k3_lamport_ar.push_slot( + peer_ptrs, mc_base, cur, nbytes, t, ar_rank, elem, self.ar_world, self.n, one_shot + ) + k3_lamport_ar.push_v4(dst, w[0], w[1], w[2], w[3], one_shot) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +@functools.cache +def _compiled_kernel(n: int, k: int, tactic: K3LowMMxTactic, use_pdl: bool): + kernel = K3Fp8wLowMMxGemmKernel(n, k, tactic, use_pdl) + dummy = 1 << 20 + ptrs = [ + make_ptr(dtype, dummy, cute.AddressSpace.gmem, assumed_align=16) + for dtype in ( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + cutlass.Uint8, + cutlass.Uint8, + cutlass.BFloat16, + ) + ] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile(kernel, *ptrs, Int32(1), stream) + + +@functools.cache +def _compiled_ar_push_kernel( + n: int, + k: int, + tactic: K3LowMMxTactic, + use_pdl: bool, + ar_push: int, + ar_world: int, + strides: tuple[int, int, int] | None = None, + push_addend: bool = False, + ag_hidden: int = 0, +): + kernel = K3Fp8wLowMMxGemmKernel( + n, + k, + tactic, + use_pdl, + ar_push=ar_push, + ar_world=ar_world, + strides=strides, + push_addend=push_addend, + ag_hidden=ag_hidden, + ) + dummy = 1 << 20 + # A column slice of the scales starts at a 4-byte boundary (K / 32 of whole 128-blocks). + xs_align = 4 if strides is not None else 16 + ptrs = [ + make_ptr(dtype, dummy, cute.AddressSpace.gmem, assumed_align=align) + for dtype, align in ( + (cutlass.Float8E4M3FN, 16), + (cutlass.Float8E4M3FN, 16), + (cutlass.Uint8, xs_align), + (cutlass.Uint8, 16), + (cutlass.BFloat16, 16), + ) + ] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + addr = Int64(dummy) + return cute.compile(kernel.call_ar_push, *ptrs, Int32(1), addr, addr, addr, Int32(0), stream) + + +def _gmem_ptr(dtype, tensor: torch.Tensor): + return make_ptr(dtype, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def can_run_mxfp8( + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + allow_partial_n_tile: bool = False, + above_cutover: bool = False, +) -> bool: + """Whether :func:`run_k3_fp8w_lowm_gemm_mxfp8` accepts these operands: the low-M kernel's + weight contract (UE8M0-resmoothed E4M3 ``w [N, K]``, DeepGEMM-packed ``w_scale``, N a + multiple of 128 unless ``allow_partial_n_tile``), ``1 <= M <= low_m_max_for(N, K)`` on + SM100/SM103 (``above_cutover``: up to :data:`MAX_M`), and ``xq`` E4M3 ``[M, K]`` / ``xs`` + uint8 ``[M, K / 32]``, contiguous, 16- and 4-byte aligned.""" + if xq.dim() != 2 or xs.dim() != 2 or w.dim() != 2 or not xq.is_cuda: + return False + if not lowm._is_supported_device(xq.device.index): + return False + m, k = xq.shape + n = w.shape[0] + if w.shape[1] != k or n % N_ALIGN or k % CTA_K or (n % MMA_M and not allow_partial_n_tile): + return False + if not 1 <= m <= (MAX_M if above_cutover else min(MAX_M, lowm.low_m_max_for(n, k))): + return False + if xq.dtype != torch.float8_e4m3fn or xs.dtype != torch.uint8: + return False + if tuple(xs.shape) != (m, k // SF_VEC) or not xq.is_contiguous() or not xs.is_contiguous(): + return False + if w.dtype != torch.float8_e4m3fn or not w.is_contiguous(): + return False + try: + validate_tactic(n, k, default_mx_tactic(m, n, k)) + except ValueError: + return False + words = _ceil_div(k // CTA_K, 4) + if ( + w_scale.dtype != torch.int32 + or tuple(w_scale.shape) != (n, words) + or w_scale.stride(0) != 1 + or (words > 1 and w_scale.stride(1) != n) + ): + return False + return xq.data_ptr() % 16 == 0 and w.data_ptr() % 16 == 0 and xs.data_ptr() % 4 == 0 + + +def run_k3_fp8w_lowm_gemm_mxfp8( + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + *, + tactic: K3LowMMxTactic | None = None, + use_pdl: bool = True, + stage_weights: bool = False, +) -> torch.Tensor: + """``out = dequant(xq, xs) @ dequant(w, w_scale)^T`` into ``out`` (bf16 ``[M, N]``) on the + current stream, bitwise the low-M kernel on the BF16 activation that ``(xq, xs)`` quantizes + (inputs are not re-validated; see :func:`can_run_mxfp8`). ``stage_weights`` (no ``tactic``): + :func:`staged_mx_tactic`, the same bits.""" + m, k = xq.shape + n = w.shape[0] + if tactic is None: + tactic = staged_mx_tactic(m, n, k) if stage_weights else default_mx_tactic(m, n, k) + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + compiled = _compiled_kernel(n, k, tactic, use_pdl) + compiled( + _gmem_ptr(cutlass.Float8E4M3FN, w), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + _gmem_ptr(cutlass.Uint8, w_scale), + _gmem_ptr(cutlass.BFloat16, out), + Int32(m), + cuda.CUstream(torch.cuda.current_stream(xq.device).cuda_stream), + ) + return out + + +def run_k3_fp8w_lowm_gemm_mxfp8_ar_push( + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + *, + tactic: K3LowMMxTactic | None = None, + use_pdl: bool = True, +) -> None: + """:func:`run_k3_fp8w_lowm_gemm_mxfp8` with the all-reduce push epilogue: stores the BF16 + output into the Lamport slots of the next all-reduce on ``ws``; + ``k3_lamport_ar.run_lamport_finish`` (or a consumer's fused finish) with the same + ``one_shot`` completes it. Same tactic and rounding as the plain MXFP8 GEMM, so the pair is + bitwise that GEMM (and the gated low-M GEMM whose quantization ``(xq, xs)`` reproduces) + followed by the MNNVL all-reduce of its output.""" + m, k = xq.shape + n = w.shape[0] + if tactic is None: + tactic = default_mx_tactic(m, n, k) + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + need = k3_lamport_ar.required_bytes(m, n, ws.world, one_shot) + if need > ws.bytes_per_buffer: + raise ValueError(f"Lamport buffer of {ws.bytes_per_buffer} B < {need} B") + mode = lowm.AR_PUSH_ONE_SHOT if one_shot else lowm.AR_PUSH_TWO_SHOT + compiled = _compiled_ar_push_kernel(n, k, tactic, use_pdl, mode, ws.world) + compiled( + _gmem_ptr(cutlass.Float8E4M3FN, w), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + _gmem_ptr(cutlass.Uint8, w_scale), + make_ptr(cutlass.BFloat16, xq.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int32(m), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(ws.rank), + cuda.CUstream(torch.cuda.current_stream(xq.device).cuda_stream), + ) + + +def run_k3_fp8w_lowm_gemm_mxfp8_slice_ar_push( + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + k0: int, + k1: int, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + *, + addend: torch.Tensor | None = None, + tactic: K3LowMMxTactic | None = None, + use_pdl: bool = True, +) -> None: + """:func:`run_k3_fp8w_lowm_gemm_mxfp8_ar_push` on the K columns ``[k0, k0 + k1)`` of wider + operands, read in place: E4M3 ``xq [M, K]`` and UE8M0 ``xs [M, K / 32]`` (e.g. the MXFP8 + output of ``k3_latent_ar_norm``) and the E4M3 weight ``w [N, K]``; ``w_scale`` is the + DeepGEMM-packed scale of the ``[N, k1]`` slice (``k3_fp8w_lowm_gemm.pack_ue8m0_exponents``). + With ``addend`` (contiguous BF16 ``[M, N]``, read after the kernel's wait) every pushed value + is ``bf16(y + addend)``. Same tactic, K-tile order and rounding as + ``k3_fp8w_lowm_gemm.run_k3_fp8w_lowm_gemm_split_ar_push`` on the BF16 rows that ``(xq, xs)`` + quantize, so the pushed bits are that GEMM's (the row-parallel MoE tail).""" + m = xq.shape[0] + n = w.shape[0] + if k0 % CTA_K or k1 % CTA_K or k0 + k1 > xq.shape[1] or xq.shape[1] != w.shape[1]: + raise ValueError(f"invalid K slice [{k0}, {k0 + k1}) of K={xq.shape[1]}") + if xs.shape != (m, xq.shape[1] // SF_VEC) or xs.stride(1) != 1 or xq.stride(1) != 1: + raise ValueError("xs must be the [M, K / 32] scales of xq, rows contiguous") + if addend is not None and ( + addend.dtype != torch.bfloat16 + or tuple(addend.shape) != (m, n) + or not addend.is_contiguous() + or addend.data_ptr() % 16 + ): + raise ValueError("addend must be a contiguous, 16-byte aligned bf16 [M, N] tensor") + if tactic is None: + tactic = default_mx_tactic(m, n, k1) + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + need = k3_lamport_ar.required_bytes(m, n, ws.world, one_shot) + if need > ws.bytes_per_buffer: + raise ValueError(f"Lamport buffer of {ws.bytes_per_buffer} B < {need} B") + mode = lowm.AR_PUSH_ONE_SHOT if one_shot else lowm.AR_PUSH_TWO_SHOT + strides = (w.stride(0), xq.stride(0), xs.stride(0)) + compiled = _compiled_ar_push_kernel( + n, k1, tactic, use_pdl, mode, ws.world, strides, addend is not None + ) + + def ptr(dtype, t, off): + return make_ptr(dtype, t.data_ptr() + off, cute.AddressSpace.gmem, assumed_align=16) + + args = [ + ptr(cutlass.Float8E4M3FN, w, k0), + ptr(cutlass.Float8E4M3FN, xq, k0), + make_ptr( + cutlass.Uint8, xs.data_ptr() + k0 // SF_VEC, cute.AddressSpace.gmem, assumed_align=4 + ), + _gmem_ptr(cutlass.Uint8, w_scale), + _gmem_ptr(cutlass.BFloat16, addend if addend is not None else xq), + Int32(m), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(ws.rank), + ] + compiled(*args, cuda.CUstream(torch.cuda.current_stream(xq.device).cuda_stream)) + + +def run_k3_fp8w_lowm_gemm_mxfp8_ag_push( + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + col0: int, + hidden: int, + *, + tactic: K3LowMMxTactic | None = None, + use_pdl: bool = True, +) -> None: + """:func:`run_k3_fp8w_lowm_gemm_mxfp8` whose BF16 output (``[M, N]``, ``N`` = rows of ``w``) + goes into columns ``[col0, col0 + N)`` of the all-gather rows (width ``hidden``) of the next + one-shot call on ``ws`` (``k3_lamport_ar`` "ag": every rank pushes its columns, the call's + finish adds them to the reduced slots). ``tactic``: e.g. the tactic of the full ``[hidden, K]`` + projection whose output columns ``w`` holds, so the pushed values are bitwise that + projection's (every 128-row output tile is computed alone, with the same split-K).""" + m, k = xq.shape + n = w.shape[0] + if tactic is None: + tactic = default_mx_tactic(m, n, k) + if m > tactic.mma_n: + raise ValueError(f"M={m} exceeds tactic.mma_n={tactic.mma_n}") + if col0 % 8 or col0 + n > hidden: + raise ValueError(f"invalid column range [{col0}, {col0 + n}) of {hidden}") + need = k3_lamport_ar.required_bytes(m, hidden, ws.world, True, ag=True) + if need > ws.bytes_per_buffer: + raise ValueError(f"Lamport buffer of {ws.bytes_per_buffer} B < {need} B") + compiled = _compiled_ar_push_kernel( + n, k, tactic, use_pdl, lowm.AR_PUSH_ONE_SHOT, ws.world, ag_hidden=hidden + ) + args = [ + _gmem_ptr(cutlass.Float8E4M3FN, w), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + _gmem_ptr(cutlass.Uint8, w_scale), + make_ptr(cutlass.BFloat16, xq.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int32(m), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(col0), + ] + compiled(*args, cuda.CUstream(torch.cuda.current_stream(xq.device).cuda_stream)) + + +def mxfp8_quantize_reference(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(xq, xs)`` of a BF16 ``[M, K]`` activation as the low-M kernel quantizes it (torch + reference for tests): per 1x32 group, ``e = biased exponent of 2^ceil(log2(amax / 448))`` + with ``amax`` floored at ~1e-4 (BF16 ``0x38D2``), codes ``satfinite_rn_e4m3(x * 2^(127 - e))`` + (the power-of-two multiply is exact in BF16).""" + m, k = x.shape + g = x.view(m, k // SF_VEC, SF_VEC) + bits = g.abs().view(torch.int16).to(torch.int32) & 0x7FFF + amax = bits.amax(dim=-1).clamp_min(lowm._AMAX_FLOOR_BF16) + e = (amax >> 7) - 8 + ((amax & 0x7F) > 0x60).to(torch.int32) + inv = torch.pow(2.0, (127 - e).float()).to(torch.bfloat16) + scaled = (g * inv.unsqueeze(-1)).float() + xq = scaled.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).view(m, k) + return xq, e.to(torch.uint8) + + +@dsl_user_op +def _amax_bf16x8_words(w0, w1, w2, w3, *, loc=None, ip=None): + """|x| max of eight BF16 values (four packed words) as a packed BF16 pair (garbage sign).""" + return Int32( + _asm( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)], + "{\n.reg .b32 m0, m1;\n" + "max.xorsign.abs.bf16x2 m0, $1, $2;\n" + "max.xorsign.abs.bf16x2 m1, $3, $4;\n" + "max.xorsign.abs.bf16x2 $0, m0, m1;\n}", + "=r,r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _amax_pair(a, b, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +# Producer-side helpers (bitwise the low-M kernel's in-kernel quantization). +amax_bf16x8_words = _amax_bf16x8_words +amax_pair = _amax_pair +ue8m0_of_amax = _ue8m0_of_amax +bf16x2_inv_scale = _bf16x2_inv_scale +quant_bf16x8_to_e4m3 = _quant_bf16x8_to_e4m3 +abs_max_bf16x32 = _abs_max_bf16x32 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fused_router.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fused_router.py new file mode 100644 index 000000000000..de229be86a80 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_fused_router.py @@ -0,0 +1,336 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared pieces of the Kimi-K3 MoE routing kernels (SM100/SM103/SM107). + +The router shape (896 experts, hidden size 7168, top 16), the ``mma.sync`` router GEMM tiling +(eight groups of 112 experts, 64-wide K chunks), inline PTX helpers and the ``noaux_tc_op`` +weight renormalization that :mod:`k3_route_tail`, :mod:`k3_router_exact` and :mod:`k3_ull_front` +build on, and :func:`sort_buffer_sizes`. The renormalization uses FP32 +division, which equals the reference's FP64 expression whenever the scale is 1 and the sum is +>= 2^-12, and the FP64 expression otherwise. +""" + +from __future__ import annotations + +import struct + +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Kimi-K3 routed experts. +NUM_EXPERTS = 896 +#: Kimi-K3 hidden size. +HIDDEN_SIZE = 7168 +#: Experts selected per token. +TOP_K = 16 + +_K_CHUNK = 64 +_NUM_K_CHUNKS = HIDDEN_SIZE // _K_CHUNK +_EXPERT_GROUPS = 8 +_EXPERTS_PER_CTA = NUM_EXPERTS // _EXPERT_GROUPS +_M_TILES = _EXPERTS_PER_CTA // 16 +_W_CHUNK_BYTES = _EXPERTS_PER_CTA * _K_CHUNK * 2 +#: L2 eviction hint for the streamed weight (CacheHintSm90::EVICT_FIRST). +_EVICT_FIRST = 0x12F0000000000000 +#: ``1e-20`` as a PTX double literal (the renormalization epsilon of ``noaux_tc_op``). +_EPS_F64 = "0d" + struct.pack(">d", 1e-20).hex().upper() + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _lv(value, *, loc=None, ip=None): + return Int64(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return Float32(value).ir_value(loc=loc, ip=ip) + + +def _unpack(res, ty, cls, n, *, loc=None, ip=None): + return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n)) + + +@dsl_user_op +def _ldsm_x4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _mma_bf16(a, b0, b1, c, *, loc=None, ip=None): + """``d[16x8] = a[16x16] @ b[16x8] + c`` (BF16 inputs, FP32 accumulate).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(x, loc=loc, ip=ip) for x in (*a, b0, b1)] + [_fv(x, loc=loc, ip=ip) for x in c], + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _st_global_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "st.global.f32 [$0], $1;", + "l,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _atom_add_relaxed_gpu(addr, value, *, loc=None, ip=None): + """Relaxed atomic add; callers order it with explicit fences.""" + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "atom.relaxed.gpu.global.add.u32 $0, [$1], $2;", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.relaxed.gpu.global.b32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.b32 $0, [$1];", "=r,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _lds_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.f32 $0, [$1];", "=f,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _sts_v4_f32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)] + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", + "r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _warp_inclusive_sum(value, *, loc=None, ip=None): + body = "".join( + f"shfl.sync.up.b32 t|p, x, {off}, 0, 0xffffffff;\n@p add.s32 x, x, t;\n" + for off in (1, 2, 4, 8, 16) + ) + return Int32( + _asm( + T.i32(), + [_iv(value, loc=loc, ip=ip)], + "{\n.reg .b32 x, t;\n.reg .pred p;\nmov.b32 x, $1;\n" + body + "mov.b32 $0, x;\n}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _red_shared_or(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "red.shared.or.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _warp_sum_f32(value, *, loc=None, ip=None): + """Butterfly warp sum, same order as ``cooperative_groups::reduce(plus)``.""" + body = "".join( + f"shfl.sync.bfly.b32 t, x, {off}, 0x1f, 0xffffffff;\nadd.rn.f32 x, x, t;\n" + for off in (16, 8, 4, 2, 1) + ) + return Float32( + _asm( + T.f32(), + [_fv(value, loc=loc, ip=ip)], + "{\n.reg .f32 x, t;\nmov.b32 x, $1;\n" + body + "mov.b32 $0, x;\n}", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _renormalize(score, total, scale_bits, *, loc=None, ip=None): + """``float(double(score) * scale / (double(total) + 1e-20))`` like ``noaux_tc_op``.""" + return Float32( + _asm( + T.f32(), + [ + _fv(score, loc=loc, ip=ip), + _fv(total, loc=loc, ip=ip), + _lv(scale_bits, loc=loc, ip=ip), + ], + "{\n" + ".reg .f64 a, b, s;\n" + "mov.b64 s, $3;\n" + "cvt.f64.f32 a, $1;\n" + "mul.rn.f64 a, a, s;\n" + "cvt.f64.f32 b, $2;\n" + f"add.rn.f64 b, b, {_EPS_F64};\n" + "div.rn.f64 a, a, b;\n" + "cvt.rn.f32.f64 $0, a;\n" + "}", + "=f,f,f,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _unit_scale_fast_path(scale_bits, total, *, loc=None, ip=None): + """1 if ``scale == 1.0`` and ``total >= 2^-12`` (see :func:`_renormalize`).""" + return Int32( + _asm( + T.i32(), + [_lv(scale_bits, loc=loc, ip=ip), _fv(total, loc=loc, ip=ip)], + "{\n.reg .pred p, q;\n" + "setp.eq.u64 p, $1, 0x3FF0000000000000;\n" + "setp.ge.f32 q, $2, 0f39800000;\n" + "and.pred p, p, q;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,l,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _div_rn_f32(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _popc(value, *, loc=None, ip=None): + return Int32( + _asm(T.i32(), [_iv(value, loc=loc, ip=ip)], "popc.b32 $0, $1;", "=r,r", loc=loc, ip=ip) + ) + + +@dsl_user_op +def _fence_acq_rel_gpu(*, loc=None, ip=None): + _asm(None, [], "fence.acq_rel.gpu;", "", loc=loc, ip=ip) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def sort_buffer_sizes(num_tokens: int, local_num_experts: int, tile_tokens_dim: int = 128): + """``(max_num_tiles, max_num_permuted_tokens)`` exactly as ``moe_sort`` allocates them.""" + expanded = num_tokens * TOP_K + filled = min(local_num_experts, expanded) + max_tiles = filled + (expanded - filled) // tile_tokens_dim + return max_tiles, max_tiles * tile_tokens_dim + + +__all__ = [ + "HIDDEN_SIZE", + "NUM_EXPERTS", + "TOP_K", + "sort_buffer_sizes", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_helix_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_helix_tail.py new file mode 100644 index 000000000000..f9e876be60fd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_helix_tail.py @@ -0,0 +1,1818 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 Helix generation tail in one kernel: split-KV fold, V up-projection, exchange, combine. + +Under Helix every CP rank decodes all ``H`` MLA heads over its ``1 / P`` slice of the KV cache, and +each rank then needs, for its ``H / P`` heads, the partial outputs and softmax statistics of all +``P`` ranks. With the V up-projection before the exchange, one +launch of :class:`K3HelixTailKernel` computes what the packed split-KV reduction +(``fp4_mla_packed_reduce.reduce_packed_splits``), the per-head V up-projection BMM, the Helix +all-to-all and ``helix_post_process_native`` compute one after another. Every CTA (16 warps) owns +one head and ``tiles_per_cta`` tiles of 16 tokens; it first sends, then receives: + +1. ``W_UV`` of its head goes into shared memory with bulk copies issued before + ``griddepcontrol.wait`` (it does not depend on the previous grid). +2. Fold. Warp ``w`` builds 256-column halves ``(row w // 2, half w % 2)`` and ``(row 8 + w // 2, + half w % 2)`` of the tile's latent rows in shared memory. A row whose request ran several + decode tasks is folded from the FP32 partials with the statistics and partials of all its tasks + (up to :data:`FOLD_TASKS`) loaded in one round, then the arithmetic of the reduction kernel's + ``_reduce_packed_row`` in its order (beyond :data:`FOLD_TASKS` tasks, ``_reduce_packed_row`` + itself): the BF16 row and its ``(max, sum)`` statistics are bitwise the reduction's. A row the + decode stored directly (single-task request, or no split partials at all) is copied with its + statistics. With ``ull2`` (``features.ull2_mla_tail``, and always for the normalized FP16 + partials of ``features.fp4_mla_f16_partials``) the requests' first tasks (plan ``begins``) are + staged in shared memory before ``griddepcontrol.wait``, rows with more tasks are folded in + rounds of :data:`FOLD_TASKS` by the same arithmetic, and the latent rows are padded by 16 bytes + so that step 3's operand loads are free of bank conflicts; the bits are the same. +3. V up-projection. ``mma.sync`` BF16 with FP32 accumulation: warp ``w`` computes output columns + ``8 w .. 8 w + 7`` of the 16 rows over K = 512. K is permuted inside every 32-column block so + that each thread's A and B fragments of two consecutive K16 steps are one 16-byte shared load; + the per-column dot products are unchanged, only their FP32 summation order differs from + cuBLAS'. +4. Push. The BF16 outputs (staged token-major in shared memory) and the statistics go with + 16-byte (8-byte) stores straight into the head owner's Lamport receive buffer, record + ``(token, this rank)``, in the record layout of ``helix_a2a_fused``. Rows of tokens this rank + owns no KV for (``zero_kv_mask``) are sent as the no-op contribution (zeros, ``(-inf, 0)``), + and no word equal to the sentinel ``0xffffffff`` is sent. +5. Receive. CTA ``b`` combines tokens ``b, b + CTAs, ..`` with the receive loop of + ``helix_a2a_fused`` (passes of 512 columns): poll the token's ``P`` records, combine them + bitwise like ``helix_post_process_native`` (same ``expf``, reciprocal, butterflies and FMA + order), store the BF16 ``[T, H / P * 128]`` output (optionally also the MXFP8 o_proj operand) + and re-arm the words read. A CTA only polls after its own pushes, and the grid is at most one + CTA per SM, so every sender is scheduled whatever the receivers wait for. Each CTA releases the + dependent grid (o_proj) between its pushes and its receive loop. + +Receive buffers, halves and the launch ticket are those of ``helix_a2a_fused`` (two halves used by +alternate launches; a word that is not the sentinel has arrived; a launch re-arms what it read +before it exits, and the next store into that half comes two launches later from a rank that has +finished the launch in between). Launches that share a workspace must be ordered on one stream on +every rank, in the same sequence on every rank. Every CTA takes its ticket after +``griddepcontrol.wait``, when every earlier launch on the stream has completed, so tickets of +consecutive launches never interleave; the ticket's latency overlaps the first fold. + +Numerics: the fold (latent rows and statistics) and the combine are bitwise those of the separate +kernels; the V up-projection rounds the same FP32 dot products to BF16 after another summation +order than the cuBLAS BMM, so the result agrees with the four-kernel path to BF16 rounding. +""" + +from __future__ import annotations + +import functools +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +from .fp4_mla_mufu16 import LOG2_E +from .fp4_mla_packed_reduce import _reduce_packed_row +from .helix_a2a_fused import ( + _NEG_INF_BITS, + MAX_CP, + SENTINEL, + _all_valid2, + _all_valid4, + _as_f32, + _asm, + _atom_add_u64, + _bf16x2_fma, + _corrected_exp, + _fadd, + _fmax, + _fmul, + _fv, + _griddepcontrol_launch_dependents, + _griddepcontrol_wait, + _iv, + _ld_nc_u8, + _ld_nc_v4, + _ld_vol_v2, + _ld_vol_v4, + _lds_b64, + _lv, + _pack_bf16x2, + _prep4, + _rcp, + _shfl_bfly, + _shfl_idx, + _st_sentinel_v2, + _st_sentinel_v4, + _st_v4, + _sts_b64, + _sts_v4, + _unpack, + record_bytes, +) +from .k3_fp8w_lowm_gemm import ( + _bf16x2_inv_scale, + _gate_mul_bf16x8, + _quant_bf16x8_to_e4m3, + _ue8m0_of_amax, +) +from .k3_fp8w_lowm_mxfp8 import amax_bf16x8_words, amax_pair + +__all__ = [ + "LATENT_DIM", + "ROWS", + "SENTINEL", + "V_HEAD_DIM", + "FoldInputs", + "can_run", + "k3_helix_tail", + "record_bytes", +] + +#: Latent (kv_lora_rank) and value (v_head_dim) widths of a head. +LATENT_DIM = 512 +V_HEAD_DIM = 128 +#: Tokens per tile (the M of one ``mma.sync`` tile). +ROWS = 16 +#: Threads per CTA: 16 warps; a warp folds two 256-column row halves and owns one n8 tile. +THREADS = 512 +_WARPS = THREADS // 32 +#: Decode tasks of a row whose loads are issued in one round; rows with more use +#: ``_reduce_packed_row`` (rounds of this many tasks with ``ull2``). +FOLD_TASKS = 8 +#: Bytes of one head's ``W_UV`` (``[V_HEAD_DIM, LATENT_DIM]`` BF16) and of one latent row. +_W_HEAD_BYTES = V_HEAD_DIM * LATENT_DIM * 2 +_LAT_ROW_BYTES = LATENT_DIM * 2 +#: Row pitch (elements) of the folded latent rows in shared memory: 16 bytes of padding, so the +#: V up-projection's A fragments (rows g and g + 8 of eight lane groups at the same column) fall +#: in eight different 16-byte bank groups instead of one. +_LAT_PITCH = LATENT_DIM + 8 +#: Request entries of the task plan (``begins``) staged in shared memory before the grid wait. +_PLAN_SMEM = 1024 +#: Row stride of ``W_UV`` in shared memory (64 bytes of padding against bank conflicts). +_W_ROW_BYTES = _LAT_ROW_BYTES + 64 +#: Warps whose lane 0 issues an equal share of ``W_UV``'s row copies: the copies of one warp +#: issue one at a time, and four issuing warps land the 128 KB about 1.5 us sooner than one. +_W_ISSUE_WARPS = 4 +#: Row stride of the token-major BF16 output staging (128 values + 16 bytes against bank conflicts). +_C_ROW_BYTES = V_HEAD_DIM * 2 + 16 +#: Poll rounds after which a debug launch traps instead of waiting for a record. +_DEBUG_POLL_LIMIT = 1 << 24 + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +# --------------------------------------------------------------------------- +# Inline PTX helpers (the rest come from helix_a2a_fused) +# --------------------------------------------------------------------------- + + +@dsl_user_op +def _lds_v4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _lds_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.b32 $0, [$1];", "=r,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _sts_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_b16(addr, value, *, loc=None, ip=None): + """Store the low 16 bits of ``value`` to shared memory.""" + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.u16.u32 h, $1;\nst.shared.b16 [$0], h;\n}", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.f32 $0, [$1];", "=f,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _ld_nc_v2(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.v2.b32 {$0, $1}, [$2];", + "=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _st_v2(addr, v0, v1, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(v0, loc=loc, ip=ip), _iv(v1, loc=loc, ip=ip)], + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mbar_init(addr, count: int, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)], + f"mbarrier.init.shared::cta.b64 [$0], {count};\nfence.mbarrier_init.release.cluster;", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mbar_expect_tx(addr, nbytes: int, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)], + f"mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], {nbytes};", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_g2s(dst, src, nbytes: int, mbar, *, loc=None, ip=None): + """Bulk copy of ``nbytes`` from global ``src`` to shared ``dst``, completing on ``mbar``.""" + _asm( + None, + [_iv(dst, loc=loc, ip=ip), _lv(src, loc=loc, ip=ip), _iv(mbar, loc=loc, ip=ip)], + f"cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], {nbytes}, [$2];", + "r,l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mbar_try_wait(addr, phase, *, loc=None, ip=None): + """1 once phase ``phase`` of the shared mbarrier at ``addr`` has completed, else 0.""" + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip), _iv(phase, loc=loc, ip=ip)], + "{\n.reg .pred p;\nmbarrier.try_wait.parity.shared::cta.b64 p, [$1], $2;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _opaque_f32(value, *, loc=None, ip=None): + """``value`` through an opaque move: the compiler cannot merge expressions on it with + expressions on ``value``.""" + return Float32( + _asm(T.f32(), [_fv(value, loc=loc, ip=ip)], "mov.b32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def _mma_bf16(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """``d[16x8] = a[16x16] @ b[16x8] + c`` (BF16 inputs, FP32 accumulation).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(x, loc=loc, ip=ip) for x in (a0, a1, a2, a3, b0, b1)] + + [_fv(x, loc=loc, ip=ip) for x in (c0, c1, c2, c3)], + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _shfl_bfly_b32(value, offset: int, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(value, loc=loc, ip=ip)], + f"shfl.sync.bfly.b32 $0, $1, {offset}, 31, -1;", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_u8(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u8 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _globaltimer(*, loc=None, ip=None): + return Int64(_asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc=loc, ip=ip)) + + +@dsl_user_op +def _st_u64(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "st.global.u64 [$0], $1;", + "l,l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _trap(*, loc=None, ip=None): + _asm(None, [], "trap;", "", loc=loc, ip=ip) + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +@cute.jit +def _fold_half( + mPartial: cute.Tensor, + mRowMax: cute.Tensor, + mRowSum: cute.Tensor, + sLat: cute.Tensor, + sMax: cute.Tensor, + sSum: cute.Tensor, + lane: Int32, + row: Int32, + local_row: Int32, + half: Int32, + task_begin: Int32, + request_tasks: Int32, + physical_m: Int32, + scale: Float32, + output_scale: Float32, +) -> None: + """``_reduce_packed_row`` of one 256-column half (``band = half`` of 256 columns) of a row + with at most :data:`FOLD_TASKS` tasks, with the statistics and partials of every task loaded + before any is consumed (the reduction issues them in three rounds). The weights, the + denominator and the numerators are accumulated in the reduction's task order with the same + operations, so the BF16 row and its statistics are bitwise the reduction's.""" + band_columns: cutlass.Constexpr = 256 + columns_per_lane: cutlass.Constexpr = 4 + band_base = half * Int32(band_columns) + lane * Int32(columns_per_lane) + gPartial_arr = cutlass.make_array_view(mPartial) + gOut_arr = cutlass.make_array_view(sLat) + sums = [] + maxima = [] + vecs = [] + for k in cutlass.range_constexpr(FOLD_TASKS): + # Out-of-range slots re-read the last valid task with zero weight (as the reduction). + task = task_begin + cutlass.min(Int32(k), request_tasks - Int32(1)) + sums.append(mRowSum[task, local_row]) + maxima.append(mRowMax[task, local_row]) + partial_row = (task * physical_m + local_row) * Int32(LATENT_DIM) + band_base + for group in cutlass.range_constexpr(2): + vecs.append( + gPartial_arr.subview(partial_row + Int32(group * 128)) + .data_ptr() + .load(count=columns_per_lane, alignment=4 * columns_per_lane) + ) + global_row_max = cutlass.Float32(-cutlass.Float32.inf) + for k in cutlass.range_constexpr(FOLD_TASKS): + if Int32(k) < request_tasks: + if sums[k] != cutlass.Float32(0.0): + global_row_max = cute.arch.fmax(global_row_max, maxima[k] * scale) + numerator0 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + numerator1 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + denominator = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(FOLD_TASKS): + task_weight = cutlass.Float32(0.0) + if Int32(k) < request_tasks: + task_row_sum = sums[k] + if task_row_sum != cutlass.Float32(0.0): + # A product of its own, as the reduction's: its contraction with the subtraction + # must not depend on the global-max product of the same task. + task_max = _opaque_f32(maxima[k]) + task_weight = cute.exp2(task_max * scale - global_row_max, fastmath=True) + denominator += task_row_sum * task_weight + # Empty tasks keep zero weight; their payload is never trusted. + if task_weight != cutlass.Float32(0.0): + numerator0 = numerator0 + vecs[2 * k] * task_weight + numerator1 = numerator1 + vecs[2 * k + 1] * task_weight + normalizer = cutlass.Float32(0.0) + if denominator != cutlass.Float32(0.0): + normalizer = output_scale / denominator + out_row = row * Int32(LATENT_DIM) + band_base + gOut_arr.store( + (numerator0 * normalizer).to(sLat.element_type), + idx=out_row, + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + gOut_arr.store( + (numerator1 * normalizer).to(sLat.element_type), + idx=out_row + Int32(128), + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + if half == Int32(0) and lane == Int32(0): + # Helix combine convention: natural-log-domain scaled max and the sum relative to it. + stats_max = cutlass.Float32(-cutlass.Float32.inf) + if denominator != cutlass.Float32(0.0): + stats_max = global_row_max * cutlass.Float32(1.0 / LOG2_E) + sMax[row] = stats_max + sSum[row] = denominator + + +@cute.jit +def _fold_round_loads( + mPartial: cute.Tensor, + mRowMax: cute.Tensor, + mRowSum: cute.Tensor, + lane: Int32, + local_row: Int32, + half: Int32, + task_begin: Int32, + request_tasks: Int32, + physical_m: Int32, +): + """The row sums, row maxima and this lane's partial vectors (two groups of 4 columns of the + 256-column ``half``) of the first :data:`FOLD_TASKS` tasks of a row; slots past + ``request_tasks`` re-read the last task (zero weight later, as the reduction).""" + columns_per_lane: cutlass.Constexpr = 4 + partial_bits: cutlass.Constexpr = mPartial.element_type.width + band_base = half * Int32(256) + lane * Int32(columns_per_lane) + gPartial_arr = cutlass.make_array_view(mPartial) + sums = [] + maxima = [] + vecs = [] + for k in cutlass.range_constexpr(FOLD_TASKS): + task = task_begin + cutlass.min(Int32(k), request_tasks - Int32(1)) + sums.append(mRowSum[task, local_row]) + maxima.append(mRowMax[task, local_row]) + partial_row = (task * physical_m + local_row) * Int32(LATENT_DIM) + band_base + for group in cutlass.range_constexpr(2): + vecs.append( + gPartial_arr.subview(partial_row + Int32(group * 128)) + .data_ptr() + .load(count=columns_per_lane, alignment=partial_bits // 8 * columns_per_lane) + ) + return sums, maxima, vecs + + +@cute.jit +def _fold_round_sums(loaded, request_tasks: Int32, scale: Float32, normalized: cutlass.Constexpr): + """The reduction's global maximum, denominator and numerators of a row with at most + :data:`FOLD_TASKS` tasks from :func:`_fold_round_loads`' values, in its task order.""" + sums, maxima, vecs = loaded + columns_per_lane: cutlass.Constexpr = 4 + global_row_max = cutlass.Float32(-cutlass.Float32.inf) + numerator0 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + numerator1 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + denominator = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(FOLD_TASKS): + if Int32(k) < request_tasks: + if sums[k] != cutlass.Float32(0.0): + global_row_max = cute.arch.fmax(global_row_max, maxima[k] * scale) + for k in cutlass.range_constexpr(FOLD_TASKS): + task_weight = cutlass.Float32(0.0) + if Int32(k) < request_tasks: + task_row_sum = sums[k] + if task_row_sum != cutlass.Float32(0.0): + # A product of its own, as the reduction's: its contraction with the subtraction + # must not depend on the global-max product of the same task. + task_max = _opaque_f32(maxima[k]) + task_weight = cute.exp2(task_max * scale - global_row_max, fastmath=True) + denominator += task_row_sum * task_weight + if cutlass.const_expr(normalized): + task_weight = task_weight * task_row_sum + # Empty tasks keep zero weight; their payload is never trusted. + if task_weight != cutlass.Float32(0.0): + v0 = vecs[2 * k] + v1 = vecs[2 * k + 1] + if cutlass.const_expr(normalized): + v0 = v0.to(cutlass.Float32) + v1 = v1.to(cutlass.Float32) + numerator0 = numerator0 + v0 * task_weight + numerator1 = numerator1 + v1 * task_weight + return global_row_max, denominator, numerator0, numerator1 + + +@cute.jit +def _fold_store( + sLat: cute.Tensor, + sMax: cute.Tensor, + sSum: cute.Tensor, + lane: Int32, + row: Int32, + half: Int32, + global_row_max: Float32, + denominator: Float32, + numerator0, + numerator1, + output_scale: Float32, + pitch: cutlass.Constexpr, +) -> None: + """The folded BF16 row half and (for half 0) its Helix statistics into shared memory.""" + columns_per_lane: cutlass.Constexpr = 4 + gOut_arr = cutlass.make_array_view(sLat) + normalizer = cutlass.Float32(0.0) + if denominator != cutlass.Float32(0.0): + normalizer = output_scale / denominator + out_row = row * Int32(pitch) + half * Int32(256) + lane * Int32(columns_per_lane) + gOut_arr.store( + (numerator0 * normalizer).to(sLat.element_type), + idx=out_row, + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + gOut_arr.store( + (numerator1 * normalizer).to(sLat.element_type), + idx=out_row + Int32(128), + vector_size=columns_per_lane, + alignment=2 * columns_per_lane, + ) + if half == Int32(0) and lane == Int32(0): + # Helix combine convention: natural-log-domain scaled max and the sum relative to it. + stats_max = cutlass.Float32(-cutlass.Float32.inf) + if denominator != cutlass.Float32(0.0): + stats_max = global_row_max * cutlass.Float32(1.0 / LOG2_E) + sMax[row] = stats_max + sSum[row] = denominator + + +@cute.jit +def _fold_half_ull2( + mPartial: cute.Tensor, + mRowMax: cute.Tensor, + mRowSum: cute.Tensor, + sLat: cute.Tensor, + sMax: cute.Tensor, + sSum: cute.Tensor, + lane: Int32, + row: Int32, + local_row: Int32, + half: Int32, + task_begin: Int32, + request_tasks: Int32, + physical_m: Int32, + scale: Float32, + output_scale: Float32, + pitch: cutlass.Constexpr = _LAT_PITCH, +) -> None: + """``ull2``: ``_reduce_packed_row`` of one 256-column half (``band = half`` of 256 columns) + of a row, into row ``row`` (rows ``pitch`` elements apart) of the shared-memory latent tile, + for any task count. With at most :data:`FOLD_TASKS` tasks the statistics and partials of + every task are loaded before any is consumed (the reduction issues them in three rounds); + with more, the row maxima are scanned first and the tasks are folded :data:`FOLD_TASKS` at a + time. The weights, the + denominator and the numerators are accumulated in the reduction's task order with the same + operations, so the BF16 row and its statistics are bitwise the reduction's. FP16 partials + (``features.fp4_mla_f16_partials``) are normalized split outputs, weighted by the max-scaled + row sums, as in the reduction.""" + band_columns: cutlass.Constexpr = 256 + columns_per_lane: cutlass.Constexpr = 4 + partial_bits: cutlass.Constexpr = mPartial.element_type.width + normalized: cutlass.Constexpr = partial_bits == 16 + band_base = half * Int32(band_columns) + lane * Int32(columns_per_lane) + gPartial_arr = cutlass.make_array_view(mPartial) + numerator0 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + numerator1 = cutlass.Vector.from_elements( + tuple(cutlass.Float32(0.0) for _ in range(columns_per_lane)), dtype=cutlass.Float32 + ) + denominator = cutlass.Float32(0.0) + global_row_max = cutlass.Float32(-cutlass.Float32.inf) + if request_tasks <= Int32(FOLD_TASKS): + loaded = _fold_round_loads( + mPartial, mRowMax, mRowSum, lane, local_row, half, task_begin, request_tasks, + physical_m, + ) # fmt: skip + global_row_max, denominator, numerator0, numerator1 = _fold_round_sums( + loaded, request_tasks, scale, normalized + ) + else: + for task_idx in cutlass.range(0, request_tasks, 1, unroll=1): + task = task_begin + task_idx + if mRowSum[task, local_row] != cutlass.Float32(0.0): + global_row_max = cute.arch.fmax(global_row_max, mRowMax[task, local_row] * scale) + for chunk_base in cutlass.range(0, request_tasks, FOLD_TASKS, unroll=1): + sums = [] + maxima = [] + vecs = [] + for k in cutlass.range_constexpr(FOLD_TASKS): + task = task_begin + cutlass.min(chunk_base + Int32(k), request_tasks - Int32(1)) + sums.append(mRowSum[task, local_row]) + maxima.append(mRowMax[task, local_row]) + partial_row = (task * physical_m + local_row) * Int32(LATENT_DIM) + band_base + for group in cutlass.range_constexpr(2): + vec = ( + gPartial_arr.subview(partial_row + Int32(group * 128)) + .data_ptr() + .load( + count=columns_per_lane, + alignment=partial_bits // 8 * columns_per_lane, + ) + ) + if cutlass.const_expr(normalized): + vec = vec.to(cutlass.Float32) + vecs.append(vec) + for k in cutlass.range_constexpr(FOLD_TASKS): + task_weight = cutlass.Float32(0.0) + if chunk_base + Int32(k) < request_tasks: + task_row_sum = sums[k] + if task_row_sum != cutlass.Float32(0.0): + task_max = _opaque_f32(maxima[k]) + task_weight = cute.exp2(task_max * scale - global_row_max, fastmath=True) + denominator += task_row_sum * task_weight + if cutlass.const_expr(normalized): + task_weight = task_weight * task_row_sum + if task_weight != cutlass.Float32(0.0): + numerator0 = numerator0 + vecs[2 * k] * task_weight + numerator1 = numerator1 + vecs[2 * k + 1] * task_weight + _fold_store(sLat, sMax, sSum, lane, row, half, global_row_max, denominator, numerator0, + numerator1, output_scale, pitch) # fmt: skip + + +class K3HelixTailKernel: + """Fold + V up-projection + exchange + combine of one Helix MLA layer (module docstring). + + Args: + cp_size: ``P`` (at most 8, or 16). + num_heads: ``H``, MLA heads every rank decodes (a multiple of ``2 * P``). + fold: The decode left FP32 split partials (fold them); else every row was stored + directly as BF16 with its statistics. + has_mask: Apply a per-token ``zero_kv_mask`` (uint8) on the sender side. + use_pdl: Launch with programmatic dependent launch. + mx_out: The receivers also store the o_proj operand: the MXFP8 form (E4M3 codes and one + UE8M0 exponent per 1x32 group) of ``bf16(out * gate)``, bitwise the gated low-M + GEMM's in-kernel quantization of it. + debug: Also store the folded latent rows, their statistics and the V up-projection + outputs, and trap instead of polling forever (tests). + timing: Thread 0 of every CTA stores ``%globaltimer`` at its phase boundaries into + ``ts[cta * 8 + phase]`` (microbenchmarks). + loopback: This rank's values of the heads of rank ``d`` go into its own receive buffer + as the record of source ``d``: every peer "sends" this rank's partials of its own heads + (the stand-in of single-GPU model tests). + ull2: (``features.ull2_mla_tail``; always with FP16 partials) the requests' first tasks + are staged in shared memory before the grid wait, rows with more than + :data:`FOLD_TASKS` tasks are folded in rounds of :data:`FOLD_TASKS` by the same code, + and the latent rows are padded (conflict-free V up-projection operands). Same bits. + """ + + def __init__( + self, + cp_size: int, + num_heads: int, + fold: bool, + has_mask: bool, + use_pdl: bool, + mx_out: bool = False, + debug: bool = False, + timing: bool = False, + loopback: bool = False, + mx_atom: bool = False, + ull2: bool = False, + ): + if not can_run(cp_size, num_heads, LATENT_DIM, V_HEAD_DIM): + raise ValueError(f"cp_size={cp_size}, num_heads={num_heads}") + self.P = cp_size + self.H = num_heads + self.H_LOC = num_heads // cp_size + self.fold = fold + self.has_mask = has_mask + self.use_pdl = use_pdl + self.mx_out = mx_out + # c240: the MXFP8 scales in the mid-M FP8 GEMM's layout (one 128-row block of tcgen05 atoms) + # instead of row-major (see ``mx_atom`` of k3_helix_tail). + self.mx_atom = mx_atom and mx_out + self.debug = debug + self.timing = timing + self.loopback = loopback + self.ull2 = ull2 + self.lat_pitch = _LAT_PITCH if ull2 else LATENT_DIM + self.plan_smem = _PLAN_SMEM if (ull2 and fold) else 1 + # Transposed V up-projection (ull2): K-half partial sums of 8 m16 tiles x 2 n8 tiles. + self.vsum_floats = (_WARPS // 2) * 2 * 32 * 4 if ull2 else 4 + self.threads = THREADS + # 16-byte vectors of a (token, source) record: the BF16 values of the H / P heads, then + # their FP32 (max, sum) pairs (two heads per vector); the layout of helix_a2a_fused. + self.DV = self.H_LOC * V_HEAD_DIM // 8 + self.SV = self.H_LOC // 2 + self.RV = self.DV + self.SV + # Columns (16-byte vectors) per head; a 16-lane group of a warp combines one head. + self.CPH = V_HEAD_DIM // 8 + # A CTA combines a token's heads in passes of THREADS columns. + self.passes = _ceil_div(self.DV, THREADS) + + @cute.jit + def __call__( + self, + partial_ptr: cute.Pointer, + maxima_ptr: cute.Pointer, + sums_ptr: cute.Pointer, + plan_ptr: cute.Pointer, + q_scale_ptr: cute.Pointer, + kv_scale_ptr: cute.Pointer, + rows_addr: Int64, + stats_addr: Int64, + w_addr: Int64, + mask_addr: Int64, + out_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + dbg_lat_addr: Int64, + dbg_v_addr: Int64, + dbg_stats_addr: Int64, + ts_addr: Int64, + gate_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + rank_stride: Int64, + half_bytes: Int64, + softmax_scale_log2: Float32, + output_scale: Float32, + num_tokens: Int32, + query_len: Int32, + physical_m: Int32, + num_tasks: Int32, + plan_len: Int32, + gate_ld: Int32, + rank: Int32, + tiles_per_cta: Int32, + n_ctas: Int32, + stream: cuda.CUstream, + ): + @cute.struct + class SharedStorage: + half: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 2], 16] + mbar: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 2], 16] + smax: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, ROWS], 16] + ssum: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, ROWS], 16] + lat: cute.struct.Align[cute.struct.MemRange[cutlass.BFloat16, ROWS * LATENT_DIM], 128] + cst: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, ROWS * _C_ROW_BYTES], 128] + wv: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int8, V_HEAD_DIM * _W_ROW_BYTES], 128 + ] + + # ull2: the staged plan, the padded latent rows and the second K half's V-up sums. + @cute.struct + class Ull2SharedStorage: + half: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 2], 16] + mbar: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 2], 16] + smax: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, ROWS], 16] + ssum: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, ROWS], 16] + plan: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.plan_smem], 16] + lat: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, ROWS * self.lat_pitch], 128 + ] + vsum: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, self.vsum_floats], 16] + cst: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, ROWS * _C_ROW_BYTES], 128] + wv: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int8, V_HEAD_DIM * _W_ROW_BYTES], 128 + ] + + self.shared_storage = SharedStorage + if cutlass.const_expr(self.ull2): + self.shared_storage = Ull2SharedStorage + mPartial = cute.make_tensor( + partial_ptr, + cute.make_layout( + (num_tasks, physical_m, LATENT_DIM), + stride=(cutlass.Int64(physical_m) * LATENT_DIM, LATENT_DIM, 1), + ), + ) + row_state = cute.make_layout((num_tasks, physical_m), stride=(physical_m, 1)) + mRowMax = cute.make_tensor(maxima_ptr, row_state) + mRowSum = cute.make_tensor(sums_ptr, row_state) + mPlan = cute.make_tensor(plan_ptr, cute.make_layout((plan_len,), stride=(1,))) + mQScale = cute.make_tensor(q_scale_ptr, cute.make_layout((1,), stride=(1,))) + mKvScale = cute.make_tensor(kv_scale_ptr, cute.make_layout((1,), stride=(1,))) + self.kernel( + mPartial, + mRowMax, + mRowSum, + mPlan, + mQScale, + mKvScale, + rows_addr, + stats_addr, + w_addr, + mask_addr, + out_addr, + ws_addr, + state_addr, + dbg_lat_addr, + dbg_v_addr, + dbg_stats_addr, + ts_addr, + gate_addr, + xq_addr, + xs_addr, + rank_stride, + half_bytes, + softmax_scale_log2, + output_scale, + num_tokens, + query_len, + physical_m, + gate_ld, + rank, + tiles_per_cta, + n_ctas, + ).launch( + grid=(n_ctas, 1, 1), + block=(self.threads, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + mPartial: cute.Tensor, + mRowMax: cute.Tensor, + mRowSum: cute.Tensor, + mPlan: cute.Tensor, + mQScale: cute.Tensor, + mKvScale: cute.Tensor, + rows_addr: Int64, + stats_addr: Int64, + w_addr: Int64, + mask_addr: Int64, + out_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + dbg_lat_addr: Int64, + dbg_v_addr: Int64, + dbg_stats_addr: Int64, + ts_addr: Int64, + gate_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + rank_stride: Int64, + half_bytes: Int64, + softmax_scale_log2: Float32, + output_scale: Float32, + num_tokens: Int32, + query_len: Int32, + physical_m: Int32, + gate_ld: Int32, + rank: Int32, + tiles_per_cta: Int32, + n_ctas: Int32, + ): + H = self.H + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = tidx // 32 + lane = tidx % 32 + ts = ts_addr + Int64(bidx) * 64 + self.stamp(tidx, ts, 0) + storage = utils.SmemAllocator().allocate(self.shared_storage) + half_slot = storage.half.data_ptr().toint() + mbar = storage.mbar.data_ptr().toint() + lat_base = storage.lat.data_ptr().toint() + smax_base = storage.smax.data_ptr().toint() + ssum_base = storage.ssum.data_ptr().toint() + wv_base = storage.wv.data_ptr().toint() + cst_base = storage.cst.data_ptr().toint() + # ull2-only shared memory; on the other path these names stand for values never read. + vsum_base = cst_base + if cutlass.const_expr(self.ull2): + vsum_base = storage.vsum.data_ptr().toint() + sLat = storage.lat.get_tensor( + cute.make_layout((ROWS, LATENT_DIM), stride=(self.lat_pitch, 1)) + ) + sPlan = mPlan + if cutlass.const_expr(self.ull2): + sPlan = storage.plan.get_tensor(cute.make_layout((self.plan_smem,), stride=(1,))) + sMax = storage.smax.get_tensor(cute.make_layout((ROWS,), stride=(1,))) + sSum = storage.ssum.get_tensor(cute.make_layout((ROWS,), stride=(1,))) + head = bidx % Int32(H) + cta_tile = bidx // Int32(H) + # W_UV of the head does not depend on the previous grid: its rows go into shared memory + # (padded stride) while this grid waits. Lane 0 of each of the last _W_ISSUE_WARPS warps + # issues an equal share of the row copies (the copies of one thread are issued one after + # another; thread 0 takes the ticket below). + if warp >= Int32(_WARPS - _W_ISSUE_WARPS): + if warp == Int32(_WARPS - 1): + if lane == 0: + _mbar_init(mbar, 1) + _mbar_expect_tx(mbar, _W_HEAD_BYTES) + cute.arch.barrier(barrier_id=1, number_of_threads=32 * _W_ISSUE_WARPS) + if lane == 0: + w_head = w_addr + Int64(head) * _W_HEAD_BYTES + rows_per_warp = V_HEAD_DIM // _W_ISSUE_WARPS + first = (warp - Int32(_WARPS - _W_ISSUE_WARPS)) * rows_per_warp + for r in cutlass.range(rows_per_warp, unroll=8): + n = first + r + _bulk_g2s( + wv_base + n * Int32(_W_ROW_BYTES), + w_head + Int64(n) * _LAT_ROW_BYTES, + _LAT_ROW_BYTES, + mbar, + ) + if cutlass.const_expr(self.ull2): + scale = Float32(0.0) + if cutlass.const_expr(self.fold): + # The packed reduction's scale expression (same operations, same result); the + # global scales are fixed tensors, read before the wait. + scale = softmax_scale_log2 / (mQScale[0] * mKvScale[0]) + # The requests' first tasks (plan ``begins``): written by the plan kernel, which + # completed before the decode this grid depends on started, so they are read + # before the wait. + num_requests = (num_tokens + query_len - Int32(1)) // query_len + for r in cutlass.range(tidx, num_requests + Int32(1), THREADS, unroll=1): + if r < Int32(_PLAN_SMEM): + sPlan[r] = mPlan[Int32(2) + r] + cute.arch.barrier() + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + # The ticket is taken after the wait (every earlier launch on the stream has completed); + # its result is first needed after the first fold. + ticket = Int64(0) + if tidx == 0: + ticket = _atom_add_u64(state_addr, Int64(1)) + self.stamp(tidx, ts, 1) + if cutlass.const_expr(not self.ull2): + scale = Float32(0.0) + if cutlass.const_expr(self.fold): + # The packed reduction's scale expression (same operations, same result). + scale = softmax_scale_log2 / (mQScale[0] * mKvScale[0]) + for tt in cutlass.range(tiles_per_cta, unroll=1): + token0 = (cta_tile * tiles_per_cta + tt) * Int32(ROWS) + if token0 < num_tokens: + # Zero-KV flag of the token this thread pushes (row tidx / 16), loaded ahead. + zero_kv = Int32(0) + if cutlass.const_expr(self.has_mask): + push_t = token0 + tidx // Int32(16) + if tidx < Int32(ROWS * 16) and push_t < num_tokens: + zero_kv = _ld_nc_u8(mask_addr + Int64(push_t)) + self.fold_tile( + warp, lane, head, token0, sLat, sMax, sSum, lat_base, smax_base, ssum_base, + mPartial, mRowMax, mRowSum, mPlan, sPlan, rows_addr, stats_addr, scale, + output_scale, num_tokens, query_len, physical_m, + ) # fmt: skip + if cutlass.const_expr(not self.ull2): + if tt == 0: + if tidx == 0: + if (ticket & Int64(0xFFFFFFFF)) == Int64(n_ctas) - Int64(1): + # Last arrival: clear the arrivals and advance the launch number. + _atom_add_u64(state_addr, Int64(1 << 32) - Int64(n_ctas)) + _sts_b64(half_slot, ((ticket >> 32) & Int64(1)) * half_bytes) + cute.arch.barrier() + self.stamp(tidx, ts, 2) + if cutlass.const_expr(self.debug): + self.dump_rows(tidx, head, token0, num_tokens, lat_base, smax_base, + ssum_base, dbg_lat_addr, dbg_stats_addr) # fmt: skip + if tt == 0: + landed = Int32(0) + while landed == 0: + landed = _mbar_try_wait(mbar, Int32(0)) + if cutlass.const_expr(self.ull2): + first_tile = Int32(0) + if tt == 0: + first_tile = Int32(1) + # The receive-buffer half is read inside, after the V up-projection + # (``rank_stride`` stands for the unused ``half_off``). + self.project_and_push( + tidx, warp, lane, head, token0, lat_base, smax_base, ssum_base, wv_base, + cst_base, vsum_base, zero_kv, ws_addr, dbg_v_addr, rank_stride, + rank_stride, half_slot, num_tokens, rank, ts, ticket, state_addr, n_ctas, + half_bytes, first_tile, + ) # fmt: skip + else: + # (``zero_kv`` stands for the unused ``first_tile``.) + self.project_and_push( + tidx, warp, lane, head, token0, lat_base, smax_base, ssum_base, wv_base, + cst_base, vsum_base, zero_kv, ws_addr, dbg_v_addr, rank_stride, + _lds_b64(half_slot), half_slot, num_tokens, rank, ts, ticket, state_addr, + n_ctas, half_bytes, zero_kv, + ) # fmt: skip + # The dependent grid (o_proj) is released once this CTA's pushes are issued: its CTAs take + # the SMs of exiting senders and stream their weights while the receivers wait for peers. + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + self.stamp(tidx, ts, 4) + self.receive( + tidx, + bidx, + out_addr, + ws_addr + Int64(rank) * rank_stride + _lds_b64(half_slot), + num_tokens, + n_ctas, + gate_addr, + gate_ld, + xq_addr, + xs_addr, + ts, + ) + self.stamp(tidx, ts, 7) + + @cute.jit + def take_half( + self, + tidx: Int32, + ticket: Int64, + state_addr: Int64, + n_ctas: Int32, + half_slot: Int32, + half_bytes: Int64, + ): + """Thread 0: the receive-buffer half of this launch from its ticket into shared memory + (read by every thread after the next CTA barrier); the last arrival advances the launch + number.""" + if tidx == 0: + if (ticket & Int64(0xFFFFFFFF)) == Int64(n_ctas) - Int64(1): + # Last arrival: clear the arrivals and advance the launch number. + _atom_add_u64(state_addr, Int64(1 << 32) - Int64(n_ctas)) + _sts_b64(half_slot, ((ticket >> 32) & Int64(1)) * half_bytes) + + @cute.jit + def stamp(self, tidx: Int32, ts: Int64, phase: int): + """``timing``: thread 0 stores ``%globaltimer`` as phase ``phase`` of this CTA.""" + if cutlass.const_expr(self.timing): + if tidx == 0: + _st_u64(ts + Int64(phase * 8), _globaltimer()) + + @cute.jit + def fold_tile( + self, + warp: Int32, + lane: Int32, + head: Int32, + token0: Int32, + sLat: cute.Tensor, + sMax: cute.Tensor, + sSum: cute.Tensor, + lat_base: Int32, + smax_base: Int32, + ssum_base: Int32, + mPartial: cute.Tensor, + mRowMax: cute.Tensor, + mRowSum: cute.Tensor, + mPlan: cute.Tensor, + sPlan: cute.Tensor, + rows_addr: Int64, + stats_addr: Int64, + scale: Float32, + output_scale: Float32, + num_tokens: Int32, + query_len: Int32, + physical_m: Int32, + ): + """Latent rows (BF16) and (max, sum) statistics of tokens ``token0 .. token0 + 15`` of + ``head`` into shared memory; rows past the batch are zeros.""" + H = self.H + half = warp % Int32(2) + for slot in cutlass.range_constexpr(2): + i = warp // Int32(2) + Int32(slot * _WARPS // 2) + t = token0 + i + # This lane's 16 bytes of the 512-byte half row. + lat_vec = lat_base + i * Int32(self.lat_pitch * 2) + half * Int32(_LAT_ROW_BYTES // 2) + lat_vec = lat_vec + lane * Int32(16) + if t < num_tokens: + row = Int64(t) * H + Int64(head) + folded = Int32(0) + if cutlass.const_expr(self.fold): + request = t // query_len + local_row = (t - request * query_len) * Int32(H) + head + if cutlass.const_expr(self.ull2): + task_begin = Int32(0) + task_end = Int32(0) + if request < Int32(_PLAN_SMEM - 1): + task_begin = sPlan[request] + task_end = sPlan[request + Int32(1)] + else: + task_begin = mPlan[Int32(2) + request] + task_end = mPlan[Int32(3) + request] + request_tasks = task_end - task_begin + else: + task_begin = mPlan[Int32(2) + request] + request_tasks = mPlan[Int32(3) + request] - task_begin + if request_tasks > Int32(1): + folded = Int32(1) + if cutlass.const_expr(self.ull2): + _fold_half_ull2( + mPartial, mRowMax, mRowSum, sLat, sMax, sSum, lane, i, local_row, + half, task_begin, request_tasks, physical_m, scale, output_scale, + _LAT_PITCH, + ) # fmt: skip + else: + if request_tasks <= Int32(FOLD_TASKS): + _fold_half( + mPartial, mRowMax, mRowSum, sLat, sMax, sSum, lane, i, + local_row, half, task_begin, request_tasks, physical_m, scale, + output_scale, + ) # fmt: skip + else: + _reduce_packed_row( + mPartial, mRowMax, mRowSum, sLat, sMax, sSum, lane, i, + local_row, half, task_begin, request_tasks, physical_m, scale, + output_scale, 256, True, + ) # fmt: skip + else: + if half == Int32(0) and lane == Int32(0): + # The packed reduction's statistics of a single-task request. + row_sum = mRowSum[task_begin, local_row] + stats_max = cutlass.Float32(-cutlass.Float32.inf) + if row_sum != cutlass.Float32(0.0): + stats_max = ( + mRowMax[task_begin, local_row] + * scale + * cutlass.Float32(1.0 / LOG2_E) + ) + sMax[i] = stats_max + sSum[i] = row_sum + else: + if half == Int32(0) and lane == Int32(0): + st = _ld_nc_v2(stats_addr + row * 8) + _sts_b32(smax_base + i * 4, st[0]) + _sts_b32(ssum_base + i * 4, st[1]) + if folded == 0: + # The decode stored this row itself: copy this lane's 16 bytes of the half. + src = rows_addr + row * _LAT_ROW_BYTES + Int64(half * (_LAT_ROW_BYTES // 2)) + _sts_v4(lat_vec, *_ld_nc_v4(src + Int64(lane) * 16)) + else: + zero = Int32(0) + _sts_v4(lat_vec, zero, zero, zero, zero) + + @cute.jit + def project_and_push( + self, + tidx: Int32, + warp: Int32, + lane: Int32, + head: Int32, + token0: Int32, + lat_base: Int32, + smax_base: Int32, + ssum_base: Int32, + wv_base: Int32, + cst_base: Int32, + vsum_base: Int32, + zero_kv: Int32, + ws_addr: Int64, + dbg_v_addr: Int64, + rank_stride: Int64, + half_off: Int64, + half_slot: Int32, + num_tokens: Int32, + rank: Int32, + ts: Int64, + ticket: Int64, + state_addr: Int64, + n_ctas: Int32, + half_bytes: Int64, + first_tile: Int32, + ): + """V up-projection of the tile (warp ``w``: output columns ``8 w .. 8 w + 7``) and the + push of its rows and statistics into the head owner's receive buffer (``zero_kv``: the + pushed token owns no KV on this rank).""" + H = self.H + g = lane // 4 + q = lane % 4 + if cutlass.const_expr(self.ull2): + self.project_transposed(warp, lane, token0, num_tokens, lat_base, wv_base, cst_base, + vsum_base) # fmt: skip + else: + # Two accumulators (the two k16 steps of every 32-column block) halve the MMA + # dependency chain; they are added at the end. + accs = [[Float32(0.0)] * 4 for _ in range(2)] + w_row = wv_base + (warp * 8 + g) * Int32(_W_ROW_BYTES) + for u in cutlass.range_constexpr(LATENT_DIM // 32): + col2 = (Int32(u * 32) + q * 8) * 2 + b = _lds_v4(w_row + col2) + a_lo = _lds_v4(lat_base + g * Int32(self.lat_pitch * 2) + col2) + a_hi = _lds_v4(lat_base + (g + 8) * Int32(self.lat_pitch * 2) + col2) + for v in cutlass.range_constexpr(2): + accs[v] = list( + _mma_bf16( + a_lo[2 * v], + a_hi[2 * v], + a_lo[2 * v + 1], + a_hi[2 * v + 1], + b[2 * v], + b[2 * v + 1], + *accs[v], + ) + ) + acc = [accs[0][c] + accs[1][c] for c in range(4)] + # The BF16 outputs are staged token-major in their own buffer, so the next tile's + # fold (after the barrier below) cannot overwrite what this tile's pushes still read; + # the statistics this thread pushes are read before that barrier for the same reason. + n = warp * 8 + q * 2 + _sts_b32(cst_base + g * Int32(_C_ROW_BYTES) + n * 2, _pack_bf16x2(acc[0], acc[1])) + _sts_b32(cst_base + (g + 8) * Int32(_C_ROW_BYTES) + n * 2, _pack_bf16x2(acc[2], acc[3])) + if cutlass.const_expr(self.ull2): + # The launch ticket is first needed by the pushes below: its atomic's latency + # overlaps the fold and the V up-projection. + if first_tile != Int32(0): + self.take_half(tidx, ticket, state_addr, n_ctas, half_slot, half_bytes) + i = tidx // 16 + vec = tidx % 16 + mx = _lds_b32(smax_base + (i % Int32(ROWS)) * 4) + sm = _lds_b32(ssum_base + (i % Int32(ROWS)) * 4) + cute.arch.barrier() + self.stamp(tidx, ts, 3) + if cutlass.const_expr(self.ull2): + # This launch's receive-buffer half (``take_half`` above, before the barrier). + half_off = _lds_b64(half_slot) + # Row i's 16 vectors (and its statistics) into the head owner's record. + if tidx < Int32(ROWS * 16): + t = token0 + i + if t < num_tokens: + words = _lds_v4(cst_base + i * Int32(_C_ROW_BYTES) + vec * 16) + if cutlass.const_expr(self.debug): + row = Int64(t) * H + Int64(head) + _st_v4(dbg_v_addr + row * (V_HEAD_DIM * 2) + Int64(vec) * 16, *words) + dst_rank = head // Int32(self.H_LOC) + local_head = head - dst_rank * Int32(self.H_LOC) + sent = _prep4(words[0], words[1], words[2], words[3], zero_kv, Int32(0)) + st = _prep4(mx, sm, Int32(0), Int32(0), zero_kv, Int32(1)) + # Record (t, this rank) of the head owner's buffer; with loopback, record + # (t, owner) of this rank's own buffer. + box_rank = dst_rank + src_rank = rank + if cutlass.const_expr(self.loopback): + box_rank = rank + src_rank = dst_rank + record = ( + ws_addr + + Int64(box_rank) * rank_stride + + half_off + + Int64(t * self.P + src_rank) * (self.RV * 16) + ) + _st_v4(record + Int64(local_head) * (V_HEAD_DIM * 2) + Int64(vec) * 16, *sent) + if vec == 0: + _st_v2(record + Int64(self.DV * 16) + Int64(local_head) * 8, st[0], st[1]) + + @cute.jit + def vup_k_half(self, w_lo: Int32, w_hi: Int32, lat_row: Int32, k_col: Int32): + """One warp's m16 x n8 tile of ``W_UV @ lat^T`` over 256 K columns starting at + ``k_col`` (this lane's 8-column offset included): two accumulators, added at the end.""" + accs = [[Float32(0.0)] * 4 for _ in range(2)] + for u in cutlass.range_constexpr(LATENT_DIM // 64): + col2 = (k_col + Int32(u * 32)) * 2 + a_lo = _lds_v4(w_lo + col2) + a_hi = _lds_v4(w_hi + col2) + b = _lds_v4(lat_row + col2) + for v in cutlass.range_constexpr(2): + accs[v] = list( + _mma_bf16( + a_lo[2 * v], + a_hi[2 * v], + a_lo[2 * v + 1], + a_hi[2 * v + 1], + b[2 * v], + b[2 * v + 1], + *accs[v], + ) + ) + return [accs[0][c] + accs[1][c] for c in range(4)] + + @cute.jit + def project_transposed( + self, + warp: Int32, + lane: Int32, + token0: Int32, + num_tokens: Int32, + lat_base: Int32, + wv_base: Int32, + cst_base: Int32, + vsum_base: Int32, + ): + """``ull2`` V up-projection as ``out^T = W_UV @ lat^T``: warp ``w`` runs the m16 tile of + output columns ``16 (w % 8) ..`` over K half ``w // 8``, for the tile's first 8 tokens and + (only if the tile has more) its last 8, so a tile of at most 8 tokens issues half the + MMAs of the token-major form. The second K half's sums go through shared memory and are + added to the first's; the BF16 results are staged token-major like the token-major form's + (the same dot products, another FP32 summation order).""" + g = lane // 4 + q = lane % 4 + m_tile = warp % Int32(_WARPS // 2) + k_half = warp // Int32(_WARPS // 2) + pitch_bytes: cutlass.Constexpr = self.lat_pitch * 2 + w_lo = wv_base + (m_tile * 16 + g) * Int32(_W_ROW_BYTES) + w_hi = w_lo + Int32(8 * _W_ROW_BYTES) + lat_row = lat_base + g * Int32(pitch_bytes) + k_col = k_half * Int32(LATENT_DIM // 2) + q * 8 + s0 = self.vup_k_half(w_lo, w_hi, lat_row, k_col) + # The tile's last 8 tokens only if it has any. + c0 = Float32(0.0) + c1 = Float32(0.0) + c2 = Float32(0.0) + c3 = Float32(0.0) + if token0 + Int32(8) < num_tokens: + r = self.vup_k_half(w_lo, w_hi, lat_row + Int32(8 * pitch_bytes), k_col) + c0 = r[0] + c1 = r[1] + c2 = r[2] + c3 = r[3] + s1 = [c0, c1, c2, c3] + sums = [s0, s1] + # Second K half: (m tile, n tile, lane, 4) partial sums into shared memory. + if k_half == Int32(1): + for n_tile in cutlass.range_constexpr(2): + slot = vsum_base + (((m_tile * 2 + Int32(n_tile)) * 32 + lane) * 4) * 4 + for c in cutlass.range_constexpr(4): + _sts_f32(slot + Int32(c * 4), sums[n_tile][c]) + cute.arch.barrier(barrier_id=2, number_of_threads=THREADS) + if k_half == Int32(0): + for n_tile in cutlass.range_constexpr(2): + slot = vsum_base + (((m_tile * 2 + Int32(n_tile)) * 32 + lane) * 4) * 4 + d = [sums[n_tile][c] + _lds_f32(slot + Int32(c * 4)) for c in range(4)] + # d0, d1: output column 16 m + g, tokens 2 q, 2 q + 1; d2, d3: column + 8. + lo = _pack_bf16x2(d[0], d[2]) + hi = _pack_bf16x2(d[1], d[3]) + col = (m_tile * 16 + g) * 2 + t0 = Int32(n_tile * 8) + q * 2 + row0 = cst_base + t0 * Int32(_C_ROW_BYTES) + col + row1 = row0 + Int32(_C_ROW_BYTES) + _sts_b16(row0, lo) + _sts_b16(row0 + Int32(16), lo >> 16) + _sts_b16(row1, hi) + _sts_b16(row1 + Int32(16), hi >> 16) + + @cute.jit + def receive( + self, + tidx: Int32, + rcta: Int32, + out_addr: Int64, + recv_base: Int64, + num_tokens: Int32, + n_recv: Int32, + gate_addr: Int64, + gate_ld: Int32, + xq_addr: Int64, + xs_addr: Int64, + ts: Int64, + ): + """Combine the records of this CTA's tokens: ``HelixA2AFusedKernel.receive`` (the same + instruction sequences, so bitwise ``helix_post_process_native``) over ``passes`` groups + of THREADS columns. Thread ``c`` of a pass owns one 16-byte column (8 output values) of one + head; lane ``l`` of a 16-lane head group polls the statistics of source rank ``l``.""" + P, RV, DV, CPH = self.P, self.RV, self.DV, self.CPH + lane = tidx % 32 + group = lane - lane % 16 + for t in cutlass.range(rcta, num_tokens, n_recv, unroll=1): + rec0 = recv_base + Int64(t * (P * RV)) * 16 + for p in cutlass.range_constexpr(self.passes): + c = tidx + Int32(p * THREADS) + h = c // CPH + col = c - h * CPH + src_l = col % 16 + valid = Int32(c < DV) + has_stats = valid * Int32(src_l < P) + # With 16 columns per head (V_HEAD_DIM 128) every group owns its head's statistics. + owns_stats = has_stats * Int32(col < 16) + st_addr = rec0 + Int64(src_l * RV + DV) * 16 + Int64(h * 8) + d_addr = rec0 + Int64(c) * 16 + stats = cute.make_rmem_tensor((2,), Int32) + stats[0] = Int32(_NEG_INF_BITS) + stats[1] = Int32(0) + vals = cute.make_rmem_tensor((P * 4,), Int32) + g = cute.make_rmem_tensor((4,), Int32) + if cutlass.const_expr(self.mx_out): + # The gate is ready before this grid starts: its load overlaps the polls. + for e in cutlass.range_constexpr(4): + g[e] = Int32(0) + if valid != 0: + gw = _ld_nc_v4(gate_addr + (Int64(t) * Int64(gate_ld) + Int64(c * 8)) * 2) + for e in cutlass.range_constexpr(4): + g[e] = gw[e] + if valid != 0: + ready = Int32(0) + spins = Int32(0) + while ready == 0: + if cutlass.const_expr(self.debug): + # Tests: a record that never arrives fails the launch, not the job. + spins = spins + 1 + if spins > Int32(_DEBUG_POLL_LIMIT): + _trap() + ready = Int32(1) + if has_stats != 0: + sv = _ld_vol_v2(st_addr) + stats[0] = sv[0] + stats[1] = sv[1] + ready = _all_valid2(sv[0], sv[1]) + for s in cutlass.range_constexpr(P): + w = _ld_vol_v4(d_addr + Int64(s * RV * 16)) + for e in cutlass.range_constexpr(4): + vals[s * 4 + e] = w[e] + ready = ready & _all_valid4(*w) + mx = _as_f32(stats[0]) + gmax = mx + for off in (1, 2, 4, 8): + gmax = _fmax(gmax, _shfl_bfly(gmax, off)) + corrected = _corrected_exp(mx, gmax, _as_f32(stats[1])) + total = corrected + for off in (1, 2, 4, 8): + total = _fadd(total, _shfl_bfly(total, off)) + weight = _fmul(corrected, _rcp(total)) + acc = [Float32(0.0)] * 8 + for s in cutlass.range_constexpr(P): + ws = _shfl_idx(weight, group + s) + for e in cutlass.range_constexpr(4): + acc[2 * e], acc[2 * e + 1] = _bf16x2_fma( + vals[s * 4 + e], ws, acc[2 * e], acc[2 * e + 1] + ) + o = [_pack_bf16x2(acc[2 * e], acc[2 * e + 1]) for e in range(4)] + if valid != 0: + _st_v4(out_addr + Int64(t * DV + c) * 16, *o) + for s in cutlass.range_constexpr(P): + _st_sentinel_v4(d_addr + Int64(s * RV * 16)) + if owns_stats != 0: + _st_sentinel_v2(st_addr) + if cutlass.const_expr(self.mx_out): + self.store_mx(o, g, valid, t, c, xq_addr, xs_addr) + if cutlass.const_expr(self.timing and p == 0): + # Thread 0 has seen every source of its token's first column. + if tidx == 0: + _st_u64(ts + Int64(16 + 8 * 3), _globaltimer()) + + @cute.jit + def store_mx( + self, + o, + g: cute.Tensor, + valid: Int32, + t: Int32, + c: Int32, + xq_addr: Int64, + xs_addr: Int64, + ): + """MXFP8 of ``bf16(out * gate)`` for column ``c`` (8 values ``o``, gate words ``g``) of + token ``t``: the gate multiply, 1x32 amax (the four threads of a group, by shuffles), + UE8M0 exponent and E4M3 codes of the gated low-M GEMM's quantization warps. Every lane of + the warp runs the shuffles; groups are either all valid or all invalid (``DV`` is a + multiple of 16).""" + K = self.DV * 8 + xg = _gate_mul_bf16x8(o[0], o[1], o[2], o[3], g[0], g[1], g[2], g[3]) + amax = amax_bf16x8_words(xg[0], xg[1], xg[2], xg[3]) + amax = amax_pair(amax, _shfl_bfly_b32(amax, 1)) + amax = amax_pair(amax, _shfl_bfly_b32(amax, 2)) + e8 = _ue8m0_of_amax(amax) + q0, q1 = _quant_bf16x8_to_e4m3(xg[0], xg[1], xg[2], xg[3], _bf16x2_inv_scale(e8)) + if valid != 0: + _st_v2(xq_addr + Int64(t) * K + Int64(c * 8), q0, q1) + if (c % 4) == 0: + if cutlass.const_expr(self.mx_atom): + # Mid-M scale atom of group g = c / 4 for token t < 128: byte (g / 4) * 512 + # + (t % 32) * 16 + (t / 32) * 4 + g % 4. + grp = c // 4 + _st_u8( + xs_addr + + Int64(grp // 4) * 512 + + Int64(t & 31) * 16 + + Int64((t >> 5) & 3) * 4 + + Int64(grp % 4), + e8, + ) + else: + _st_u8(xs_addr + Int64(t) * (K // 32) + Int64(c // 4), e8) + + @cute.jit + def dump_rows( + self, + tidx: Int32, + head: Int32, + token0: Int32, + num_tokens: Int32, + lat_base: Int32, + smax_base: Int32, + ssum_base: Int32, + dbg_lat_addr: Int64, + dbg_stats_addr: Int64, + ): + """Tests: store the tile's latent rows ``[T * H, 512]`` and statistics ``[T * H, 2]``.""" + H = self.H + for k in cutlass.range_constexpr(ROWS * _LAT_ROW_BYTES // 16 // THREADS): + idx = tidx + Int32(k * THREADS) + i = idx // Int32(_LAT_ROW_BYTES // 16) + vec = idx % Int32(_LAT_ROW_BYTES // 16) + t = token0 + i + if t < num_tokens: + words = _lds_v4(lat_base + i * Int32(self.lat_pitch * 2) + vec * 16) + row = Int64(t) * H + Int64(head) + _st_v4(dbg_lat_addr + row * _LAT_ROW_BYTES + Int64(vec) * 16, *words) + if tidx < Int32(ROWS): + t = token0 + tidx + if t < num_tokens: + mx = _lds_b32(smax_base + tidx * 4) + sm = _lds_b32(ssum_base + tidx * 4) + row = Int64(t) * H + Int64(head) + _st_v2(dbg_stats_addr + row * 8, mx, sm) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +class FoldInputs: + """The packed decode's split state that the tail folds instead of the reduction kernel + (the arguments of ``reduce_packed_splits``). + + Args: + partial: FP32 ``[tasks, physical_m, 512]`` split partials, or the FP16 normalized split + outputs of ``features.fp4_mla_f16_partials`` (``output_scale`` 1). + maxima: FP32 ``[tasks, physical_m]`` raw row maxima. + sums: FP32 ``[tasks, physical_m]`` row sums relative to them. + plan: int32 task plan (``plan[2 + r]`` = first task of request ``r``). + q_scale: FP32 scalar device tensor, the FP4 Q global scale. + kv_scale: FP32 scalar device tensor, the FP4 KV global scale. + num_requests: Requests of the packed decode. + valid_m: ``query_len * H`` rows per request. + physical_m: ``valid_m`` padded to 128. + sm_scale: Softmax scale. + output_scale: Scale of the normalized output. + """ + + def __init__( + self, + partial: torch.Tensor, + maxima: torch.Tensor, + sums: torch.Tensor, + plan: torch.Tensor, + q_scale: torch.Tensor, + kv_scale: torch.Tensor, + *, + num_requests: int, + valid_m: int, + physical_m: int, + sm_scale: float, + output_scale: float, + ): + self.partial = partial + self.maxima = maxima + self.sums = sums + self.plan = plan + self.q_scale = q_scale + self.kv_scale = kv_scale + self.num_requests = num_requests + self.valid_m = valid_m + self.physical_m = physical_m + self.sm_scale = sm_scale + self.output_scale = output_scale + + +def can_run(cp_size: int, num_heads: int, latent_dim: int, v_head_dim: int) -> bool: + """Whether :func:`k3_helix_tail` supports this Helix MLA geometry.""" + if latent_dim != LATENT_DIM or v_head_dim != V_HEAD_DIM: + return False + if not (1 <= cp_size <= 8 or cp_size == MAX_CP): + return False + return num_heads % (2 * cp_size) == 0 + + +@functools.cache +def _kernel( + cp_size: int, + num_heads: int, + fold: bool, + has_mask: bool, + use_pdl: bool, + mx_out: bool, + debug: bool, + timing: bool, + loopback: bool, + mx_atom: bool = False, + ull2: bool = False, +): + return K3HelixTailKernel( + cp_size, num_heads, fold, has_mask, use_pdl, mx_out, debug, timing, loopback, mx_atom, ull2 + ) + + +_COMPILED: dict = {} + + +def _gmem(dtype, tensor: torch.Tensor, align: int): + return make_ptr(dtype, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=align) + + +def k3_helix_tail( + rows: torch.Tensor, + stats: torch.Tensor | None, + w_uv: torch.Tensor, + out: torch.Tensor, + *, + fold: FoldInputs | None, + zero_kv_mask: torch.Tensor | None, + cp_rank: int, + cp_size: int, + ws_addr: int, + rank_stride: int, + half_bytes: int, + state: torch.Tensor, + max_ctas: int, + use_pdl: bool = True, + mx_out: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + debug: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + timestamps: torch.Tensor | None = None, + loopback: bool = False, + mx_atom: bool = False, + ull2: bool = False, +) -> None: + """Launch the Helix tail of ``T`` tokens into ``out``. + + Args: + rows: BF16 ``[T * H, 512]`` latent rows of the decode (row ``t * H + h``). With ``fold``, + only the rows of single-task requests are read (the decode stored them). + stats: FP32 ``[T * H, 2]`` ``(max, sum)`` of ``rows`` (natural-log max); read only + without ``fold``. + w_uv: BF16 ``[H, 128, 512]`` V up-projection of every head. + out: BF16 ``[T, H / P * 128]`` combined output of this rank's heads. + fold: The packed decode's split state to fold, or None. + zero_kv_mask: Optional bool ``[T]``, True where this rank owns no KV of the token. + cp_rank, cp_size: This rank and ``P``. + ws_addr, rank_stride, half_bytes: Rank 0's receive buffer in this process's address space + (rank ``r``'s at ``ws_addr + r * rank_stride``), two halves of ``half_bytes`` each, + every word armed with the sentinel. + state: int64 launch counter shared by all launches on the workspace. + max_ctas: Cap on the CTAs (at most one per SM: every CTA polls after its own pushes, so + all of them must be resident together). A CTA folds and sends ``tiles_per_cta`` tiles + of 16 tokens of one head, the fewest that keep the grid within the cap. + use_pdl: Programmatic dependent launch. + mx_out: Optional ``(gate, xq, xs)``: the BF16 ``[T, H / P * 128]`` gate (row stride a + multiple of 8), and outputs for the MXFP8 form of ``bf16(out * gate)``: E4M3 ``xq`` + ``[T, H / P * 128]`` and UE8M0 (uint8) ``xs`` ``[T, H / P * 4]``. + debug: Optional ``(latent [T * H, 512] BF16, values [T * H, 128] BF16, stats [T * H, 2] + FP32)`` outputs of the sender stages (tests). + timestamps: Optional int64 ``[CTAs, 8]`` (``CTAs`` at most ``max_ctas``): + ``%globaltimer`` at the phase boundaries of every CTA (0 start, 1 after the wait, 2 + fold, 3 V up-projection, 4 pushes done, 5 first column combined, 7 end; 2 and 3 of + the CTA's last tile). + loopback: Send to this rank's own buffer, as source ``d``, the values of rank ``d``'s heads + (single-GPU model tests). + mx_atom: (c240, with ``mx_out``, ``T <= 128``) ``xs`` is ``uint8 [(H / P * 128 / 128) * + 512]`` in the mid-M FP8 GEMM's scale layout (``k3_fp8w_lowm_gemm.run_mid_prequantized`` + with one 128-row block); the same exponents. + ull2: The ``features.ull2_mla_tail`` sender path (:class:`K3HelixTailKernel`); always + taken for FP16 split partials, which the original fold does not read. + """ + num_heads = w_uv.shape[0] + num_tokens = rows.shape[0] // num_heads + if num_tokens == 0: + return + if num_tokens * cp_size * record_bytes(num_heads // cp_size, V_HEAD_DIM) > half_bytes: + raise ValueError("the Helix tail receive buffer is too small for this many tokens") + mx_atom = mx_atom and mx_out is not None + if mx_atom and num_tokens > 128: + raise ValueError(f"mx_atom covers at most 128 tokens, got {num_tokens}") + has_mask = zero_kv_mask is not None + flags = (fold is not None, has_mask, use_pdl, mx_out is not None, debug is not None, + timestamps is not None, loopback, mx_atom) # fmt: skip + fold16 = fold is not None and fold.partial.dtype == torch.float16 + ull2 = ull2 or fold16 + kernel = _kernel(cp_size, num_heads, *flags, ull2) + n_tiles = _ceil_div(num_tokens, ROWS) + tiles_per_cta = _ceil_div(num_heads * n_tiles, max(max_ctas, num_heads)) + n_ctas = num_heads * _ceil_div(n_tiles, tiles_per_cta) + if fold is not None: + ptrs = ( + _gmem(cutlass.Float16 if fold16 else cutlass.Float32, fold.partial, 32), + _gmem(cutlass.Float32, fold.maxima, 32), + _gmem(cutlass.Float32, fold.sums, 32), + _gmem(cutlass.Int32, fold.plan, 4), + _gmem(cutlass.Float32, fold.q_scale, 4), + _gmem(cutlass.Float32, fold.kv_scale, 4), + ) + query_len = fold.valid_m // num_heads + dims = (fold.physical_m, fold.partial.shape[0], fold.plan.numel()) + scales = (Float32(fold.sm_scale * LOG2_E), Float32(fold.output_scale)) + else: + # Never dereferenced without ``fold``; any valid pointers keep one signature. + ptrs = ( + _gmem(cutlass.Float32, stats, 4), + _gmem(cutlass.Float32, stats, 4), + _gmem(cutlass.Float32, stats, 4), + _gmem(cutlass.Int32, stats, 4), + _gmem(cutlass.Float32, stats, 4), + _gmem(cutlass.Float32, stats, 4), + ) + query_len, dims = 1, (1, 1, 1) + scales = (Float32(1.0), Float32(1.0)) + dbg = (0, 0, 0) if debug is None else tuple(t.data_ptr() for t in debug) + mx = (0, 0, 0) if mx_out is None else tuple(t.data_ptr() for t in mx_out) + gate_ld = 0 if mx_out is None else mx_out[0].stride(0) + args = ( + *ptrs, + Int64(rows.data_ptr()), + Int64(0 if stats is None else stats.data_ptr()), + Int64(w_uv.data_ptr()), + Int64(zero_kv_mask.data_ptr() if has_mask else 0), + Int64(out.data_ptr()), + Int64(ws_addr), + Int64(state.data_ptr()), + *(Int64(a) for a in dbg), + Int64(0 if timestamps is None else timestamps.data_ptr()), + *(Int64(a) for a in mx), + Int64(rank_stride), + Int64(half_bytes), + *scales, + Int32(num_tokens), + Int32(query_len), + *(Int32(d) for d in dims), + Int32(gate_ld), + Int32(cp_rank), + Int32(tiles_per_cta), + Int32(n_ctas), + cuda.CUstream(torch.cuda.current_stream(rows.device).cuda_stream), + ) + key = (cp_size, num_heads, *flags, ull2, fold16) + compiled = _COMPILED.get(key) + if compiled is None: + compiled = cute.compile(kernel, *args) + _COMPILED[key] = compiled + compiled(*args) + + +def reference_v_up(latent: torch.Tensor, w_uv: torch.Tensor) -> torch.Tensor: + """FP64 V up-projection of BF16 latent rows ``[T * H, 512]`` with ``w_uv`` ``[H, 128, 512]``, + rounded to BF16 (tests).""" + num_heads = w_uv.shape[0] + lat = latent.view(-1, num_heads, LATENT_DIM).double() + out = torch.einsum("tjk,jnk->tjn", lat, w_uv.double()) + return out.reshape(-1, V_HEAD_DIM).to(torch.bfloat16) + + +def bf16_ulp(x: torch.Tensor) -> torch.Tensor: + """Spacing of BF16 values at ``x`` (tests).""" + mag = x.float().abs().clamp_min(2.0**-126) + return torch.exp2(torch.floor(torch.log2(mag)) - 7) + + +assert math.isclose(LOG2_E, math.log2(math.e)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_kda_fab.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_kda_fab.py new file mode 100644 index 000000000000..ab588e2f3cfb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_kda_fab.py @@ -0,0 +1,2040 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused Kimi-K3 KDA ``[f_a | b] -> f_b`` decode projection (SM100/SM103/SM107). + +One launch computes, for ``x`` of ``1 <= M <= 256`` token rows, + +* ``bfa = bf16(x @ W_bfa^T)`` (``W_bfa`` = ``[f_a_proj; b_proj]`` rows, FP32 accumulation); +* ``beta = bfa[:, 128 : 128 + H]``; +* ``forget_gate = bf16(bfa[:, :128] @ W_fb^T)`` (FP32 accumulation), + +i.e. what a ``[f_a | b]`` GEMM, a slice copy and an ``f_b`` GEMM produce in three to five +launches. + +The kernel runs on a second stream next to the KDA ``qkvg`` GEMM, which fills 96-120 SMs with +~220 KB of shared memory each (clusters of 4-5 CTAs). It therefore uses 27-28 large CTAs that fit +in the SMs the GEMM leaves free; more CTAs (or small ones that land between the GEMM's clusters) +delay the GEMM. + +Design: + +* GEMM 1, ``M <= 64`` (``mma.sync``). CTA ``(g, r, s)`` owns token group ``g`` (``8 * ntg`` + tokens), row group ``r`` (``9 / R`` 16-row tiles of ``W_bfa``, rows past ``N1`` zero-filled by + TMA) and the ``s``-th slice of K in 64-wide chunks. One thread TMA-loads the CTA's whole weight + slice (before ``griddepcontrol.wait``) and activation slice up front, one mbarrier per chunk. + MMA warps tile the CTA's rows x tokens (``mt`` x ``nt`` m16n8k16 tiles each, rows on the MMA M + side) and split the chunks into ``kg`` groups, summed through shared memory in a fixed order. +* GEMM 1, larger M (``tcgen05`` UMMA). CTA ``(g, s)`` owns 128 tokens (UMMA M) and all ``N1`` + rows (UMMA N = 144) of the ``s``-th K slice: the weight slice is TMA-loaded up front, the + activation streams through a ring, the accumulator goes TMEM -> shared memory -> partial slot. +* Hand-off. Every CTA stores one FP32 partial tile ``[tokens, N1]`` into slots of a per-device + workspace that are filled with a sentinel (``0xffffffff``, which no FP32 operation produces); + consumers poll the data itself, so no fence, flag or counter sits between the producers' stores + and the consumers' loads. The workspace has two halves used by alternate launches. A control + warp per CTA takes a ticket from a packed 64-bit counter (high half: launch number, low half: + CTAs of the launch that have arrived; the last arrival moves it to the next launch), caches the + scalar parameters in shared memory, warms L2 with a share of ``W_fb``, and refills with the + sentinel its share of the region the previous launch dirtied in the other half. +* Tail. The CTAs with the last ``T`` ranks (``T`` = tail work items) become workers. Worker + ``(tb, rb)`` TMA-loads its ``W_fb`` row block, polls the ``R * S`` partials of its tokens (in + register-sized batches with all loads of a batch in flight), sums them in split order, rounds + ``f_a`` to BF16 into shared memory (worker ``rb == 0`` also writes ``beta``) and computes its + forget-gate rows with ``mma.sync``. A worker only waits for CTAs with higher ranks or that + already run, so progress needs ``T + 1`` co-resident CTAs, not grid-wide co-residency. +* Launches sharing a device's counter and workspace must not overlap (stream order; with PDL the + ticket and the refill happen after ``griddepcontrol.wait``). + +``f_a`` mode (:func:`k3_kda_fab_fa`) runs the same GEMM 1 and hand-off but stops at +``bfa = bf16([f_a | b])``: every CTA becomes a worker that sums an equal share of the +``[token, 4-row unit]`` partials (same split order, same rounding, so bitwise the ``f_a`` and +``beta`` of the full kernel) and stores them as BF16 rows. The KDA decode then computes the +forget gate itself (``kda_mtp_decode_v2`` ``FUSED_FB``), which removes ``W_fb``, the f_b MMA and +the forget-gate round trip from this kernel. A worker polls the partials of every CTA, so all +(at most 28) CTAs must become co-resident; next to the ``qkvg`` GEMM, which does not depend on +this kernel, they do at the latest when that GEMM retires. +""" + +from __future__ import annotations + +import dataclasses +import functools +from typing import Optional + +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 +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Width of ``f_a`` (the ``W_fb`` reduction size). +FA_DIM = 128 +#: Largest token count. +MAX_M = 256 +#: Largest ``[f_a | b]`` row count (nine 16-row MMA tiles). +MAX_N1 = 144 + +_K_CHUNK = 64 +_ROW_BYTES = _K_CHUNK * 2 +_M_TILE = 16 +_N_TILE = 8 +_N1_TILES = MAX_N1 // _M_TILE +#: UMMA GEMM 1: tokens per CTA (UMMA M) and tensor-memory columns (144 accumulator columns). +_UMMA_M = 128 +_UMMA_TMEM_COLS = 256 +#: Row pitch (floats) of the UMMA accumulator staging tile (conflict-free 16-byte accesses). +_UMMA_STAGE_PITCH = MAX_N1 + 4 +#: The tail probes one unit per producer before loading every partial when a thread loads at +#: least this many partial units per round. +_PROBE_MIN_LOADS = 32 +#: Partial floats a tail thread keeps in registers per polling round. +_POLL_FLOATS = 128 +#: The tail stages its forget-gate tile in shared memory for 16-byte stores at this size. +_STAGE_MIN_ELEMS = 4096 +_SMEM_CAPACITY = utils.get_smem_capacity_in_bytes("sm_100") +#: L2 eviction hint for the streamed weight (CacheHintSm90::EVICT_FIRST). +_EVICT_FIRST = 0x12F0000000000000 +#: Partial-tile workspace per device: two halves used by alternate launches. +_WS_HALF_BYTES = 4 << 20 +#: Device state words: packed ticket, then the dirtied bytes of each workspace half. +_STATE_WORDS = 8 +_TICKET_WORD, _DIRTY_WORD = 0, 1 +#: The sentinel that marks an empty FP32 workspace word. +_SENTINEL = -1 + +# Shared-memory cache of scalar kernel parameters (constant-cache misses in the tail are slow). +(_P_M, _P_BETA, _P_FG, _P_HALF, _P_RANK, _P_PREV, _P_OTHER) = range(7) +_NUM_PARAMS = 8 + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +@dataclasses.dataclass(frozen=True) +class FabConfig: + """Tiling of one specialization. + + Attributes: + ntg: GEMM-1 token tiles (of 8) per CTA. + r1: GEMM-1 row groups (1, 3 or 9). + s: GEMM-1 K splits. + mt: 16-row MMA tiles per warp (divides ``9 // r1``). + nt: 8-token MMA tiles per warp. + kg: K groups of warps (summed through shared memory). + ntb: Tail token tiles per worker (divides ``ntg``). + mtb: Tail ``W_fb`` 16-row tiles per worker. + umma: GEMM 1 on ``tcgen05`` UMMA (128-token CTAs: ``ntg = 16``, ``r1 = 1``; ``mt``, + ``nt`` and ``kg`` unused) instead of ``mma.sync``. + """ + + ntg: int + r1: int + s: int + mt: int + nt: int + kg: int + ntb: int + mtb: int + umma: bool = False + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _lv(value, *, loc=None, ip=None): + return Int64(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return Float32(value).ir_value(loc=loc, ip=ip) + + +def _unpack(res, ty, cls, n, *, loc=None, ip=None): + return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n)) + + +@dsl_user_op +def _ldsm_x4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _mma_bf16(a, b0, b1, c, *, loc=None, ip=None): + """``d[16x8] = a[16x16] @ b[16x8] + c`` (BF16 inputs, FP32 accumulate).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(x, loc=loc, ip=ip) for x in (*a, b0, b1)] + [_fv(x, loc=loc, ip=ip) for x in c], + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _sts_v4_f32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)] + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", + "r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_bf16x4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + """Round four FP32 values to BF16 (nearest even) and store them as 8 bytes.""" + _asm( + None, + [_iv(addr, loc=loc, ip=ip)] + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n.reg .b32 lo, hi;\ncvt.rn.bf16x2.f32 lo, $2, $1;\ncvt.rn.bf16x2.f32 hi, $4, $3;\n" + "st.shared.v2.b32 [$0], {lo, hi};\n}", + "r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_b64(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "st.shared.b64 [$0], $1;", + "r,l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_b64(addr, *, loc=None, ip=None): + return Int64( + _asm( + T.i64(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.b64 $0, [$1];", "=l,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _ld_relaxed_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.relaxed.gpu.global.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _st_global_bf16(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\nst.global.b16 [$0], h;\n}", + "l,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_global_bf16x4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + """Round four FP32 values to BF16 (nearest even) and store them as 8 bytes.""" + _asm( + None, + [_lv(addr, loc=loc, ip=ip)] + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n.reg .b32 lo, hi;\ncvt.rn.bf16x2.f32 lo, $2, $1;\ncvt.rn.bf16x2.f32 hi, $4, $3;\n" + "st.global.v2.b32 [$0], {lo, hi};\n}", + "l,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _atom_add_u64(addr, value, *, loc=None, ip=None): + return Int64( + _asm( + T.i64(), + [_lv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "atom.relaxed.gpu.global.add.u64 $0, [$1], $2;", + "=l,l,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _red_add_u64(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "red.relaxed.gpu.global.add.u64 [$0], $1;", + "l,l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _lv(value, loc=loc, ip=ip)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_v4_f32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)] + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.relaxed.gpu.global.v4.f32 [$0], {$1, $2, $3, $4};", + "l,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_sentinel_v4(addr, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)], + "{\n.reg .b32 t;\nmov.b32 t, -1;\nst.global.v4.b32 [$0], {t, t, t, t};\n}", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _all_valid4(v0, v1, v2, v3, *, loc=None, ip=None): + """1 if none of the four FP32 words is the empty-slot sentinel ``0xffffffff``.""" + return Int32( + _asm( + T.i32(), + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n.reg .pred p;\n.reg .b32 t;\n" + "mov.b32 t, $1;\nsetp.ne.b32 p, t, -1;\n" + "mov.b32 t, $2;\nsetp.ne.and.b32 p, t, -1, p;\n" + "mov.b32 t, $3;\nsetp.ne.and.b32 p, t, -1, p;\n" + "mov.b32 t, $4;\nsetp.ne.and.b32 p, t, -1, p;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return Int64( + _asm( + T.i64(), + [_lv(addr, loc=loc, ip=ip)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _min_s32(a, b, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "min.s32 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_shared(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.shared::cta;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _prefetch_l2(addr, nbytes, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(nbytes, loc=loc, ip=ip)], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_bf16(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\nst.shared.b16 [$0], h;\n}", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _copy_s2g_16(dst, src, *, loc=None, ip=None): + """16 bytes from shared ``src`` to global ``dst``.""" + _asm( + None, + [_lv(dst, loc=loc, ip=ip), _iv(src, loc=loc, ip=ip)], + "{\n.reg .b32 a, b, c, d;\nld.shared.v4.b32 {a, b, c, d}, [$1];\n" + "st.global.v4.b32 [$0], {a, b, c, d};\n}", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.wait;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.launch_dependents;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _lds_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.b32 $0, [$1];", "=r,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _tmem_ld16(taddr, *, loc=None, ip=None): + """16 consecutive FP32 accumulator columns of this thread's TMEM lane.""" + regs = ", ".join(f"${i}" for i in range(16)) + res = _asm( + llvm.StructType.get_literal([T.f32()] * 16), + [_iv(taddr, loc=loc, ip=ip)], + "tcgen05.fence::after_thread_sync;\n" + f"tcgen05.ld.sync.aligned.32x32b.x16.b32 {{{regs}}}, [$16];\n" + "tcgen05.wait::ld.sync.aligned;", + ",".join(["=f"] * 16) + ",r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 16, loc=loc, ip=ip) + + +@dsl_user_op +def _tmem_fence_before_sync(*, loc=None, ip=None): + _asm(None, [], "tcgen05.fence::before_thread_sync;", "", loc=loc, ip=ip) + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +class _FabTail: + """Control warp, hand-off and tail shared by the GEMM-1 variants (see the module docstring). + + Subclasses set the geometry attributes (``num_ctas``, ``items``, ``tg``, ``threads``, ...), + run :meth:`control_warp` and :meth:`refill` on one warp, store their partial tiles into the + workspace half cached in shared memory and then run :meth:`tail` on the worker CTAs. + """ + + def _init_common( + self, cfg, n_tiles, n1, n2, k, num_heads, use_pdl, geo, fa_mode=False, warm_w2=True + ): + self.cfg = cfg + self.n_tiles = n_tiles + self.n1 = n1 + self.n2 = n2 + self.k = k + self.num_heads = num_heads + self.use_pdl = use_pdl + self.fa_mode = fa_mode + # Bytes of W_fb the control warps warm in L2 (the tail's operand; in f_a mode the KDA + # decode's, whose per-head TMA load would otherwise come from HBM). + self.warm_bytes = n2 * FA_DIM * 2 if warm_w2 else 0 + for name, value in geo.items(): + setattr(self, name, value) + if fa_mode: + # Every CTA is a worker over an equal share of the [token, 4-row unit] partials. + self.items = self.num_ctas + units_max = n_tiles * _N_TILE * (n1 // 4) + self.fa_units_per_thread = _ceil_div(_ceil_div(units_max, self.items), self.threads) + self.dirty_bytes = self.groups * cfg.s * self.tg * n1 * 4 + self.beta_units = _ceil_div(num_heads, 4) + # Tail units (4 consecutive rows of one token) a thread polls at once. + self.units_per_thread = _ceil_div( + cfg.ntb * _N_TILE * (FA_DIM // 4 + self.beta_units), self.threads + ) + + @cute.jit + def make_w2_tma(self, w2_ptr: cute.Pointer): + """TMA atom of ``W_fb`` boxes and their shared-memory layout (the tail's row block).""" + mW2 = cute.make_tensor(w2_ptr, cute.make_layout((self.n2, FA_DIM), stride=(FA_DIM, 1))) + atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mW2, + cute.tile_to_shape(atom, (self.w2_box_rows, _K_CHUNK), order=(0, 1)), + (self.w2_box_rows, _K_CHUNK), + ) + layout = cute.tile_to_shape(atom, (self.w2_box_rows, _K_CHUNK, self.w2_boxes), (0, 1, 2)) + return tma_atom, tma_tensor, layout + + @cute.jit + def control_warp( + self, + params: Int32, + bidx: Int32, + lane: Int32, + m: Int32, + w2_addr: Int64, + beta_addr: Int64, + fg_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + ): + """Ticket, parameter cache (in shared memory) and ``W_fb`` L2 warm-up.""" + if lane == 0: + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + ticket = _atom_add_u64(state_addr + Int64(_TICKET_WORD * 8), Int64(1)) + launch = ticket >> 32 + rank = ticket & Int64(0xFFFFFFFF) + if rank == self.num_ctas - 1: + # Last arrival: clear the arrivals and advance the launch number in one add. + _red_add_u64(state_addr + Int64(_TICKET_WORD * 8), Int64((1 << 32) - self.num_ctas)) + parity = launch & Int64(1) + prev = _ld_relaxed_u64(state_addr + Int64((_DIRTY_WORD + 1 - parity) * 8)) + if rank == 0: + _st_relaxed_u64( + state_addr + Int64((_DIRTY_WORD + parity) * 8), Int64(self.dirty_bytes) + ) + cached = ( + Int64(m), + beta_addr, + fg_addr, + ws_addr + parity * _WS_HALF_BYTES, + rank, + prev, + ws_addr + (Int64(1) - parity) * _WS_HALF_BYTES, + ) + for i in cutlass.range_constexpr(len(cached)): + _sts_b64(params + i * 8, cached[i]) + # Warm L2 with this CTA's share of W_fb for the tail. + total = self.warm_bytes + if cutlass.const_expr(total > 0): + share = _ceil_div(_ceil_div(total, self.num_ctas), 16) * 16 + off = bidx * share + if off < total: + nbytes = Int32(share) + if off + share > total: + nbytes = total - off + _prefetch_l2(w2_addr + Int64(off), nbytes) + cute.arch.sync_warp() + + @cute.jit + def refill(self, params: Int32, lane: Int32): + """Refill this CTA's share of what the previous launch dirtied in the other half.""" + prev = _lds_b64(params + _P_PREV * 8) + other = _lds_b64(params + _P_OTHER * 8) + rank = _lds_b64(params + _P_RANK * 8) + share = ((prev + self.num_ctas - 1) // self.num_ctas + 15) // 16 * 16 + begin = rank * share + end = begin + share + if end > prev: + end = prev + off = begin + Int64(lane * 16) + while off < end: + _st_sentinel_v4(other + off) + off = off + 512 + + @cute.jit + def tail( + self, + region: Int32, + sW2: cute.Tensor, + w2_full, + tma_atom_w2: cute.CopyAtom, + mW2t: cute.Tensor, + params: Int32, + tidx: Int32, + warp: Int32, + lane: Int32, + worker: Int32, + m: Int32, + half: Int64, + ): + cfg = self.cfg + S = cfg.s + beta_addr = _lds_b64(params + _P_BETA * 8) + fg_addr = _lds_b64(params + _P_FG * 8) + tb = worker % self.token_blocks + rb = worker // self.token_blocks + blk_rows = cfg.mtb * _M_TILE + sw2 = region + sfa = sw2 + self.w2_smem_bytes + tb_rows = cfg.ntb * _N_TILE + sout = sfa + tb_rows * FA_DIM * 2 + out_pitch = blk_rows + 8 + stage = tb_rows * blk_rows >= _STAGE_MIN_ELEMS + + # W_fb row block -> shared memory: TMA boxes of w2_box_rows x 64 columns (128B + # swizzle), rows past N2 zero-filled. + row0 = rb * blk_rows + if warp == 0: + # The block reuses shared memory last accessed through the generic proxy. + _fence_proxy_async_shared() + gW2 = cute.local_tile(mW2t, (self.w2_box_rows, _K_CHUNK), (None, None)) + tW2s, tW2g = cpasync.tma_partition( + tma_atom_w2, + 0, + cute.make_layout(1), + cute.group_modes(sW2, 0, 2), + cute.group_modes(gW2, 0, 2), + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(w2_full, self.w2_smem_bytes) + box0 = row0 // self.w2_box_rows + for b in cutlass.range_constexpr(self.w2_boxes): + cute.copy( + tma_atom_w2, + tW2g[(None, box0 + b // 2, b % 2)], + tW2s[(None, b)], + tma_bar_ptr=w2_full, + ) + + # f_a (and, for rb == 0, beta) of the block's tokens: poll the R * S partials of each + # 4-row unit until no slot holds the sentinel, then sum them in split order. + t0 = tb * tb_rows + g = t0 // self.tg + upt = FA_DIM // 4 + self.beta_units * Int32(rb == 0) + total = tb_rows * upt + slot_stride = self.tg * self.n1 * 4 + U = self.units_per_thread + valid = [] + srcs = [] + for it in cutlass.range_constexpr(U): + idx = tidx + it * self.threads + tl = idx // upt + u = idx % upt + # Unit u covers rows [4 u, 4 u + 4): f_a for u < 32, beta (rows 128 + ...) after. + row = 4 * u + t = t0 + tl + valid.append(Int32(idx < total) * Int32(t < m)) + srcs.append(half + Int64(((g * S * self.tg + t - g * self.tg) * self.n1 + row) * 4)) + # Cheap wait first: lane i of warp 0 polls one unit of producer slot i (split, row + # group) until it is set; the full loads below then rarely need a second round. + if cutlass.const_expr(U * S >= _PROBE_MIN_LOADS): + self.probe(warp, lane, t0, tb_rows, g, m, half) + # Units are polled in batches of `batch` (all their loads in flight together) so the + # partials of a batch fit in registers. + batch = max(1, _POLL_FLOATS // (4 * S)) + for b0 in cutlass.range_constexpr(0, U, batch): + nb = min(batch, U - b0) + parts = cute.make_rmem_tensor((nb * S * 4,), Float32) + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for i in cutlass.range_constexpr(nb): + if valid[b0 + i] != 0: + for sp in cutlass.range_constexpr(S): + v = _ld_relaxed_v4_f32(srcs[b0 + i] + Int64(sp * slot_stride)) + for e in cutlass.range_constexpr(4): + parts[(i * S + sp) * 4 + e] = v[e] + ready = ready & _all_valid4(*v) + for i in cutlass.range_constexpr(nb): + if valid[b0 + i] != 0: + idx = tidx + (b0 + i) * self.threads + tl = idx // upt + u = idx % upt + t = t0 + tl + v = [parts[i * S * 4 + e] for e in range(4)] + for sp in cutlass.range_constexpr(1, S): + for e in cutlass.range_constexpr(4): + v[e] = v[e] + parts[(i * S + sp) * 4 + e] + if u < FA_DIM // 4: + kc = u // 16 + unit = (u % 16) // 2 + dst = sfa + kc * (tb_rows * _ROW_BYTES) + tl * _ROW_BYTES + dst = dst + ((unit ^ (tl & 7)) << 4) + (u % 2) * 8 + _sts_bf16x4(dst, *v) + else: + for e in cutlass.range_constexpr(4): + h = 4 * (u - FA_DIM // 4) + e + if h < self.num_heads: + _st_global_bf16( + beta_addr + Int64((t * self.num_heads + h) * 2), v[e] + ) + cute.arch.mbarrier_wait(w2_full, 0) + cute.arch.barrier() + + # forget_gate[block tokens, block rows] = f_a @ W_fb^T. + mat = lane >> 3 + r = lane & 7 + a_unit = mat >> 1 + b_unit = mat & 1 + ntb = cfg.ntb + b_rows = [ + ((2 * p + (mat >> 1) * min(1, ntb - 1 - 2 * p)) * _N_TILE + r) * _ROW_BYTES + for p in range(_ceil_div(ntb, 2)) + ] + # The block's f_a fragments are loaded once per warp; a warp's m-tiles (warp, warp + + # num_warps, ...) accumulate independently. + nk = FA_DIM // 16 + bfrag = cute.make_rmem_tensor((nk * _ceil_div(ntb, 2) * 4,), Int32) + for kk in cutlass.range_constexpr(nk): + b_chunk = sfa + (kk // 4) * (tb_rows * _ROW_BYTES) + for p in cutlass.range_constexpr(_ceil_div(ntb, 2)): + b = _ldsm_x4(b_chunk + b_rows[p] + ((((kk % 4) * 2 + b_unit) ^ r) << 4)) + for i in cutlass.range_constexpr(4): + bfrag[(kk * _ceil_div(ntb, 2) + p) * 4 + i] = b[i] + mpw = _ceil_div(cfg.mtb, self.num_warps) + acc = cute.make_rmem_tensor((mpw * ntb * 4,), Float32) + acc.fill(0.0) + a_rows = [] + for w in cutlass.range_constexpr(mpw): + # Warps past the block's last m-tile redo it (their results are not stored). + mi = _min_s32(warp + w * self.num_warps, cfg.mtb - 1) + row_l = mi * _M_TILE + (mat & 1) * 8 + r + a_rows.append( + sw2 + + (row_l // self.w2_box_rows) * 2 * (self.w2_box_rows * _ROW_BYTES) + + (row_l % self.w2_box_rows) * _ROW_BYTES + ) + for kk in cutlass.range_constexpr(nk): + for w in cutlass.range_constexpr(mpw): + a = _ldsm_x4( + a_rows[w] + + (kk // 4) * (self.w2_box_rows * _ROW_BYTES) + + ((((kk % 4) * 2 + a_unit) ^ r) << 4) + ) + for j in cutlass.range_constexpr(ntb): + p = j // 2 + q = j % 2 + o = (w * ntb + j) * 4 + base = (kk * _ceil_div(ntb, 2) + p) * 4 + d = _mma_bf16( + a, + bfrag[base + 2 * q], + bfrag[base + 2 * q + 1], + (acc[o], acc[o + 1], acc[o + 2], acc[o + 3]), + ) + for e in cutlass.range_constexpr(4): + acc[o + e] = d[e] + for w in cutlass.range_constexpr(mpw): + mi = warp + w * self.num_warps + if mi < cfg.mtb: + for j in cutlass.range_constexpr(ntb): + for e in cutlass.range_constexpr(4): + tl = j * _N_TILE + 2 * (lane & 3) + (e & 1) + row_l = mi * _M_TILE + (lane >> 2) + (e >> 1) * 8 + val = acc[(w * ntb + j) * 4 + e] + if cutlass.const_expr(stage): + _sts_bf16(sout + (tl * out_pitch + row_l) * 2, val) + elif t0 + tl < m and row0 + row_l < self.n2: + _st_global_bf16( + fg_addr + Int64(((t0 + tl) * self.n2 + row0 + row_l) * 2), val + ) + if cutlass.const_expr(stage): + # Coalesced 16-byte stores of the staged [tokens, rows] tile. + cute.arch.barrier() + for idx in cutlass.range(tidx, tb_rows * (blk_rows // 8), self.threads, unroll=1): + tl = idx // (blk_rows // 8) + uu = idx % (blk_rows // 8) + t = t0 + tl + row = row0 + uu * 8 + if t < m and row < self.n2: + _copy_s2g_16( + fg_addr + Int64((t * self.n2 + row) * 2), + sout + (tl * out_pitch + uu * 8) * 2, + ) + + @cute.jit + def tail_fa(self, params: Int32, tidx: Int32, worker: Int32, m: Int32, half: Int64): + """f_a mode: sum this worker's share of the partials and store ``bf16`` rows of ``bfa``. + + Unit ``idx`` is rows ``[4 u, 4 u + 4)`` of token ``t`` (``idx = t * N1 / 4 + u``); its + ``S`` split partials are polled until set and summed in split order, as in :meth:`tail`. + """ + S = self.cfg.s + out_addr = _lds_b64(params + _P_FG * 8) + upt = self.n1 // 4 + total = m * upt + per = (total + self.items - 1) // self.items + begin = worker * per + end = begin + per + if end > total: + end = total + slot_stride = self.tg * self.n1 * 4 + U = self.fa_units_per_thread + batch = max(1, _POLL_FLOATS // (4 * S)) + for b0 in cutlass.range_constexpr(0, U, batch): + nb = min(batch, U - b0) + valid = [] + srcs = [] + dsts = [] + for i in cutlass.range_constexpr(nb): + idx = begin + tidx + (b0 + i) * self.threads + t = idx // upt + u = idx - t * upt + g = t // self.tg + valid.append(Int32(idx < end)) + srcs.append( + half + Int64(((g * S * self.tg + t - g * self.tg) * self.n1 + 4 * u) * 4) + ) + dsts.append(out_addr + Int64((t * self.n1 + 4 * u) * 2)) + parts = cute.make_rmem_tensor((nb * S * 4,), Float32) + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for i in cutlass.range_constexpr(nb): + if valid[i] != 0: + for sp in cutlass.range_constexpr(S): + v = _ld_relaxed_v4_f32(srcs[i] + Int64(sp * slot_stride)) + for e in cutlass.range_constexpr(4): + parts[(i * S + sp) * 4 + e] = v[e] + ready = ready & _all_valid4(*v) + for i in cutlass.range_constexpr(nb): + if valid[i] != 0: + v = [parts[i * S * 4 + e] for e in range(4)] + for sp in cutlass.range_constexpr(1, S): + for e in cutlass.range_constexpr(4): + v[e] = v[e] + parts[(i * S + sp) * 4 + e] + _st_global_bf16x4(dsts[i], *v) + + @cute.jit + def probe( + self, warp: Int32, lane: Int32, t0: Int32, tb_rows: int, g: Int32, m: Int32, half: Int64 + ): + """Lane i of warp 0 polls one unit of producer slot i (split, row group) until set.""" + cfg = self.cfg + S = cfg.s + if warp == 0: + if lane < S * cfg.r1: + sp = lane // cfg.r1 + row_end = (lane - sp * cfg.r1 + 1) * self.rows_cta + if row_end > self.n1: + row_end = Int32(self.n1) + t_last = t0 + tb_rows - 1 + if t_last > m - 1: + t_last = m - 1 + probe = half + Int64( + (((g * S + sp) * self.tg + t_last - g * self.tg) * self.n1 + row_end - 4) * 4 + ) + v = _ld_relaxed_v4_f32(probe) + ok = _all_valid4(*v) + while ok == 0: + v = _ld_relaxed_v4_f32(probe) + ok = _all_valid4(*v) + cute.arch.barrier() + + +class K3KdaFabKernel(_FabTail): + """``beta``, ``forget_gate`` from ``x``, ``W_bfa`` and ``W_fb`` (see the module docstring). + + Args: + cfg: Tiling. + n_tiles: Token tiles of 8 (``ceil(M / 8)``) of the specialization. + n1: Rows of ``W_bfa`` (``128 + H`` padded to a multiple of 8, at most 144). + n2: Rows of ``W_fb`` (a multiple of 16). + k: Hidden size (a multiple of 64). + num_heads: ``H``; ``beta`` is rows ``[128, 128 + H)`` of ``bfa``. + use_pdl: Launch with programmatic dependent launch. + fa_mode: Stop at ``bfa`` (written through the forget-gate address; see the module + docstring). + warm_w2: Warm L2 with ``W_fb`` (``n2`` rows at ``w2_addr``). + """ + + def __init__( + self, + cfg: FabConfig, + n_tiles: int, + n1: int, + n2: int, + k: int, + num_heads: int, + use_pdl: bool, + fa_mode: bool = False, + warm_w2: bool = True, + ): + validate_config(cfg, n_tiles, n1, n2, k, num_heads) + self._init_common( + cfg, + n_tiles, + n1, + n2, + k, + num_heads, + use_pdl, + _geometry(cfg, n_tiles, n2, k), + fa_mode, + warm_w2, + ) + + @cute.jit + def __call__( + self, + w1_ptr: cute.Pointer, + x_ptr: cute.Pointer, + w2_ptr: cute.Pointer, + m: Int32, + w2_addr: Int64, + beta_addr: Int64, + fg_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + stream: cuda.CUstream, + ): + mW = cute.make_tensor(w1_ptr, cute.make_layout((self.n1, self.k), stride=(self.k, 1))) + mX = cute.make_tensor(x_ptr, cute.make_layout((m, self.k), stride=(self.k, 1))) + atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + w_layout = cute.tile_to_shape(atom, (self.rows_cta, _K_CHUNK, self.max_chunks), (0, 1, 2)) + x_layout = cute.tile_to_shape(atom, (self.tg, _K_CHUNK, self.max_chunks), (0, 1, 2)) + tma_atom_w, tma_w = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mW, + cute.tile_to_shape(atom, (self.rows_cta, _K_CHUNK), order=(0, 1)), + (self.rows_cta, _K_CHUNK), + ) + tma_atom_x, tma_x = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mX, + cute.tile_to_shape(atom, (self.tg, _K_CHUNK), order=(0, 1)), + (self.tg, _K_CHUNK), + ) + tma_atom_w2, tma_w2, w2_layout = self.make_w2_tma(w2_ptr) + + @cute.struct + class SharedStorage: + mbar: cute.struct.MemRange[cutlass.Int64, self.max_chunks] + w2_full: cute.struct.MemRange[cutlass.Int64, 1] + params: cute.struct.MemRange[cutlass.Int64, _NUM_PARAMS] + # The weight and activation chunk buffers; once the GEMM-1 MMAs are done, the same + # bytes hold the warp-partial reduction and then the tail's operands. + sW: cute.struct.Align[cute.struct.MemRange[cutlass.BFloat16, self.sw_elems], 1024] + sX: cute.struct.Align[cute.struct.MemRange[cutlass.BFloat16, self.sx_elems], 1024] + + self.shared_storage = SharedStorage + self.kernel( + tma_atom_w, + tma_w, + tma_atom_x, + tma_x, + w_layout, + x_layout, + tma_atom_w2, + tma_w2, + w2_layout, + m, + w2_addr, + beta_addr, + fg_addr, + ws_addr, + state_addr, + ).launch( + grid=(self.num_ctas, 1, 1), + block=(self.threads, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + tma_atom_w: cute.CopyAtom, + mWt: cute.Tensor, + tma_atom_x: cute.CopyAtom, + mXt: cute.Tensor, + w_layout: cute.ComposedLayout, + x_layout: cute.ComposedLayout, + tma_atom_w2: cute.CopyAtom, + mW2t: cute.Tensor, + w2_layout: cute.ComposedLayout, + m: Int32, + w2_addr: Int64, + beta_addr: Int64, + fg_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + ): + cfg = self.cfg + S = cfg.s + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + group = bidx % self.groups + rest = bidx // self.groups + rgrp = rest % cfg.r1 + split = rest // cfg.r1 + base_chunks = self.num_k_chunks // S + rem_chunks = self.num_k_chunks % S + chunk_count = Int32(base_chunks) + chunk_start = split * base_chunks + if cutlass.const_expr(rem_chunks > 0): + if split < rem_chunks: + chunk_count = chunk_count + 1 + chunk_start = chunk_start + split + else: + chunk_start = chunk_start + rem_chunks + + storage = utils.SmemAllocator().allocate(self.shared_storage) + mbar = storage.mbar.data_ptr() + params = storage.params.data_ptr().toint() + if warp == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(self.max_chunks): + cute.arch.mbarrier_init(mbar + i, 1) + cute.arch.mbarrier_init(storage.w2_full.data_ptr(), 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + # ---- GEMM-1 loads: the weight never depends on the previous grid. ---- + if warp == 0: + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + cpasync.prefetch_descriptor(tma_atom_w) + cpasync.prefetch_descriptor(tma_atom_x) + sW = storage.sW.get_tensor(w_layout.outer, swizzle=w_layout.inner) + sX = storage.sX.get_tensor(x_layout.outer, swizzle=x_layout.inner) + gW = cute.local_tile(mWt, (self.rows_cta, _K_CHUNK), (None, None)) + gX = cute.local_tile(mXt, (self.tg, _K_CHUNK), (None, None)) + tWs, tWg = cpasync.tma_partition( + tma_atom_w, + 0, + cute.make_layout(1), + cute.group_modes(sW, 0, 2), + cute.group_modes(gW, 0, 2), + ) + tXs, tXg = cpasync.tma_partition( + tma_atom_x, + 0, + cute.make_layout(1), + cute.group_modes(sX, 0, 2), + cute.group_modes(gX, 0, 2), + ) + for c in cutlass.range(chunk_count, unroll=1): + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + mbar + c, self.w_chunk_bytes + self.x_chunk_bytes + ) + cute.copy( + tma_atom_w, + tWg[(None, rgrp, chunk_start + c)], + tWs[(None, c)], + tma_bar_ptr=mbar + c, + cache_policy=Int64(_EVICT_FIRST), + ) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + for c in cutlass.range(chunk_count, unroll=1): + cute.copy( + tma_atom_x, + tXg[(None, group, chunk_start + c)], + tXs[(None, c)], + tma_bar_ptr=mbar + c, + ) + + if warp == self.num_mma_warps: + self.control_warp( + params, bidx, lane, m, w2_addr, beta_addr, fg_addr, ws_addr, state_addr + ) + self.refill(params, lane) + else: + self.gemm1(storage, mbar, warp, lane, chunk_count) + + # ---- Warp partials were written to shared memory; sum them in K-group order and + # store the CTA's FP32 partial tile into the current workspace half. ---- + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + cute.arch.barrier() + m = Int32(_lds_b64(params + _P_M * 8)) + half = _lds_b64(params + _P_HALF * 8) + red = storage.sW.data_ptr().toint() + units = self.rows_cta // 4 + part_base = half + Int64((group * S + split) * self.tg * self.n1 * 4) + for o in cutlass.range(tidx, self.tg * units, self.threads, unroll=1): + token = o // units + u = o % units + row = rgrp * self.rows_cta + 4 * u + if group * self.tg + token < m and row < self.n1: + v = list(_lds_v4_f32(red + (token * self.red_stride + 4 * u) * 4)) + for w in cutlass.range_constexpr(1, cfg.kg): + v2 = _lds_v4_f32(red + ((w * self.tg + token) * self.red_stride + 4 * u) * 4) + for e in cutlass.range_constexpr(4): + v[e] = v[e] + v2[e] + _st_relaxed_v4_f32(part_base + Int64((token * self.n1 + row) * 4), *v) + cute.arch.barrier() + + # ---- Tail: the last T ranks are workers. ---- + worker = Int32(_lds_b64(params + _P_RANK * 8)) - (self.num_ctas - self.items) + sW2 = storage.sW.get_tensor(w2_layout.outer, swizzle=w2_layout.inner) + if cutlass.const_expr(self.fa_mode): + self.tail_fa(params, tidx, worker, m, half) + elif worker >= 0: + self.tail( + red, + sW2, + storage.w2_full.data_ptr(), + tma_atom_w2, + mW2t, + params, + tidx, + warp, + lane, + worker, + m, + half, + ) + + @cute.jit + def gemm1(self, storage, mbar, warp: Int32, lane: Int32, chunk_count: Int32): + """GEMM-1 MMAs of one MMA warp; its partial tile goes to shared memory.""" + cfg = self.cfg + mt = cfg.mt + nt = cfg.nt + wk = warp % cfg.kg + wmn = warp // cfg.kg + wn = wmn % self.warps_n + wm = wmn // self.warps_n + acc = cute.make_rmem_tensor((mt * nt * 4,), Float32) + acc.fill(0.0) + w_base = storage.sW.data_ptr().toint() + x_base = storage.sX.data_ptr().toint() + mat = lane >> 3 + r = lane & 7 + a_unit = mat >> 1 + b_unit = mat & 1 + a_rows = [((wm * mt + i) * _M_TILE + (mat & 1) * 8 + r) * _ROW_BYTES for i in range(mt)] + # Token tiles of this warp, clamped to the CTA's last tile (clamped tiles are not stored). + tiles = [_min_s32(wn * nt + j, cfg.ntg - 1) for j in range(nt)] + # ldmatrix.x4 row addresses of each tile pair (an odd last tile loads itself twice). + b_rows = [] + for p in cutlass.range_constexpr(_ceil_div(nt, 2)): + second = tiles[min(2 * p + 1, nt - 1)] + tile = tiles[2 * p] + (mat >> 1) * (second - tiles[2 * p]) + b_rows.append((tile * _N_TILE + r) * _ROW_BYTES) + for c in cutlass.range(wk, chunk_count, cfg.kg, unroll=1): + cute.arch.mbarrier_wait(mbar + c, 0) + a_chunk = w_base + c * self.w_chunk_bytes + x_chunk = x_base + c * self.x_chunk_bytes + for ks in cutlass.range_constexpr(_K_CHUNK // 16): + a = [ + _ldsm_x4(a_chunk + a_rows[i] + (((ks * 2 + a_unit) ^ r) << 4)) + for i in range(mt) + ] + for p in cutlass.range_constexpr(_ceil_div(nt, 2)): + b = _ldsm_x4(x_chunk + b_rows[p] + (((ks * 2 + b_unit) ^ r) << 4)) + for q in cutlass.range_constexpr(min(2, nt - 2 * p)): + j = 2 * p + q + for i in cutlass.range_constexpr(mt): + o = (i * nt + j) * 4 + d = _mma_bf16( + a[i], + b[2 * q], + b[2 * q + 1], + (acc[o], acc[o + 1], acc[o + 2], acc[o + 3]), + ) + for e in cutlass.range_constexpr(4): + acc[o + e] = d[e] + # All MMA warps must be done with the chunk buffers before they become the reduction + # buffer (named barrier over the MMA warps only). + cute.arch.barrier(barrier_id=1, number_of_threads=32 * self.num_mma_warps) + red = w_base + for j in cutlass.range_constexpr(nt): + if wn * nt + j < cfg.ntg: + for i in cutlass.range_constexpr(mt): + for e in cutlass.range_constexpr(4): + token = (wn * nt + j) * _N_TILE + 2 * (lane & 3) + (e & 1) + row = (wm * mt + i) * _M_TILE + (lane >> 2) + (e >> 1) * 8 + _sts_f32( + red + ((wk * self.tg + token) * self.red_stride + row) * 4, + acc[(i * nt + j) * 4 + e], + ) + + +class K3KdaFabUmmaKernel(_FabTail): + """GEMM 1 on ``tcgen05`` UMMA for larger M; same hand-off and tail as :class:`K3KdaFabKernel`. + + CTA ``(g, s)`` owns 128 tokens (UMMA M: ``x`` is operand A) and every ``W_bfa`` row (UMMA N = + 144: operand B, rows past ``N1`` zero-filled by TMA) of the ``s``-th K slice in 64-wide + chunks. Warp 0 TMA-loads the CTA's whole weight slice before ``griddepcontrol.wait`` and + streams the activation chunks through a ring; warp 1 allocates tensor memory and issues the + UMMAs; warp 2 is the control warp; warps 4-7 copy the accumulator (one token per thread) into + the partial slots. + + Args: see :class:`K3KdaFabKernel` (``cfg.ntg`` must be 16 and ``cfg.r1`` 1). + """ + + _THREADS = 256 + _CONTROL_WARP = 2 + _EPI_WARP0 = 4 + + def __init__( + self, + cfg: FabConfig, + n_tiles: int, + n1: int, + n2: int, + k: int, + num_heads: int, + use_pdl: bool, + fa_mode: bool = False, + warm_w2: bool = True, + ): + validate_config(cfg, n_tiles, n1, n2, k, num_heads) + self._init_common( + cfg, + n_tiles, + n1, + n2, + k, + num_heads, + use_pdl, + _geometry(cfg, n_tiles, n2, k), + fa_mode, + warm_w2, + ) + self.mma_tiler = (_UMMA_M, MAX_N1, _K_CHUNK) + + @cute.jit + def __call__( + self, + w1_ptr: cute.Pointer, + x_ptr: cute.Pointer, + w2_ptr: cute.Pointer, + m: Int32, + w2_addr: Int64, + beta_addr: Int64, + fg_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + stream: cuda.CUstream, + ): + mW = cute.make_tensor( + w1_ptr, cute.make_layout((self.n1, self.k, 1), stride=(self.k, 1, self.n1 * self.k)) + ) + mX = cute.make_tensor( + x_ptr, cute.make_layout((m, self.k, 1), stride=(self.k, 1, m * self.k)) + ) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + cutlass.BFloat16, + cutlass.BFloat16, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + self.mma_tiler[:2], + ) + a_layout = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, cutlass.BFloat16, self.x_stages + ) + b_layout = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, cutlass.BFloat16, self.max_chunks + ) + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) + ) + tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mX, + cute.slice_(a_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mW, + cute.slice_(b_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_w2, tma_w2, w2_layout = self.make_w2_tma(w2_ptr) + + @cute.struct + class SharedStorage: + w_full: cute.struct.MemRange[cutlass.Int64, self.max_chunks] + w2_full: cute.struct.MemRange[cutlass.Int64, 1] + x_full: cute.struct.MemRange[cutlass.Int64, self.x_stages] + x_empty: cute.struct.MemRange[cutlass.Int64, self.x_stages] + acc_full: cute.struct.MemRange[cutlass.Int64, 1] + tmem_ready: cute.struct.MemRange[cutlass.Int64, 1] + tmem_free: cute.struct.MemRange[cutlass.Int64, 1] + tmem_holding: cute.struct.MemRange[cutlass.Int32, 1] + params: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, _NUM_PARAMS], 16] + # Weight chunks, then the activation ring; the tail reuses both. + sB: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, max(cute.cosize(b_layout.outer), self.sw_elems) + ], + 1024, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, max(cute.cosize(a_layout.outer), self.sx_elems) + ], + 1024, + ] + + self.shared_storage = SharedStorage + self.kernel( + tiled_mma, + tma_atom_a, + tma_a, + tma_atom_b, + tma_b, + a_layout, + b_layout, + tma_atom_w2, + tma_w2, + w2_layout, + m, + w2_addr, + beta_addr, + fg_addr, + ws_addr, + state_addr, + ).launch( + grid=(self.num_ctas, 1, 1), + block=(self._THREADS, 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB: cute.Tensor, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + tma_atom_w2: cute.CopyAtom, + mW2t: cute.Tensor, + w2_layout: cute.ComposedLayout, + m: Int32, + w2_addr: Int64, + beta_addr: Int64, + fg_addr: Int64, + ws_addr: Int64, + state_addr: Int64, + ): + S = self.cfg.s + stages = self.x_stages + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + group = bidx % self.groups + split = bidx // self.groups + base_chunks = self.num_k_chunks // S + rem_chunks = self.num_k_chunks % S + chunk_count = Int32(base_chunks) + chunk_start = split * base_chunks + if cutlass.const_expr(rem_chunks > 0): + if split < rem_chunks: + chunk_count = chunk_count + 1 + chunk_start = chunk_start + split + else: + chunk_start = chunk_start + rem_chunks + + storage = utils.SmemAllocator().allocate(self.shared_storage) + w_full = storage.w_full.data_ptr() + x_full = storage.x_full.data_ptr() + x_empty = storage.x_empty.data_ptr() + acc_full = storage.acc_full.data_ptr() + tmem_ready = storage.tmem_ready.data_ptr() + tmem_free = storage.tmem_free.data_ptr() + tmem_holding = storage.tmem_holding.data_ptr() + params = storage.params.data_ptr().toint() + if warp == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(self.max_chunks): + cute.arch.mbarrier_init(w_full + i, 1) + for i in cutlass.range_constexpr(stages): + cute.arch.mbarrier_init(x_full + i, 1) + cute.arch.mbarrier_init(x_empty + i, 1) + cute.arch.mbarrier_init(acc_full, 1) + cute.arch.mbarrier_init(tmem_ready, 1) + cute.arch.mbarrier_init(tmem_free, 4) + cute.arch.mbarrier_init(storage.w2_full.data_ptr(), 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + if warp == 0: + # Weight first (independent of the previous grid), then the activation ring. + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + gA = cute.local_tile( + mA, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB = cute.local_tile( + mB, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + thr_mma = tiled_mma.get_slice(0) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(thr_mma.partition_A(gA), 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(thr_mma.partition_B(gB), 0, 3), + ) + for c in cutlass.range(chunk_count, unroll=1): + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(w_full + c, self.w_chunk_bytes) + cute.copy( + tma_atom_b, + tBgB[(None, 0, chunk_start + c, 0)], + tBsB[(None, c)], + tma_bar_ptr=w_full + c, + cache_policy=Int64(_EVICT_FIRST), + ) + # Dependents may launch only once this CTA holds its tensor memory: a dependent CTA + # that took the columns first and then waited on this grid would deadlock it. + cute.arch.mbarrier_wait(tmem_ready, 0) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + _griddepcontrol_wait() + for c in cutlass.range(chunk_count, unroll=1): + stage = c % stages + if c >= stages: + cute.arch.mbarrier_wait(x_empty + stage, ((c // stages) & 1) ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(x_full + stage, self.x_chunk_bytes) + cute.copy( + tma_atom_a, + tAgA[(None, group, chunk_start + c, 0)], + tAsA[(None, stage)], + tma_bar_ptr=x_full + stage, + ) + elif warp == 1: + cute.arch.alloc_tmem(_UMMA_TMEM_COLS, tmem_holding) + cute.arch.relinquish_tmem_alloc_permit() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(tmem_ready) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=tmem_holding + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tiled_mma.make_fragment_C(acc_shape).layout) + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + num_kphases = cute.size(tCrA, mode=[2]) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for c in cutlass.range(chunk_count, unroll=1): + stage = c % stages + cute.arch.mbarrier_wait(w_full + c, 0) + cute.arch.mbarrier_wait(x_full + stage, (c // stages) & 1) + for kp in cutlass.range_constexpr(num_kphases): + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[(None, None, kp, stage)], + tCrB[(None, None, kp, c)], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + with cute.arch.elect_one(): + tcgen05.commit(x_empty + stage) + with cute.arch.elect_one(): + tcgen05.commit(acc_full) + cute.arch.mbarrier_wait(tmem_free, 0) + cute.arch.dealloc_tmem(acc_tmem_ptr, _UMMA_TMEM_COLS) + elif warp == self._CONTROL_WARP: + self.control_warp( + params, bidx, lane, m, w2_addr, beta_addr, fg_addr, ws_addr, state_addr + ) + self.refill(params, lane) + elif warp >= self._EPI_WARP0: + # TMEM lane = token: warp 4 + i reads lanes [32 i, 32 i + 32), all 144 columns. + cute.arch.mbarrier_wait(tmem_ready, 0) + tmem_base = _lds_b32(tmem_holding.toint()) + token = (warp - self._EPI_WARP0) * 32 + lane + lane_base = tmem_base + (((warp - self._EPI_WARP0) * 32) << 16) + # The accumulator goes through shared memory (the drained weight chunks) so the + # partial-tile stores below are whole lines: one token row of N1 floats per thread. + stage = storage.sB.data_ptr().toint() + token * (_UMMA_STAGE_PITCH * 4) + cute.arch.mbarrier_wait(acc_full, 0) + for j in cutlass.range_constexpr(MAX_N1 // 16): + v = _tmem_ld16(lane_base + j * 16) + for q in cutlass.range_constexpr(4): + if cutlass.const_expr(j * 16 + q * 4 < self.n1): + _sts_v4_f32( + stage + (j * 16 + q * 4) * 4, + v[4 * q], + v[4 * q + 1], + v[4 * q + 2], + v[4 * q + 3], + ) + _tmem_fence_before_sync() + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(tmem_free) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + cute.arch.barrier() + # The staged tile -> the CTA's partial slot, by every thread. + m = Int32(_lds_b64(params + _P_M * 8)) + half = _lds_b64(params + _P_HALF * 8) + units = self.n1 // 4 + rows = m - group * _UMMA_M + if rows > _UMMA_M: + rows = Int32(_UMMA_M) + part_base = half + Int64(((group * S + split) * _UMMA_M) * self.n1 * 4) + stage0 = storage.sB.data_ptr().toint() + for o in cutlass.range(tidx, rows * units, self._THREADS, unroll=4): + t = o // units + u = o % units + w = _lds_v4_f32(stage0 + (t * _UMMA_STAGE_PITCH + 4 * u) * 4) + _st_relaxed_v4_f32(part_base + Int64((t * self.n1 + 4 * u) * 4), *w) + cute.arch.barrier() + m = Int32(_lds_b64(params + _P_M * 8)) + half = _lds_b64(params + _P_HALF * 8) + worker = Int32(_lds_b64(params + _P_RANK * 8)) - (self.num_ctas - self.items) + region = storage.sB.data_ptr().toint() + sW2 = storage.sB.get_tensor(w2_layout.outer, swizzle=w2_layout.inner) + if cutlass.const_expr(self.fa_mode): + self.tail_fa(params, tidx, worker, m, half) + elif worker >= 0: + self.tail( + region, + sW2, + storage.w2_full.data_ptr(), + tma_atom_w2, + mW2t, + params, + tidx, + warp, + lane, + worker, + m, + half, + ) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def _geometry(cfg: FabConfig, n_tiles: int, n2: int, k: int) -> dict: + tg = cfg.ntg * _N_TILE + groups = _ceil_div(n_tiles, cfg.ntg) + num_k_chunks = k // _K_CHUNK + max_chunks = _ceil_div(num_k_chunks, cfg.s) + token_blocks = _ceil_div(n_tiles, cfg.ntb) + row_blocks = _ceil_div(n2 // _M_TILE, cfg.mtb) + w2_smem_bytes = cfg.mtb * _M_TILE * FA_DIM * 2 + w2_box_rows = min(128, cfg.mtb * _M_TILE) + fa_smem_bytes = cfg.ntb * _N_TILE * FA_DIM * 2 + out_smem_bytes = cfg.ntb * _N_TILE * (cfg.mtb * _M_TILE + 8) * 2 + geo = { + "tg": tg, + "groups": groups, + "num_k_chunks": num_k_chunks, + "max_chunks": max_chunks, + "token_blocks": token_blocks, + "row_blocks": row_blocks, + "items": token_blocks * row_blocks, + "num_ctas": groups * cfg.r1 * cfg.s, + "w2_smem_bytes": w2_smem_bytes, + "w2_box_rows": w2_box_rows, + "w2_boxes": (cfg.mtb * _M_TILE // w2_box_rows) * (FA_DIM // _K_CHUNK), + "workspace_bytes": groups * cfg.s * tg * MAX_N1 * 4, + } + if cfg.umma: + w_chunk_bytes = MAX_N1 * _ROW_BYTES + x_chunk_bytes = _UMMA_M * _ROW_BYTES + sw_bytes = max_chunks * w_chunk_bytes + budget = _SMEM_CAPACITY - sw_bytes - 4096 + x_stages = max(1, min(max_chunks, budget // x_chunk_bytes)) + sx_bytes = x_stages * x_chunk_bytes + alias_need = max( + w2_smem_bytes + fa_smem_bytes + out_smem_bytes, _UMMA_M * _UMMA_STAGE_PITCH * 4 + ) + sw_bytes = max(sw_bytes, _ceil_div(max(alias_need - sx_bytes, 0), 1024) * 1024) + geo.update( + rows_cta=MAX_N1, + num_mma_warps=0, + num_warps=K3KdaFabUmmaKernel._THREADS // 32, + threads=K3KdaFabUmmaKernel._THREADS, + x_stages=x_stages, + w_chunk_bytes=w_chunk_bytes, + x_chunk_bytes=x_chunk_bytes, + sw_elems=sw_bytes // 2, + sx_elems=sx_bytes // 2, + smem_bytes=sw_bytes + + sx_bytes + + 8 * (max_chunks + 2 * x_stages + 4) + + 8 * _NUM_PARAMS + + 2048, + ) + return geo + mtc = _N1_TILES // cfg.r1 + rows_cta = mtc * _M_TILE + warps_m = mtc // cfg.mt + warps_n = _ceil_div(cfg.ntg, cfg.nt) + num_mma_warps = warps_m * warps_n * cfg.kg + red_stride = rows_cta + 4 + w_chunk_bytes = rows_cta * _ROW_BYTES + x_chunk_bytes = tg * _ROW_BYTES + red_bytes = cfg.kg * tg * red_stride * 4 + sw_bytes = max_chunks * w_chunk_bytes + sx_bytes = max_chunks * x_chunk_bytes + alias_need = max(red_bytes, w2_smem_bytes + fa_smem_bytes + out_smem_bytes) + sw_bytes = max(sw_bytes, _ceil_div(max(alias_need - sx_bytes, 0), 1024) * 1024) + geo.update( + mtc=mtc, + rows_cta=rows_cta, + warps_m=warps_m, + warps_n=warps_n, + num_mma_warps=num_mma_warps, + # One control warp (ticket, parameters, workspace refill) besides the MMA warps. + num_warps=num_mma_warps + 1, + threads=32 * (num_mma_warps + 1), + red_stride=red_stride, + w_chunk_bytes=w_chunk_bytes, + x_chunk_bytes=x_chunk_bytes, + sw_elems=sw_bytes // 2, + sx_elems=sx_bytes // 2, + smem_bytes=sw_bytes + sx_bytes + 8 * max_chunks + 8 * _NUM_PARAMS + 2048, + ) + return geo + + +def validate_config(cfg: FabConfig, n_tiles: int, n1: int, n2: int, k: int, num_heads: int) -> None: + """Raise ``ValueError`` if ``cfg`` cannot run the given problem.""" + if not 1 <= n_tiles <= MAX_M // _N_TILE: + raise ValueError(f"n_tiles={n_tiles} must be in [1, {MAX_M // _N_TILE}]") + if n1 % 8 or not FA_DIM < n1 <= MAX_N1 or not 1 <= num_heads <= n1 - FA_DIM: + raise ValueError(f"N1={n1} / num_heads={num_heads} unsupported") + if n2 % _M_TILE or k % _K_CHUNK: + raise ValueError(f"N2={n2} must be a multiple of 16 and K={k} of 64") + if cfg.umma: + if cfg.ntg != _UMMA_M // _N_TILE or cfg.r1 != 1: + raise ValueError(f"UMMA tiling needs ntg={_UMMA_M // _N_TILE} and r1=1: {cfg}") + elif cfg.r1 not in (1, 3, 9) or (_N1_TILES // cfg.r1) % cfg.mt: + raise ValueError(f"r1={cfg.r1} / mt={cfg.mt} invalid") + elif cfg.ntg < 1 or cfg.nt < 1 or cfg.nt > cfg.ntg or cfg.kg < 1: + raise ValueError(f"invalid token tiling {cfg}") + if cfg.ntb < 1 or cfg.ntg % cfg.ntb or cfg.mtb < 1: + raise ValueError(f"invalid tail tiling {cfg}") + if not 1 <= cfg.s <= k // _K_CHUNK: + raise ValueError(f"s={cfg.s} invalid for K={k}") + geo = _geometry(cfg, n_tiles, n2, k) + if (cfg.mtb * _M_TILE) % min(128, cfg.mtb * _M_TILE): + raise ValueError(f"{cfg}: the W_fb block must be whole 128-row boxes") + if cfg.s * cfg.r1 > 32: + raise ValueError(f"{cfg}: more than 32 producer slots per token group") + if geo["num_warps"] > 16: + raise ValueError(f"{cfg} needs {geo['num_warps']} warps") + if geo["items"] > geo["num_ctas"]: + raise ValueError(f"{cfg}: {geo['items']} tail items > {geo['num_ctas']} CTAs") + if geo["smem_bytes"] > _SMEM_CAPACITY: + raise ValueError(f"{cfg} needs {geo['smem_bytes']} B of shared memory") + if geo["workspace_bytes"] > _WS_HALF_BYTES: + raise ValueError(f"{cfg} needs {geo['workspace_bytes']} B of workspace") + + +def _base_config(n_tiles: int) -> FabConfig: + if n_tiles == 1: + return FabConfig(ntg=1, r1=3, s=9, mt=1, nt=1, kg=4, ntb=1, mtb=2) + if n_tiles == 2: + return FabConfig(ntg=2, r1=3, s=9, mt=1, nt=2, kg=4, ntb=1, mtb=4) + if n_tiles <= 4: + return FabConfig(ntg=n_tiles, r1=3, s=9, mt=3, nt=2, kg=4, ntb=1, mtb=8) + if n_tiles <= 8: + return FabConfig(ntg=n_tiles, r1=3, s=9, mt=3, nt=4, kg=2, ntb=1, mtb=16) + umma = {"ntg": _UMMA_M // _N_TILE, "r1": 1, "mt": 1, "nt": 1, "kg": 1, "umma": True} + if n_tiles <= 16: + return FabConfig(s=28, ntb=1 if n_tiles <= 14 else 2, mtb=24, **umma) + return FabConfig(s=14, ntb=1 if n_tiles <= 28 else 2, mtb=48, **umma) + + +@functools.cache +def default_config(n_tiles: int, n2: int = 768, k: int = 7168) -> FabConfig: + """Tiling for ``ceil(M / 8) = n_tiles``, tuned on GB300 next to the concurrent ``qkvg`` GEMM. + + The concurrent ``qkvg`` GEMM (``k3_fp8w_lowm_gemm``, clusters of 4-5 CTAs) occupies 96-120 of + the 152 SMs; more than ~28 CTAs here break its single-wave cluster placement. So the grid is + 27 CTAs (``mma.sync``, M <= 64) or 28 CTAs (UMMA, larger M), with at most one tail item each. + Other shapes keep the tiling with at most ``k / 64`` K splits and enough ``W_fb`` rows per + tail item to cover ``n2`` with the CTAs available. + """ + cfg = _base_config(n_tiles) + cfg = dataclasses.replace(cfg, s=min(cfg.s, k // _K_CHUNK)) + geo = _geometry(cfg, n_tiles, n2, k) + row_blocks = max(1, geo["num_ctas"] // geo["token_blocks"]) + mtb = max(cfg.mtb, _ceil_div(n2 // _M_TILE, row_blocks)) + if mtb > 128 // _M_TILE: + mtb = _ceil_div(mtb, 128 // _M_TILE) * (128 // _M_TILE) + return dataclasses.replace(cfg, mtb=mtb) + + +@functools.cache +def _compiled( + cfg: FabConfig, + n_tiles: int, + n1: int, + n2: int, + k: int, + num_heads: int, + use_pdl: bool, + fa_mode: bool = False, + warm_w2: bool = True, +): + kernel_cls = K3KdaFabUmmaKernel if cfg.umma else K3KdaFabKernel + kernel = kernel_cls(cfg, n_tiles, n1, n2, k, num_heads, use_pdl, fa_mode, warm_w2) + dummy = 1 << 20 + ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + addr = Int64(dummy) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + kernel, ptr, ptr, ptr, Int32(8 * n_tiles), addr, addr, addr, addr, addr, stream + ) + + +_STATE: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + + +def _device_index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() + + +def _device_state(device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """Per-device ``(state words, sentinel-filled partial workspace)``. + + Created outside CUDA graph capture (a capture would allocate it from the graph's private + pool); :func:`can_use_k3_kda_fab` rejects capturing callers until it exists. + """ + index = _device_index(device) + state = _STATE.get(index) + if state is None: + words = torch.zeros(_STATE_WORDS, dtype=torch.int64, device=device) + ws = torch.full((2 * _WS_HALF_BYTES // 4,), _SENTINEL, dtype=torch.int32, device=device) + state = (words, ws) + _STATE[index] = state + return state + + +@functools.cache +def _num_sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +@functools.cache +def _is_supported_device(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3), (10, 7)) + + +def _check_inputs(x: torch.Tensor, w_bfa: torch.Tensor, w_fb: torch.Tensor, num_heads: int): + for name, t in (("x", x), ("w_bfa", w_bfa), ("w_fb", w_fb)): + if t.dtype != torch.bfloat16 or t.dim() != 2 or not t.is_contiguous(): + raise ValueError(f"{name} must be a contiguous 2D bfloat16 tensor") + if not t.is_cuda or t.device != x.device or t.data_ptr() % 16: + raise ValueError(f"{name} must be a 16-byte aligned tensor on x's CUDA device") + m, k = x.shape + n1 = w_bfa.shape[0] + n2 = w_fb.shape[0] + if w_bfa.shape[1] != k or w_fb.shape[1] != FA_DIM: + raise ValueError( + f"shape mismatch: x {tuple(x.shape)}, w_bfa {tuple(w_bfa.shape)}, " + f"w_fb {tuple(w_fb.shape)}" + ) + if not 1 <= m <= MAX_M: + raise ValueError(f"M={m} must be in [1, {MAX_M}]") + if n1 % 8 or not FA_DIM < n1 <= MAX_N1 or not 1 <= num_heads <= n1 - FA_DIM: + raise ValueError(f"N1={n1} / num_heads={num_heads} unsupported") + if n2 % _M_TILE or k % _K_CHUNK: + raise ValueError(f"N2={n2} must be a multiple of 16 and K={k} of 64") + + +def run_k3_kda_fab( + x: torch.Tensor, + w_bfa: torch.Tensor, + w_fb: torch.Tensor, + num_heads: int, + beta: torch.Tensor, + forget_gate: torch.Tensor, + *, + config: FabConfig | None = None, + use_pdl: bool = True, +) -> None: + """Launch into ``beta`` (``bf16 [M, H]``) and ``forget_gate`` (``bf16 [M, N2]``).""" + m, k = x.shape + n_tiles = _ceil_div(m, _N_TILE) + cfg = config if config is not None else default_config(n_tiles, w_fb.shape[0], k) + compiled = _compiled(cfg, n_tiles, w_bfa.shape[0], w_fb.shape[0], k, num_heads, use_pdl) + words, ws = _device_state(x.device) + compiled( + make_ptr(cutlass.BFloat16, w_bfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, w_fb.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int32(m), + Int64(w_fb.data_ptr()), + Int64(beta.data_ptr()), + Int64(forget_gate.data_ptr()), + Int64(ws.data_ptr()), + Int64(words.data_ptr()), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + + +def run_k3_kda_fab_fa( + x: torch.Tensor, + w_bfa: torch.Tensor, + num_heads: int, + bfa: torch.Tensor, + w_fb: torch.Tensor | None = None, + *, + config: FabConfig | None = None, + use_pdl: bool = True, +) -> None: + """f_a mode: launch into ``bfa`` (``bf16 [M, N1]`` contiguous): ``bf16(x @ W_bfa^T)``, bitwise + the ``f_a`` and ``beta`` of :func:`run_k3_kda_fab` with the same ``config``. ``w_fb`` + (``bf16 [N2, 128]``, optional) is only warmed in L2 for the decode that applies it; ``N2`` + (768 without it) also selects the default tiling like in the full kernel.""" + m, k = x.shape + n_tiles = _ceil_div(m, _N_TILE) + n2 = w_fb.shape[0] if w_fb is not None else 768 + cfg = config if config is not None else default_config(n_tiles, n2, k) + compiled = _compiled( + cfg, n_tiles, w_bfa.shape[0], n2, k, num_heads, use_pdl, True, w_fb is not None + ) + words, ws = _device_state(x.device) + w1 = make_ptr(cutlass.BFloat16, w_bfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + w2_addr = w_fb.data_ptr() if w_fb is not None else w_bfa.data_ptr() + compiled( + w1, + make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + # The W_fb TMA operand is unused in f_a mode: any 16-byte aligned buffer of at least + # N2 x 128 BF16 (W_bfa without w_fb). + w1, + Int32(m), + Int64(w2_addr), + Int64(bfa.data_ptr()), + Int64(bfa.data_ptr()), + Int64(ws.data_ptr()), + Int64(words.data_ptr()), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + + +def k3_kda_fab_fa( + x: torch.Tensor, w_bfa: torch.Tensor, num_heads: int, w_fb: Optional[torch.Tensor] = None +) -> torch.Tensor: + """``bfa = bf16(x @ w_bfa^T)`` (``bf16 [M, N1]``) by the fused fab kernel's GEMM 1 and + reduction (f_a mode): columns ``[0, 128)`` are ``f_a``, ``[128, 128 + num_heads)`` ``beta``, + bitwise as inside :func:`k3_kda_fab_proj` (same ``W_fb`` row count). ``w_fb`` is only warmed + in L2. Accepts what :func:`can_use_k3_kda_fab_fa` does.""" + _check_fa_inputs(x, w_bfa, num_heads, w_fb) + bfa = torch.empty((x.shape[0], w_bfa.shape[0]), dtype=torch.bfloat16, device=x.device) + run_k3_kda_fab_fa(x, w_bfa, num_heads, bfa, w_fb) + return bfa + + +def _check_fa_inputs( + x: torch.Tensor, w_bfa: torch.Tensor, num_heads: int, w_fb: torch.Tensor | None = None +) -> None: + """:func:`_check_inputs` with an optional ``W_fb`` operand.""" + if w_fb is not None: + _check_inputs(x, w_bfa, w_fb, num_heads) + return + for name, t in (("x", x), ("w_bfa", w_bfa)): + if t.dtype != torch.bfloat16 or t.dim() != 2 or not t.is_contiguous(): + raise ValueError(f"{name} must be a contiguous 2D bfloat16 tensor") + if not t.is_cuda or t.device != x.device or t.data_ptr() % 16: + raise ValueError(f"{name} must be a 16-byte aligned tensor on x's CUDA device") + m, k = x.shape + n1 = w_bfa.shape[0] + if w_bfa.shape[1] != k or not 1 <= m <= MAX_M or k % _K_CHUNK: + raise ValueError(f"unsupported shapes: x {tuple(x.shape)}, w_bfa {tuple(w_bfa.shape)}") + if n1 % 8 or not FA_DIM < n1 <= MAX_N1 or not 1 <= num_heads <= n1 - FA_DIM: + raise ValueError(f"N1={n1} / num_heads={num_heads} unsupported") + + +def can_use_k3_kda_fab_fa( + x: torch.Tensor, w_bfa: torch.Tensor, num_heads: int, w_fb: torch.Tensor | None = None +) -> bool: + """Whether :func:`k3_kda_fab_fa` accepts these inputs (non-raising check).""" + if not x.is_cuda or not _is_supported_device(x.device.index): + return False + if torch.cuda.is_current_stream_capturing() and _device_index(x.device) not in _STATE: + return False + try: + _check_fa_inputs(x, w_bfa, num_heads, w_fb) + n_tiles = _ceil_div(x.shape[0], _N_TILE) + k = x.shape[1] + n2 = w_fb.shape[0] if w_fb is not None else 768 + validate_config(default_config(n_tiles, n2, k), n_tiles, w_bfa.shape[0], n2, k, num_heads) + except ValueError: + return False + return True + + +def k3_kda_fab_proj( + x: torch.Tensor, w_bfa: torch.Tensor, w_fb: torch.Tensor, num_heads: int +) -> tuple[torch.Tensor, torch.Tensor]: + """``(beta, forget_gate)`` of the KDA ``[f_a | b] -> f_b`` projection in one launch. + + Args: + x: ``bf16 [M, K]``, ``1 <= M <= 256``, ``K`` a multiple of 64. + w_bfa: ``bf16 [N1, K]``: ``f_a_proj`` rows ``[0, 128)``, ``b_proj`` rows + ``[128, 128 + num_heads)``, zero padding (``N1`` a multiple of 8, at most 144). + w_fb: ``bf16 [N2, 128]`` (``f_b_proj``), ``N2`` a multiple of 16. + num_heads: Rows of ``b_proj``. + + Returns: + ``beta`` ``bf16 [M, num_heads]`` and ``forget_gate`` ``bf16 [M, N2]``, equal to + ``bfa = bf16(x @ w_bfa^T)``, ``bfa[:, 128:128 + num_heads]`` and + ``bf16(bfa[:, :128] @ w_fb^T)`` (FP32 accumulation). + """ + _check_inputs(x, w_bfa, w_fb, num_heads) + beta = torch.empty((x.shape[0], num_heads), dtype=torch.bfloat16, device=x.device) + forget_gate = torch.empty((x.shape[0], w_fb.shape[0]), dtype=torch.bfloat16, device=x.device) + run_k3_kda_fab(x, w_bfa, w_fb, num_heads, beta, forget_gate) + return beta, forget_gate + + +def can_use_k3_kda_fab( + x: torch.Tensor, w_bfa: torch.Tensor, w_fb: torch.Tensor, num_heads: int +) -> bool: + """Whether :func:`k3_kda_fab_proj` accepts these inputs (non-raising check).""" + if not x.is_cuda or not _is_supported_device(x.device.index): + return False + if torch.cuda.is_current_stream_capturing() and _device_index(x.device) not in _STATE: + return False + try: + _check_inputs(x, w_bfa, w_fb, num_heads) + n_tiles = _ceil_div(x.shape[0], _N_TILE) + n2, k = w_fb.shape[0], x.shape[1] + validate_config(default_config(n_tiles, n2, k), n_tiles, w_bfa.shape[0], n2, k, num_heads) + except ValueError: + return False + return True + + +__all__ = [ + "FA_DIM", + "FabConfig", + "K3KdaFabKernel", + "MAX_M", + "MAX_N1", + "can_use_k3_kda_fab", + "can_use_k3_kda_fab_fa", + "default_config", + "k3_kda_fab_fa", + "k3_kda_fab_proj", + "run_k3_kda_fab", + "run_k3_kda_fab_fa", + "validate_config", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_l2_prefetch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_l2_prefetch.py new file mode 100644 index 000000000000..8f1cfeaaf66d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_l2_prefetch.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""L2 prefetch of weight byte ranges (SM100/SM103). + +A small kernel that issues ``cp.async.bulk.prefetch.L2`` for up to :data:`MAX_RANGES` +contiguous byte ranges and exits. The prefetches are hints: nothing waits for them and no +result depends on them. The Kimi-K3 decode uses it on a side stream to pull the next layer's +projection weights into L2 while the current layer's latency-bound kernels (recurrence, +all-reduce) leave HBM idle, so the next layer's GEMMs read them from L2 (their weight loads carry +an evict-first hint, so prefetched lines with the default policy survive those weight streams). + +Each thread prefetches whole chunks of :data:`CHUNK_BYTES` (the last chunk of a range may be +shorter); one warp per CTA, as many CTAs as it takes to give every thread about one chunk, at +most ``max_ctas``. A CTA stays resident until its prefetches have been served, and the prefetch +rate grows with the number of SMs that issue them (about 4 TB/s from 64-120 CTAs). + +Usage notes: + +* Launch it on a side stream behind an event recorded after the last kernel that streams weights + from HBM, without the PDL attribute: its dependencies then stay full ones and it does not + compete with them for HBM. +* It triggers its programmatic dependents at once, but a programmatic dependent that also waits + for another stream loses its early launch: join the side stream where no kernel's head start + matters (or let a later user of the stream join it). +""" + +from __future__ import annotations + +import functools +from typing import Sequence + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op + +#: Largest number of byte ranges per launch. +MAX_RANGES = 8 +#: Bytes per ``cp.async.bulk.prefetch.L2`` instruction. +CHUNK_BYTES = 8 << 10 +#: Default upper bound on the CTAs (one warp each) of a launch. +DEFAULT_MAX_CTAS = 64 +_THREADS = 32 + + +@dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "griddepcontrol.launch_dependents;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_prefetch_l2(addr, nbytes, *, loc=None, ip=None): + llvm.inline_asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Int32(nbytes).ir_value(loc=loc, ip=ip)], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class K3L2PrefetchKernel: + """``ctas`` CTAs of one warp prefetch the first ``num_ranges`` of :data:`MAX_RANGES` + ``(address, bytes)`` ranges in :data:`CHUNK_BYTES` chunks (addresses 16-byte aligned, sizes + multiples of 16).""" + + def __init__(self, num_ranges: int, ctas: int): + if not 1 <= num_ranges <= MAX_RANGES or ctas < 1: + raise ValueError(f"num_ranges={num_ranges} ctas={ctas}") + self.num_ranges = num_ranges + self.ctas = ctas + + @cute.jit + def __call__( + self, + a0: Int64, + n0: Int64, + a1: Int64, + n1: Int64, + a2: Int64, + n2: Int64, + a3: Int64, + n3: Int64, + a4: Int64, + n4: Int64, + a5: Int64, + n5: Int64, + a6: Int64, + n6: Int64, + a7: Int64, + n7: Int64, + stream: cuda.CUstream, + ): + self.kernel(a0, n0, a1, n1, a2, n2, a3, n3, a4, n4, a5, n5, a6, n6, a7, n7).launch( + grid=(self.ctas, 1, 1), block=(_THREADS, 1, 1), stream=stream + ) + + @cute.kernel + def kernel( + self, + a0: Int64, + n0: Int64, + a1: Int64, + n1: Int64, + a2: Int64, + n2: Int64, + a3: Int64, + n3: Int64, + a4: Int64, + n4: Int64, + a5: Int64, + n5: Int64, + a6: Int64, + n6: Int64, + a7: Int64, + n7: Int64, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + # Nothing reads what this kernel does: a programmatic dependent (e.g. a kernel that joins + # this stream) need not wait for the prefetches to land. + _griddepcontrol_launch_dependents() + thread = bidx * _THREADS + tidx + total = self.ctas * _THREADS + ranges = ((a0, n0), (a1, n1), (a2, n2), (a3, n3), (a4, n4), (a5, n5), (a6, n6), (a7, n7)) + for r in cutlass.range_constexpr(self.num_ranges): + base, nbytes = ranges[r] + chunks = Int32((nbytes + (CHUNK_BYTES - 1)) // CHUNK_BYTES) + for c in cutlass.range(thread, chunks, total, unroll=1): + off = Int64(c) * CHUNK_BYTES + size = nbytes - off + if size > CHUNK_BYTES: + size = Int64(CHUNK_BYTES) + _bulk_prefetch_l2(base + off, Int32(size)) + + +@functools.cache +def _compiled(num_ranges: int, ctas: int): + zero = Int64(0) + return cute.compile( + K3L2PrefetchKernel(num_ranges, ctas), + *([zero] * (2 * MAX_RANGES)), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def tensor_range(t: torch.Tensor) -> tuple[int, int]: + """``(address, bytes)`` spanned by a dense tensor (contiguous up to a permutation of its + dimensions, e.g. a transposed scale buffer), aligned inwards to 16 bytes.""" + if t.numel() == 0: + return t.data_ptr(), 0 + if any(s < 0 for s in t.stride()): + raise ValueError("prefetch ranges need non-negative strides") + span = (1 + sum((n - 1) * s for n, s in zip(t.shape, t.stride()))) * t.element_size() + if span != t.numel() * t.element_size(): + raise ValueError("prefetch ranges must be dense tensors") + start = t.data_ptr() + first = (start + 15) // 16 * 16 + return first, max(0, (start + span - first) // 16 * 16) + + +def launch_ctas(ranges: Sequence[tuple[int, int]], max_ctas: int = DEFAULT_MAX_CTAS) -> int: + """CTAs that give every thread about one chunk (at least 1, at most ``max_ctas``): the + prefetch rate scales with the SMs issuing it (each SM's TMA unit handles its own requests).""" + chunks = sum((n + CHUNK_BYTES - 1) // CHUNK_BYTES for _, n in ranges) + return max(1, min(max_ctas, (chunks + _THREADS - 1) // _THREADS)) + + +def run_k3_l2_prefetch( + ranges: Sequence[tuple[int, int]], *, max_ctas: int = DEFAULT_MAX_CTAS +) -> None: + """Prefetch ``ranges`` (``(address, bytes)``, 16-byte aligned, at most :data:`MAX_RANGES`) + into L2 on the current stream. Empty ranges are dropped; nothing is launched if none is + left.""" + ranges = [(a, n) for a, n in ranges if n > 0] + if not ranges: + return + if len(ranges) > MAX_RANGES: + raise ValueError(f"at most {MAX_RANGES} ranges, got {len(ranges)}") + for a, n in ranges: + if a % 16 or n % 16: + raise ValueError(f"range ({a:#x}, {n}) is not 16-byte aligned") + args = [] + for a, n in ranges: + args += [Int64(a), Int64(n)] + args += [Int64(0)] * (2 * MAX_RANGES - len(args)) + _compiled(len(ranges), launch_ctas(ranges, max_ctas))( + *args, cuda.CUstream(torch.cuda.current_stream().cuda_stream) + ) + + +__all__ = ["CHUNK_BYTES", "MAX_RANGES", "K3L2PrefetchKernel", "run_k3_l2_prefetch", "tensor_range"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_lamport_ar.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_lamport_ar.py new file mode 100644 index 000000000000..ce5a89d38e9d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_lamport_ar.py @@ -0,0 +1,1081 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MNNVL Lamport all-reduce whose input the producer pushes (SM100 / SM103 / SM107). + +The MNNVL all-reduce (``trtllm::mnnvl_fusion_allreduce``, ``mnnvlAllreduceKernels.cu``) starts +by copying every rank's BF16 input row slices into Lamport buffers: in the two-shot kernel each +token row goes to slot ``(t // W, rank)`` of the scatter stage of its owner ``t % W``, in the +one-shot kernel to slot ``(t, rank)`` of every rank's buffer (multicast). This module splits +that copy off: a producer stores its output straight into those slots (the K3 low-M FP8 GEMM's +``ar_push`` epilogue, or :func:`run_lamport_push` for any BF16 tensor), and +:class:`K3LamportFinishKernel` runs the rest of the all-reduce on the same workspace: + +* two-shot: the owner of token ``t`` polls its ``W`` slots until no 32-bit word holds the + sentinel ``0x80000000``, sums them in FP32 in rank order, rounds to BF16, replaces ``-0`` by + ``+0`` and multicasts the row into the broadcast stage of every rank; every rank then polls the + broadcast rows it needs and copies them to the output; +* one-shot: every rank polls the ``W`` slots of each token, sums them (chunks of eight ranks + above eight ranks, like the C++ kernel) and rounds to BF16. + +The finish kernel clears the previous call's buffer back to the sentinel and advances the +workspace's rotation state (current / dirty buffer, bytes to clear, CTA arrival counter) exactly +like the C++ kernels, so fused and C++ calls can alternate on one workspace. Values are +sanitized (``-0`` sent as ``+0``) before they enter a Lamport buffer, which the sentinel needs; +the sums use the C++ order and round to nearest even, so the output is bitwise that of the C++ +all-reduce of the same input. Either transport can reproduce either C++ kernel's result (see +:class:`K3LamportFinishKernel`). + +Ordering (details in :class:`K3LamportFinishKernel`): the producer reads the rotation state +after its own ``griddepcontrol.wait``; the finish kernel's workers read it and poll without +waiting for the producer, and only its updater CTA waits for the producer before it advances +the state. Calls that share a workspace must be ordered by streams (or events) on every rank, +every rank must make the same sequence of calls, and every PDL kernel between two calls must +execute ``griddepcontrol.wait``. + +A consumer kernel can also finish a call itself after its ``griddepcontrol.wait`` +(``k3_attn_res`` with ``lamport_finish``): :func:`poll_reduce_v4` sums a token's one-shot slots +in the finish kernel's order, :func:`finish_two_shot_v4` reduces and multicasts an owned two-shot +row or polls another rank's, :func:`clear_dirty_stages` re-arms the previous call's buffer and +:func:`arrive_finish` advances the rotation state once every CTA has read it. + +All-gather row (``ag``, one-shot transport only): a call of ``m`` rows can carry, after its +``m x W`` slots, ``m`` more rows of the same width that the ranks fill column-wise (each rank +multicasts its own columns of every row; :func:`ag_row_addr`). The finish then returns, per +element, ``bf16(y + bf16(sum of the slots))`` with ``y`` the all-gathered value: the BF16 add of +a column-parallel projection's output and an all-reduced partial (the K3 MoE tail with +latent_up split by output columns). The rotation state then re-arms ``m x (W + 1)`` rows. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Threads per CTA of the finish and push kernels; one thread owns one 16-byte vector. +THREADS = 128 +#: 32-bit word of an empty Lamport slot (``-0.0`` FP32; a BF16 pair ``(+0, -0)``). +SENTINEL = -2147483648 # 0x80000000 +#: Number of Lamport buffers in a workspace (rotation period). +NUM_BUFFERS = 3 +#: Word of the CTA arrival counter in the workspace's state flags (uint32 [9]: current buffer, +#: dirty buffer, bytes per buffer, dirty stages, bytes to clear per stage x4, arrivals). +_FLAG_ARRIVALS = 8 + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _lv(value, *, loc=None, ip=None): + return Int64(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return cutlass.Float32(value).ir_value(loc=loc, ip=ip) + + +def _unpack(res, ty, cls, n, *, loc=None, ip=None): + return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n)) + + +@dsl_user_op +def ld_volatile_v4(addr, *, loc=None, ip=None): + """16 bytes of a Lamport buffer (or the state flags), read past L1.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.volatile.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def ld_nc_u64(addr, *, loc=None, ip=None): + """A 64-bit word of a tensor no kernel writes (the peer pointer table).""" + return Int64( + _asm( + T.i64(), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.u64 $0, [$1];", + "=l,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ld_v4(addr, *, loc=None, ip=None): + """16 bytes of a tensor the previous grid wrote (read after ``griddepcontrol.wait``).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def st_v4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + """16-byte store to local memory or to a peer's unicast mapping.""" + _asm( + None, + [_lv(addr, loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def multimem_st_v4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + """16-byte store to the address in every rank's buffer (multicast mapping).""" + _asm( + None, + [_lv(addr, loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "multimem.st.relaxed.sys.global.v4.f32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def st_sentinel_v4(addr, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)], + "{\n.reg .b32 t;\nmov.b32 t, 0x80000000;\nst.global.v4.b32 [$0], {t, t, t, t};\n}", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def clean4(w0, w1, w2, w3, *, loc=None, ip=None): + """1 if none of the four words is the sentinel (the vector has arrived).""" + return Int32( + _asm( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)], + "{\n.reg .pred p;\n" + "setp.ne.b32 p, $1, 0x80000000;\nsetp.ne.and.b32 p, $2, 0x80000000, p;\n" + "setp.ne.and.b32 p, $3, 0x80000000, p;\nsetp.ne.and.b32 p, $4, 0x80000000, p;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def sanitize_bf16x2(word, *, loc=None, ip=None): + """Replace each ``-0`` half (0x8000) of a BF16 pair by ``+0``.""" + return Int32( + _asm( + T.i32(), + [_iv(word, loc=loc, ip=ip)], + "{\n.reg .pred p, q;\n.reg .b32 l, h;\n" + "and.b32 l, $1, 0xffff;\nsetp.eq.b32 p, l, 0x8000;\nselp.b32 l, 0, l, p;\n" + "and.b32 h, $1, 0xffff0000;\nsetp.eq.b32 q, h, 0x80000000;\nselp.b32 h, 0, h, q;\n" + "or.b32 $0, l, h;\n}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def acc_bf16x2(word, acc_lo, acc_hi, *, loc=None, ip=None): + """``(acc_lo + lo, acc_hi + hi)`` in FP32 (round to nearest) for the BF16 pair ``word``.""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_iv(word, loc=loc, ip=ip), _fv(acc_lo, loc=loc, ip=ip), _fv(acc_hi, loc=loc, ip=ip)], + "{\n.reg .b32 l, h;\n.reg .f32 fl, fh;\n" + "shl.b32 l, $2, 16;\nand.b32 h, $2, 0xffff0000;\n" + "mov.b32 fl, l;\nmov.b32 fh, h;\n" + "add.rn.f32 $0, $3, fl;\nadd.rn.f32 $1, $4, fh;\n}", + "=f,=f,r,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), cutlass.Float32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def add_f32(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def pack_bf16x2(lo, hi, *, loc=None, ip=None): + """Round two FP32 values to BF16 (nearest even), ``lo`` in the low half.""" + return Int32( + _asm( + T.i32(), + [_fv(lo, loc=loc, ip=ip), _fv(hi, loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ag_add_bf16x2(acc_lo, acc_hi, yword, *, loc=None, ip=None): + """``bf16(y + bf16(acc))`` for the pair ``(acc_lo, acc_hi)`` of FP32 slot sums and the BF16 + pair ``yword`` (round to nearest even throughout), packed with ``lo`` in the low half.""" + return Int32( + _asm( + T.i32(), + [_fv(acc_lo, loc=loc, ip=ip), _fv(acc_hi, loc=loc, ip=ip), _iv(yword, loc=loc, ip=ip)], + "{\n.reg .b32 s, sl, sh, yl, yh;\n.reg .f32 fsl, fsh, fyl, fyh;\n" + "cvt.rn.bf16x2.f32 s, $2, $1;\n" + "shl.b32 sl, s, 16;\nand.b32 sh, s, 0xffff0000;\n" + "shl.b32 yl, $3, 16;\nand.b32 yh, $3, 0xffff0000;\n" + "mov.b32 fsl, sl;\nmov.b32 fsh, sh;\nmov.b32 fyl, yl;\nmov.b32 fyh, yh;\n" + "add.rn.f32 fyl, fyl, fsl;\nadd.rn.f32 fyh, fyh, fsh;\n" + "cvt.rn.bf16x2.f32 $0, fyh, fyl;\n}", + "=r,f,f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def red_release_add(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def atom_add_acq_rel_u32(addr, value, *, loc=None, ip=None): + """``atomicAdd`` (acquire-release, GPU scope); returns the old value.""" + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def atom_add_relaxed_u32(addr, value, *, loc=None, ip=None): + """``atomicAdd`` (relaxed, GPU scope); returns the old value. The thread stalls only where it + reads the result.""" + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "atom.relaxed.gpu.global.add.u32 $0, [$1], $2;", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ld_acquire_u32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def st_u32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def griddepcontrol_wait(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.wait;", "", loc=loc, ip=ip) + + +@dsl_user_op +def griddepcontrol_launch_dependents(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.launch_dependents;", "", loc=loc, ip=ip) + + +@cute.jit +def lamport_state(flags_addr: Int64): + """``(cur, dirty, bytes_per_buffer, dirty_stages, clear0, clear1, clear2, clear3)`` of the + workspace, read past L1. Valid once the previous call on the workspace has completed (see + the ordering rules of the kernels).""" + f = ld_volatile_v4(flags_addr) + c = ld_volatile_v4(flags_addr + Int64(16)) + return f[0], f[1], f[2], f[3], c[0], c[1], c[2], c[3] + + +@cute.jit +def stage_addr(base: Int64, index: Int32, stage: Int32, num_stages: int, bytes_per_buffer: Int32): + """Address of ``stage`` of Lamport buffer ``index`` (``LamportBufferLayout::getStagePtr``).""" + stage_bytes = Int64(bytes_per_buffer // num_stages) + return base + Int64(index * num_stages + stage) * stage_bytes + + +@cute.jit +def ag_row_addr( + base: Int64, cur: Int32, nbytes: Int32, m: Int32, tok: Int32, world: int, hidden: int +): + """Address of row ``tok`` of the all-gather rows of a one-shot call of ``m`` rows in buffer + ``cur`` (they follow the call's ``m x world`` slots); ``base``: a unicast or the multicast + base of the workspace.""" + return stage_addr(base, cur, Int32(0), 1, nbytes) + Int64(m * world + tok) * (hidden * 2) + + +@cute.jit +def poll_v4(addr: Int64): + """The 16-byte Lamport vector at ``addr``, polled until it has arrived.""" + w = ld_volatile_v4(addr) + while clean4(*w) == 0: + w = ld_volatile_v4(addr) + return w + + +@cute.jit +def clear_stage(addr: Int64, nbytes: Int32, gtid: Int32, nthreads: Int32): + """Re-arm ``nbytes`` (rounded up to 16) at ``addr`` with the sentinel, grid-strided.""" + nvec = (nbytes + 15) // 16 + for i in cutlass.range(gtid, nvec, nthreads, unroll=1): + st_sentinel_v4(addr + Int64(i) * 16) + + +@cute.jit +def push_slot( + peer_ptrs: Int64, + mc_base: Int64, + cur: Int32, + nbytes: Int32, + tok: Int32, + rank: Int32, + elem: Int32, + world: int, + hidden: int, + one_shot: bool, +): + """Address where ``rank`` stores its values ``[elem, elem + 8)`` of token ``tok``: slot + ``(tok // W, rank)`` of the owner's scatter stage (two-shot, through the owner's unicast + mapping), or slot ``(tok, rank)`` of every rank's buffer (one-shot, multicast address).""" + row_bytes = hidden * 2 + addr = Int64(0) + if cutlass.const_expr(one_shot): + addr = stage_addr(mc_base, cur, Int32(0), 1, nbytes) + Int64(tok * world + rank) * row_bytes + else: + owner = ld_nc_u64(peer_ptrs + Int64(tok % world) * 8) + addr = ( + stage_addr(owner, cur, Int32(0), 2, nbytes) + + Int64((tok // world) * world + rank) * row_bytes + ) + return addr + Int64(elem) * 2 + + +@cute.jit +def push_v4(addr: Int64, w0: Int32, w1: Int32, w2: Int32, w3: Int32, one_shot: bool): + """Store one sanitized 16-byte vector into a Lamport slot.""" + s = [sanitize_bf16x2(w) for w in (w0, w1, w2, w3)] + if cutlass.const_expr(one_shot): + multimem_st_v4(addr, *s) + else: + st_v4(addr, *s) + + +def reduce_chunk(world: int, cpp_one_shot: bool) -> int: + """Ranks per partial sum of the C++ kernel a reduction reproduces (see + :class:`K3LamportFinishKernel`): chunks of eight ranks for the one-shot kernel above eight + ranks, up to sixteen ranks in one chunk for the two-shot kernel.""" + return (world if world <= 8 else 8) if cpp_one_shot else min(world, 16) + + +@cute.jit +def poll_reduce_v4(base: Int64, slot_bytes: int, world: int, chunk: int): + """Poll the ``world`` 16-byte slots at ``base + r * slot_bytes`` until no word holds the + sentinel, then sum them in FP32 in rank order, ``chunk`` ranks per partial sum (the order of + :meth:`K3LamportFinishKernel.reduce_slots`); returns eight FP32 values (four BF16 pairs).""" + words = cute.make_rmem_tensor((world * 4,), Int32) + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for r in cutlass.range_constexpr(world): + w = ld_volatile_v4(base + Int64(r * slot_bytes)) + for e in cutlass.range_constexpr(4): + words[r * 4 + e] = w[e] + ready = ready & clean4(*w) + acc = [cutlass.Float32(0.0)] * 8 + for c0 in cutlass.range_constexpr(0, world, chunk): + part = [cutlass.Float32(0.0)] * 8 + for r in cutlass.range_constexpr(c0, min(c0 + chunk, world)): + for q in cutlass.range_constexpr(4): + part[2 * q], part[2 * q + 1] = acc_bf16x2( + words[r * 4 + q], part[2 * q], part[2 * q + 1] + ) + for i in cutlass.range_constexpr(8): + acc[i] = add_f32(acc[i], part[i]) + return acc + + +@cute.jit +def clear_dirty_stages( + uc_local: Int64, + dirty: Int32, + nbytes: Int32, + dstages: Int32, + c0: Int32, + c1: Int32, + c2: Int32, + c3: Int32, + gtid: Int32, + nthreads: Int32, +): + """Re-arm every stage the previous call on the workspace used (its buffer ``dirty``, ``dstages`` + stages, ``c0..c3`` bytes each, from the rotation state), grid-strided over ``nthreads``.""" + clears = (c0, c1, c2, c3) + for s in cutlass.range_constexpr(4): + if s < dstages: + clear_stage( + uc_local + Int64(dirty * dstages + s) * Int64(nbytes // dstages), + clears[s], + gtid, + nthreads, + ) + + +@cute.jit +def arrive_finish( + flags: Int64, + total: Int32, + cur: Int32, + nbytes: Int32, + m: Int32, + world: int, + hidden: int, + one_shot: bool, + extra_rows: int = 0, +): + """Count the calling CTA (one thread per CTA, after every thread of the CTA has read the + rotation state) for a call of ``m`` rows that ``total`` CTAs finish; the last CTA to arrive + writes the next call's state (buffer ``cur`` becomes the dirty one with the stages and bytes + the transport used) and resets the counter, as :class:`K3LamportFinishKernel`'s updater does. + The caller has passed ``griddepcontrol.wait``: the producer, which reads the state after its + own wait, has completed. ``extra_rows``: rows per token past the ``world`` slots (the + all-gather row) that the next call re-arms too (one-shot transport).""" + arrivals = flags + Int64(_FLAG_ARRIVALS * 4) + old = atom_add_acq_rel_u32(arrivals, Int32(1)) + arrive_rotate(flags, old, total, cur, nbytes, m, world, hidden, one_shot, extra_rows) + + +@cute.jit +def arrive_count(flags: Int64): + """The counting half of :func:`arrive_finish` (one thread per CTA, after every thread of the + CTA has read the rotation state and the producer has read it too): a relaxed atomic whose + result only :func:`arrive_rotate` reads, so the thread can go on until then.""" + return atom_add_relaxed_u32(flags + Int64(_FLAG_ARRIVALS * 4), Int32(1)) + + +@cute.jit +def arrive_rotate( + flags: Int64, + old: Int32, + total: Int32, + cur: Int32, + nbytes: Int32, + m: Int32, + world: int, + hidden: int, + one_shot: bool, + extra_rows: int = 0, +): + """The rotation half of :func:`arrive_finish`: the CTA whose count ``old`` was the last of + ``total`` writes the next call's state and resets the counter (read by the next call, which + starts after this grid has completed). ``extra_rows`` as in :func:`arrive_finish`.""" + if old == total - 1: + row_bytes = hidden * 2 + if cutlass.const_expr(one_shot): + st_v4(flags, (cur + 1) % NUM_BUFFERS, cur, nbytes, Int32(1)) + st_v4( + flags + Int64(16), + m * ((world + extra_rows) * row_bytes), + Int32(0), + Int32(0), + Int32(0), + ) + else: + st_v4(flags, (cur + 1) % NUM_BUFFERS, cur, nbytes, Int32(2)) + st_v4( + flags + Int64(16), + ((m + (world - 1)) // world) * (world * row_bytes), + m * row_bytes, + Int32(0), + Int32(0), + ) + st_u32(arrivals_addr(flags), Int32(0)) + + +def arrivals_addr(flags): + """Address of the CTA arrival counter in the state words.""" + return flags + Int64(_FLAG_ARRIVALS * 4) + + +@cute.jit +def finish_two_shot_v4( + uc_local: Int64, + mc_base: Int64, + cur: Int32, + nbytes: Int32, + tok: Int32, + byte_off: Int32, + rank: Int32, + store: Int32, + world: int, + chunk: int, + hidden: int, +): + """16 bytes at ``byte_off`` of row ``tok`` of the two-shot call in buffer ``cur``, finished by + the caller: on the row's owner (``tok % world == rank``) the rank-order sum of its scatter + slots (``chunk`` ranks per partial sum), rounded, ``-0`` replaced by ``+0`` and multicast into + every rank's broadcast stage (where ``store != 0``); elsewhere the broadcast row polled until it + arrived. Both give the bits :class:`K3LamportFinishKernel` broadcasts.""" + w0 = Int32(0) + w1 = Int32(0) + w2 = Int32(0) + w3 = Int32(0) + if tok % world == rank: + scatter = stage_addr(uc_local, cur, Int32(0), 2, nbytes) + acc = poll_reduce_v4( + scatter + Int64((tok // world) * world * hidden) * 2 + Int64(byte_off), + hidden * 2, + world, + chunk, + ) + w0 = sanitize_bf16x2(pack_bf16x2(acc[0], acc[1])) + w1 = sanitize_bf16x2(pack_bf16x2(acc[2], acc[3])) + w2 = sanitize_bf16x2(pack_bf16x2(acc[4], acc[5])) + w3 = sanitize_bf16x2(pack_bf16x2(acc[6], acc[7])) + if store != 0: + bcast_mc = stage_addr(mc_base, cur, Int32(1), 2, nbytes) + multimem_st_v4(bcast_mc + Int64(tok * hidden) * 2 + Int64(byte_off), w0, w1, w2, w3) + else: + src = stage_addr(uc_local, cur, Int32(1), 2, nbytes) + Int64(tok * hidden) * 2 + w = ld_volatile_v4(src + Int64(byte_off)) + while clean4(*w) == 0: + w = ld_volatile_v4(src + Int64(byte_off)) + w0 = w[0] + w1 = w[1] + w2 = w[2] + w3 = w[3] + return w0, w1, w2, w3 + + +@cute.jit +def two_shot_token(v: Int32, m: Int32, rank: Int32, world: int): + """Row handled by the ``v``-th CTA group of a grid that finishes a two-shot call of ``m`` + rows on ``rank``: the rows this rank owns (``rank``, ``rank + world``, ...) first, then the + others in order, so the owners' reductions are dispatched before the CTAs that poll for + them.""" + n_own = (m + (world - 1) - rank) // world + t = rank + v * world + if v >= n_own: + j = v - n_own + q = j // (world - 1) + rem = j - q * (world - 1) + t = q * world + rem + if rem >= rank: + t = t + 1 + return t + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- + + +class K3LamportPushKernel: + """Producer half for a BF16 tensor that is already in memory: copy ``x [m, hidden]`` into + the Lamport slots (the C++ kernels' first phase as a kernel of its own). Grid + ``(m, ceil(hidden / 8 / 128))``, one 16-byte vector per thread.""" + + def __init__(self, world: int, hidden: int, one_shot: bool, use_pdl: bool): + if hidden % 8: + raise ValueError(f"hidden={hidden} must be a multiple of 8") + self.W = world + self.H = hidden + self.one_shot = one_shot + self.use_pdl = use_pdl + self.vecs = hidden // 8 + self.cols = _ceil_div(self.vecs, THREADS) + + @cute.jit + def __call__( + self, + x: Int64, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + m: Int32, + rank: Int32, + stream: cuda.CUstream, + ): + self.kernel(x, peer_ptrs, mc_base, flags, m, rank).launch( + grid=(m, self.cols, 1), block=(THREADS, 1, 1), stream=stream, use_pdl=self.use_pdl + ) + + @cute.kernel + def kernel( + self, x: Int64, peer_ptrs: Int64, mc_base: Int64, flags: Int64, m: Int32, rank: Int32 + ): + tidx, _, _ = cute.arch.thread_idx() + tok, by, _ = cute.arch.block_idx() + v = by * THREADS + tidx + if cutlass.const_expr(self.use_pdl): + griddepcontrol_wait() + cur, _, nbytes, _, _, _, _, _ = lamport_state(flags) + if cutlass.const_expr(self.use_pdl): + griddepcontrol_launch_dependents() + if v < self.vecs: + w = ld_v4(x + Int64(tok * self.H + v * 8) * 2) + dst = push_slot( + peer_ptrs, mc_base, cur, nbytes, tok, rank, v * 8, self.W, self.H, self.one_shot + ) + push_v4(dst, *w, self.one_shot) + + +class K3LamportFinishKernel: + """Rest of an MNNVL all-reduce of ``[m, hidden]`` BF16 rows the producer pushes. + + Grid ``m * ceil(hidden / 8 / 128) + 1`` CTAs of 128 threads. Worker CTA ``b`` owns token + ``b // cols`` and its 16-byte vectors ``(b % cols) * 128 + tid``; the last CTA advances the + rotation. + + Ordering. The workers read the rotation state and start polling as soon as they are + resident, without ``griddepcontrol.wait``: the rows they poll validate themselves (no + sentinel word), so they react the moment the last rank's data lands instead of after this + rank's producer has completed. Reading the state early is safe because the producer + releases this kernel only after its own ``griddepcontrol.wait`` (the low-M GEMM's push + variant and the push kernel both do): by then every earlier grid has completed (each + waits for its predecessor before it completes), the previous all-reduce on the workspace + included, so the state is this call's. Every CTA releases its dependents at its start, so + they become resident once the producer has passed its wait; consumers that read data + before their own wait (attn_res reads the prefix sum and snapshot rows, GEMMs their + weights) read only data older than the producer, which is complete by the same argument. + The updater CTA alone waits for the producer: the producer reads the state after its own + wait, so the state may only change once it has completed; the updater then waits until + every worker has counted itself (having read the state) and writes the next call's state, + like ``LamportFlags::waitAndUpdate``. Its wait also makes this kernel's completion imply + the producer's, so dependents that wait for this kernel see everything before it. + ``wait_first`` makes the workers wait for the producer first (no early polling). + + Numerics. ``cpp_one_shot`` names the C++ kernel whose result the call reproduces (default: + the transport's own): the one-shot kernel sums chunks of eight ranks above eight ranks and + stores the rounded sum as is; the two-shot kernel sums up to sixteen ranks in one chunk and + replaces ``-0`` by ``+0`` (its result passes through a Lamport buffer). Either transport + reproduces either result. The two-shot transport replaces ``-0`` in a one-shot result too, + which has none: the inputs carry no ``-0``, and every partial sum of BF16 values in FP32 is + a multiple of ``2**-133``, the smallest BF16 denormal (FP32 rounds only where its spacing + exceeds that), so a nonzero sum rounds to a nonzero BF16 value and an exact zero is ``+0`` + under round-to-nearest. + """ + + def __init__( + self, + world: int, + hidden: int, + one_shot: bool, + use_pdl: bool, + wait_first: bool = False, + cpp_one_shot: bool | None = None, + ag: bool = False, + ): + if hidden % 8: + raise ValueError(f"hidden={hidden} must be a multiple of 8") + if cpp_one_shot is None: + cpp_one_shot = one_shot + if ag and not one_shot: + raise ValueError("the all-gather row goes with the one-shot transport") + # The call carries an all-gather row per token: the output is bf16(y + bf16(sum)). + self.ag = ag + self.W = world + self.H = hidden + self.one_shot = one_shot + self.use_pdl = use_pdl + self.wait_first = wait_first + self.vecs = hidden // 8 + self.cols = _ceil_div(self.vecs, THREADS) + # Ranks per partial sum of the reproduced C++ kernel, and whether the one-shot transport + # sanitizes its output like the two-shot kernel (the two-shot transport always does). + self.chunk = (world if world <= 8 else 8) if cpp_one_shot else min(world, 16) + self.sanitize_out = not cpp_one_shot + + @cute.jit + def __call__( + self, + uc_local: Int64, + mc_base: Int64, + flags: Int64, + out: Int64, + m: Int32, + rank: Int32, + stream: cuda.CUstream, + ): + self.kernel(uc_local, mc_base, flags, out, m, rank).launch( + grid=(m * self.cols + 1, 1, 1), + block=(THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.jit + def reduce_slots(self, base: Int64, slot_bytes: int): + """Poll the ``W`` 16-byte slots at ``base + r * slot_bytes`` until every word has arrived, + then sum them in the C++ order; returns eight FP32 values (four BF16 pairs).""" + W = self.W + words = cute.make_rmem_tensor((W * 4,), Int32) + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for r in cutlass.range_constexpr(W): + w = ld_volatile_v4(base + Int64(r * slot_bytes)) + for e in cutlass.range_constexpr(4): + words[r * 4 + e] = w[e] + ready = ready & clean4(*w) + acc = [cutlass.Float32(0.0)] * 8 + for c0 in cutlass.range_constexpr(0, W, self.chunk): + part = [cutlass.Float32(0.0)] * 8 + for r in cutlass.range_constexpr(c0, c0 + self.chunk): + for q in cutlass.range_constexpr(4): + part[2 * q], part[2 * q + 1] = acc_bf16x2( + words[r * 4 + q], part[2 * q], part[2 * q + 1] + ) + for i in cutlass.range_constexpr(8): + acc[i] = add_f32(acc[i], part[i]) + return acc + + @cute.kernel + def kernel( + self, uc_local: Int64, mc_base: Int64, flags: Int64, out: Int64, m: Int32, rank: Int32 + ): + W, H = self.W, self.H + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + workers = m * self.cols + arrivals = flags + Int64(_FLAG_ARRIVALS * 4) + if bidx == workers: + # Updater: the next call's state, once the producer is done with this call's. + if cutlass.const_expr(self.use_pdl): + griddepcontrol_launch_dependents() + griddepcontrol_wait() + if tidx == 0: + cur, _, nbytes, _, _, _, _, _ = lamport_state(flags) + arrived = ld_acquire_u32(arrivals) + while arrived < workers: + arrived = ld_acquire_u32(arrivals) + row_bytes = H * 2 + clear0 = m * ((W + (1 if self.ag else 0)) * row_bytes) + clear1 = Int32(0) + if cutlass.const_expr(not self.one_shot): + clear0 = ((m + (W - 1)) // W) * (W * row_bytes) + clear1 = m * row_bytes + st_v4(flags, (cur + 1) % NUM_BUFFERS, cur, nbytes, Int32(1 if self.one_shot else 2)) + st_v4(flags + Int64(16), clear0, clear1, Int32(0), Int32(0)) + st_u32(arrivals, Int32(0)) + else: + tok = bidx // self.cols + by = bidx - tok * self.cols + v = by * THREADS + tidx + if cutlass.const_expr(self.use_pdl and self.wait_first): + griddepcontrol_wait() + cur, dirty, nbytes, dstages, c0, c1, c2, c3 = lamport_state(flags) + if cutlass.const_expr(self.use_pdl): + griddepcontrol_launch_dependents() + # Count this CTA once every thread has read the state. + cute.arch.barrier() + if tidx == 0: + red_release_add(arrivals, Int32(1)) + row_bytes = H * 2 + out_v = out + Int64(tok * H + v * 8) * 2 + if cutlass.const_expr(self.one_shot): + buf = stage_addr(uc_local, cur, Int32(0), 1, nbytes) + if v < self.vecs: + acc = self.reduce_slots(buf + Int64(tok * W * H + v * 8) * 2, row_bytes) + if cutlass.const_expr(self.ag): + # bf16(y + bf16(sum)): the sum rounds as the C++ kernel stores it (no -0 + # arises, see Numerics), then the all-gathered value is added. + y = poll_v4( + ag_row_addr(uc_local, cur, nbytes, m, tok, W, H) + Int64(v * 16) + ) + o = [ag_add_bf16x2(acc[2 * q], acc[2 * q + 1], y[q]) for q in range(4)] + else: + o = [pack_bf16x2(acc[2 * q], acc[2 * q + 1]) for q in range(4)] + if cutlass.const_expr(self.sanitize_out): + o = [sanitize_bf16x2(w) for w in o] + st_v4(out_v, *o) + else: + if v < self.vecs: + if tok % W == rank: + scatter = stage_addr(uc_local, cur, Int32(0), 2, nbytes) + acc = self.reduce_slots( + scatter + Int64((tok // W) * W * H + v * 8) * 2, row_bytes + ) + o = [ + sanitize_bf16x2(pack_bf16x2(acc[2 * q], acc[2 * q + 1])) + for q in range(4) + ] + bcast_mc = stage_addr(mc_base, cur, Int32(1), 2, nbytes) + multimem_st_v4(bcast_mc + Int64(tok * H + v * 8) * 2, *o) + # Re-arm the previous call's buffer (every stage it used) while the rows travel. + gtid = bidx * THREADS + tidx + nthreads = workers * THREADS + clears = (c0, c1, c2, c3) + for s in cutlass.range_constexpr(4): + if s < dstages: + clear_stage( + uc_local + Int64(dirty * dstages + s) * Int64(nbytes // dstages), + clears[s], + gtid, + nthreads, + ) + if cutlass.const_expr(not self.one_shot): + if v < self.vecs: + src = ( + stage_addr(uc_local, cur, Int32(1), 2, nbytes) + Int64(tok * H + v * 8) * 2 + ) + w = ld_volatile_v4(src) + while clean4(*w) == 0: + w = ld_volatile_v4(src) + st_v4(out_v, *w) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class LamportWorkspace: + """Device addresses of one MNNVL Lamport workspace (``get_or_scale_allreduce_mnnvl_workspace``). + + ``peer_ptrs``: int64 CUDA tensor with every rank's unicast mapping of the workspace (this + rank's own at index ``rank``); ``flags``: the uint32 rotation state shared with the C++ + kernels; ``bytes_per_buffer``: size of one of the three Lamport buffers. + """ + + world: int + rank: int + uc_local: int + mc_base: int + peer_ptrs: torch.Tensor + flags: torch.Tensor + bytes_per_buffer: int + + @staticmethod + def from_mnnvl(workspace: dict, world: int, rank: int) -> "LamportWorkspace": + handle = workspace["handle"] + peers = [handle.get_uc_buffer(r, (1,), torch.uint8, 0).data_ptr() for r in range(world)] + local = workspace["uc_buffer"].data_ptr() + if peers[rank] != local: + raise RuntimeError("MNNVL workspace: rank's unicast mapping is not the local buffer") + return LamportWorkspace( + world=world, + rank=rank, + uc_local=local, + mc_base=handle.get_mc_buffer((1,), torch.uint8, 0).data_ptr(), + peer_ptrs=torch.tensor(peers, dtype=torch.int64, device=workspace["uc_buffer"].device), + flags=workspace["buffer_flags"], + bytes_per_buffer=int(workspace["buffer_size_bytes"]), + ) + + +def one_shot_for(num_tokens: int, hidden: int, world: int, threshold_bytes: int) -> bool: + """The C++ rule: one-shot when ``tokens * hidden * world * 2 <= threshold``.""" + return num_tokens * hidden * world * 2 <= threshold_bytes + + +def required_bytes( + num_tokens: int, hidden: int, world: int, one_shot: bool, ag: bool = False +) -> int: + """Bytes of one Lamport buffer the call needs (as ``MNNVLAllReduce.get_required_workspace_size``; + ``ag``: with the all-gather rows, one-shot transport).""" + if one_shot: + return num_tokens * hidden * (world + (1 if ag else 0)) * 2 + return 2 * _ceil_div(num_tokens, world) * world * hidden * 2 + + +def _stream() -> cuda.CUstream: + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + +@functools.cache +def _compiled_finish( + world: int, + hidden: int, + one_shot: bool, + use_pdl: bool, + wait_first: bool = False, + cpp_one_shot: bool | None = None, + ag: bool = False, +): + kernel = K3LamportFinishKernel(world, hidden, one_shot, use_pdl, wait_first, cpp_one_shot, ag) + a = Int64(1 << 20) + return cute.compile(kernel, a, a, a, a, Int32(1), Int32(0), _stream()) + + +@functools.cache +def _compiled_push(world: int, hidden: int, one_shot: bool, use_pdl: bool): + kernel = K3LamportPushKernel(world, hidden, one_shot, use_pdl) + a = Int64(1 << 20) + return cute.compile(kernel, a, a, a, a, Int32(1), Int32(0), _stream()) + + +def _check( + ws: LamportWorkspace, num_tokens: int, hidden: int, one_shot: bool, ag: bool = False +) -> None: + if hidden % 8: + raise ValueError(f"hidden={hidden} must be a multiple of 8") + need = required_bytes(num_tokens, hidden, ws.world, one_shot, ag) + if need > ws.bytes_per_buffer: + raise ValueError( + f"Lamport buffer of {ws.bytes_per_buffer} B too small for {num_tokens} x {hidden} " + f"({need} B)" + ) + + +def run_lamport_push( + ws: LamportWorkspace, x: torch.Tensor, one_shot: bool, use_pdl: bool = True +) -> None: + """Push the contiguous BF16 ``x [m, hidden]`` into the Lamport slots of the next call.""" + m, hidden = x.shape + _check(ws, m, hidden, one_shot) + _compiled_push(ws.world, hidden, one_shot, use_pdl)( + Int64(x.data_ptr()), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int32(m), + Int32(ws.rank), + _stream(), + ) + + +def run_lamport_finish( + ws: LamportWorkspace, + out: torch.Tensor, + one_shot: bool, + use_pdl: bool = True, + wait_first: bool = False, + cpp_one_shot: bool | None = None, + ag: bool = False, +) -> torch.Tensor: + """Complete the all-reduce whose input every rank pushed for this call into ``out`` + (contiguous BF16 ``[m, hidden]``); returns ``out``. ``one_shot`` is the transport the + producer used; ``cpp_one_shot`` the C++ kernel whose result to reproduce (default: the + same; see :class:`K3LamportFinishKernel`). ``ag``: the call carries all-gather rows (see the + module docstring); ``out`` is then ``bf16(y + bf16(sum))``.""" + m, hidden = out.shape + _check(ws, m, hidden, one_shot, ag) + _compiled_finish(ws.world, hidden, one_shot, use_pdl, wait_first, cpp_one_shot, ag)( + Int64(ws.uc_local), + Int64(ws.mc_base), + Int64(ws.flags.data_ptr()), + Int64(out.data_ptr()), + Int32(m), + Int32(ws.rank), + _stream(), + ) + return out + + +__all__ = [ + "K3LamportFinishKernel", + "K3LamportPushKernel", + "LamportWorkspace", + "ag_add_bf16x2", + "ag_row_addr", + "arrive_finish", + "clear_dirty_stages", + "finish_two_shot_v4", + "one_shot_for", + "poll_reduce_v4", + "poll_v4", + "push_slot", + "push_v4", + "reduce_chunk", + "required_bytes", + "two_shot_token", + "run_lamport_finish", + "run_lamport_push", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_latent_ar_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_latent_ar_norm.py new file mode 100644 index 000000000000..08da772a47a6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_latent_ar_norm.py @@ -0,0 +1,1298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 routed latent all-reduce + RMSNorm (MNNVL one-shot, SM100/SM103/SM107) with an early +PDL trigger, bit-identical to ``oneshotAllreduceFusionKernel``. + +The C++ kernel (``mnnvlAllreduceKernels.cu``) serves the one-shot sizes of the latent all-reduce +(``M`` rows of 3584 BF16; ``M <= 38`` launches clusters of 4 CTAs x 112 threads per row). It +triggers its programmatic dependents only after the cross-rank reduction, so latent_up cannot +start loading its weight while the all-reduce waits for the other ranks. +:class:`K3LatentArNormKernel` does the same work with the same arithmetic and triggers right after +its own ``griddepcontrol.wait``: + +* Thread ``t`` of cluster CTA ``c`` owns the eight values ``[8 (112 c + t), + 8)`` of its row. It + replaces BF16 ``-0`` by ``+0`` (the Lamport sentinel), stores the 16 bytes into slot + ``(row, rank)`` of every rank's current Lamport buffer through the multicast address, clears + this rank's dirty buffer (the bytes the previous launch on the workspace left to clear), and + polls its eight values of all ``WS`` slots in the local buffer. Ranks are summed in FP32 in + chunks of eight like the C++ kernel (each chunk from 0, chunks added in rank order) and rounded + to BF16. +* RMSNorm: BF16 squares (``mul.rn.bf16x2``) summed in element order; butterfly warp sums (the + half warp first adds 0 for its missing upper half); the block sum ``(w0 + w2) + (w1 + w3)`` of + the reference's warp-0 butterfly over four warp sums; the four CTA sums of the row added in CTA + order from 0 through distributed shared memory; ``rsqrt.approx(div.rn(sum, 3584) + eps)``; then + ``bf16((x * r) * gamma)``. +* Lamport state: the same flag words as ``LamportFlags`` (current / dirty buffer, buffer bytes, + dirty stages, bytes to clear per stage, arrival counter), so the C++ one- and two-shot kernels + can use the same workspace before or after it. Each warp counts one arrival once its lanes + have read the flags; at the end, thread 0 of CTA (0, 0) waits for all warps of the grid, rotates + the buffers and resets the counter. + +Optionally the kernel also stores the MXFP8 form of its output for latent_up +(``k3_fp8w_lowm_mxfp8``): per 1x32 group (four consecutive threads) the low-M FP8 GEMM's own +amax / UE8M0 exponent / E4M3 conversion, so that GEMM on it is bitwise the low-M GEMM on the BF16 +output. + +PDL: ``griddepcontrol.wait`` first, and the dependents are triggered right after it (they may +only start reading what they do not wait for, e.g. latent_up its weight), not after the reduction +like the C++ kernel. + +:class:`K3LatentArRowKernel` computes the same result, bit for bit, +with one 512-thread CTA per row: four groups of 128 threads play the four cluster CTAs (same warp +boundaries and reduction trees; every thread forms the group sums after one CTA barrier instead of +a cluster barrier); the rank's own slot is summed from registers instead of being polled back +from the multicast loopback, after the first polling round only the slots that have not arrived +are loaded again, and CTA 0 rotates the buffers while it waits for the other ranks instead of at +the end. +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from .k3_fp8w_lowm_gemm import _bf16x2_inv_scale, _quant_bf16x8_to_e4m3, _ue8m0_of_amax +from .k3_route_tail import _griddepcontrol_launch_dependents, _griddepcontrol_wait + +#: Row width (the latent) and values per thread (one 16-byte access). +HIDDEN = 3584 +ELTS = 8 +#: Cluster of CTAs per row and threads per CTA, as the C++ launcher picks them for 3584 BF16 values +#: and ``M <= 38`` rows (adjustGridConfig): 4 x 112. +CLUSTER = 4 +THREADS = HIDDEN // ELTS // CLUSTER +#: Warps per CTA (the last one is a half warp). +WARPS = (THREADS + 31) // 32 +#: Largest row count with that launch shape. +MAX_M = 38 +#: Ranks summed per chunk (``accumulateLamportRanksChunked``). +RANK_CHUNK = 8 +#: World sizes the C++ launcher dispatches. +WORLD_SIZES = (2, 4, 8, 16, 32, 64) +#: Lamport buffers per workspace. +NUM_LAMPORT_BUFFERS = 3 + +assert THREADS * CLUSTER * ELTS == HIDDEN and THREADS % 16 == 0 + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _lv(value, *, loc=None, ip=None): + return Int64(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return Float32(value).ir_value(loc=loc, ip=ip) + + +def _unpack(res, ty, cls, n, *, loc=None, ip=None): + return tuple(cls(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(n)) + + +@dsl_user_op +def _ld_v4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _ld_vol_v4(addr, *, loc=None, ip=None): + """16 bytes of the Lamport buffer, read past L1 (peers write it over NVLink).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.volatile.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _st_v4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_sentinel_v4(addr, *, loc=None, ip=None): + """Four FP32 ``-0`` words (``getPackedLamportInit``).""" + _asm( + None, + [_lv(addr, loc=loc, ip=ip)], + "{\n.reg .b32 t;\nmov.b32 t, 0x80000000;\nst.global.v4.b32 [$0], {t, t, t, t};\n}", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sanitize(word, *, loc=None, ip=None): + """A BF16 pair with each ``-0`` (0x8000) replaced by ``+0``.""" + return Int32( + _asm( + T.i32(), + [_iv(word, loc=loc, ip=ip)], + "{\n.reg .pred p;\n.reg .b32 lo, hi;\n" + "and.b32 lo, $1, 0xffff;\nsetp.eq.u32 p, lo, 0x8000;\nselp.b32 lo, 0, lo, p;\n" + "and.b32 hi, $1, 0xffff0000;\nsetp.eq.u32 p, hi, 0x80000000;\nselp.b32 hi, 0, hi, p;\n" + "or.b32 $0, lo, hi;\n}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _clean4(w0, w1, w2, w3, *, loc=None, ip=None): + """1 if none of the four words is the sentinel (``isLamportDirty`` is false).""" + return Int32( + _asm( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)], + "{\n.reg .pred p;\n" + "setp.ne.b32 p, $1, 0x80000000;\nsetp.ne.and.b32 p, $2, 0x80000000, p;\n" + "setp.ne.and.b32 p, $3, 0x80000000, p;\nsetp.ne.and.b32 p, $4, 0x80000000, p;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bf16x2_f32(word, *, loc=None, ip=None): + """(low, high) BF16 of a pair as FP32.""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_iv(word, loc=loc, ip=ip)], + "{\n.reg .b32 lo, hi;\nshl.b32 lo, $2, 16;\nand.b32 hi, $2, 0xffff0000;\n" + "mov.b32 $0, lo;\nmov.b32 $1, hi;\n}", + "=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _pack_bf16x2(lo, hi, *, loc=None, ip=None): + """Round two FP32 values to a BF16 pair (``lo`` in the low half).""" + return Int32( + _asm( + T.i32(), + [_fv(lo, loc=loc, ip=ip), _fv(hi, loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sq_bf16x2(word, *, loc=None, ip=None): + """Elementwise BF16 square of a pair (``mul.rn.bf16x2``).""" + return Int32( + _asm( + T.i32(), + [_iv(word, loc=loc, ip=ip)], + "mul.rn.bf16x2 $0, $1, $1;", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fadd(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fmul(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _rms_scale(total, eps, *, loc=None, ip=None): + """``rsqrtf(total / 3584 + eps)`` as the C++ kernel computes it (IEEE division, then add).""" + return Float32( + _asm( + T.f32(), + [_fv(total, loc=loc, ip=ip), _fv(eps, loc=loc, ip=ip)], + "{\n.reg .f32 q;\ndiv.rn.f32 q, $1, 0f45600000;\nadd.rn.f32 q, q, $2;\n" + "rsqrt.approx.f32 $0, q;\n}", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_bfly(value, offset: int, mask: int, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(value, loc=loc, ip=ip)], + f"shfl.sync.bfly.b32 $0, $1, {offset}, 0x1f, {mask};", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _amax_bf16x8(w0, w1, w2, w3, *, loc=None, ip=None): + """|x| max of eight BF16 values (four packed words) as a packed BF16 pair (garbage sign bit, + masked by ``_ue8m0_of_amax``).""" + return Int32( + _asm( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)], + "{\n.reg .b32 m0, m1;\n" + "max.xorsign.abs.bf16x2 m0, $1, $2;\n" + "max.xorsign.abs.bf16x2 m1, $3, $4;\n" + "max.xorsign.abs.bf16x2 $0, m0, m1;\n}", + "=r,r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _amax_pair(a, b, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_bfly_b32(value, offset: int, mask: int, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(value, loc=loc, ip=ip)], + f"shfl.sync.bfly.b32 $0, $1, {offset}, 0x1f, {mask};", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_v2(addr, v0, v1, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(v0, loc=loc, ip=ip), _iv(v1, loc=loc, ip=ip)], + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_u8(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u8 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.f32 $0, [$1];", "=f,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _st_cluster_f32(local_addr, peer, value, *, loc=None, ip=None): + """Store ``value`` at ``local_addr`` of cluster CTA ``peer``'s shared memory.""" + _asm( + None, + [_iv(local_addr, loc=loc, ip=ip), _iv(peer, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "{\n.reg .b32 ra;\nmapa.shared::cluster.u32 ra, $0, $1;\n" + "st.shared::cluster.f32 [ra], $2;\n}", + "r,r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cluster_sync_release_acquire(*, loc=None, ip=None): + _asm( + None, + [], + "barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;", + "", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _warp_sync(half_mask: int, is_half, *, loc=None, ip=None): + """``bar.warp.sync`` over the warp's existing lanes (the last warp of the CTA may be a half + warp: ``half_mask``).""" + _asm( + None, + [_iv(is_half, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.s32 p, $0, 0;\n" + f"@p bar.warp.sync {half_mask};\n@!p bar.warp.sync 0xffffffff;\n}}", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bar_sync(*, loc=None, ip=None): + _asm(None, [], "bar.sync 0;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _red_add_release_gpu(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_acquire_gpu_u32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_u32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +class K3LatentArNormKernel: + """One-shot MNNVL all-reduce + RMSNorm of ``M`` rows of 3584 BF16 (see the module docstring).""" + + def __init__( + self, + world_size: int, + use_pdl: bool, + norm: bool = True, + mxfp8: bool = False, + ): + if world_size not in WORLD_SIZES: + raise ValueError(f"unsupported world size {world_size}") + if mxfp8 and not norm: + raise ValueError("the MXFP8 output is the normalized output's") + self.world_size = world_size + self.use_pdl = use_pdl + # Without ``norm`` the output is the BF16 sum (``kAllReduce``), e.g. to check the rank order. + self.norm = norm + # Also store the MXFP8 form of the output (E4M3 codes, one UE8M0 exponent per 1x32 group). + self.mxfp8 = mxfp8 + self.rank_chunk = min(RANK_CHUNK, world_size) + + @cute.jit + def __call__( + self, + x_addr: Int64, + gamma_addr: Int64, + out_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + uc_addr: Int64, + mc_addr: Int64, + flags_addr: Int64, + rank: Int32, + m: Int32, + eps: Float32, + stream: cuda.CUstream, + ): + @cute.struct + class SharedStorage: + # Warp sums, the block sum, and the row's CTA sums (written by the cluster peers). + warp_sum: cute.struct.MemRange[cutlass.Float32, WARPS] + block_sum: cute.struct.MemRange[cutlass.Float32, 1] + cta_sum: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, CLUSTER], 16] + + self.shared_storage = SharedStorage + self.kernel( + x_addr, + gamma_addr, + out_addr, + xq_addr, + xs_addr, + uc_addr, + mc_addr, + flags_addr, + rank, + m, + eps, + ).launch( + grid=(m, CLUSTER, 1), + block=(THREADS, 1, 1), + cluster=(1, CLUSTER, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + x_addr: Int64, + gamma_addr: Int64, + out_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + uc_addr: Int64, + mc_addr: Int64, + flags_addr: Int64, + rank: Int32, + m: Int32, + eps: Float32, + ): + ws = self.world_size + tidx, _, _ = cute.arch.thread_idx() + row, cy, _ = cute.arch.block_idx() + crank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + storage = utils.SmemAllocator().allocate(self.shared_storage) + warp_sum = storage.warp_sum.data_ptr().toint() + block_sum = storage.block_sum.data_ptr().toint() + cta_sum = storage.cta_sum.data_ptr().toint() + + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + _griddepcontrol_launch_dependents() + + # LamportFlags: [current, dirty, bytes per buffer, dirty stages], bytes to clear x4. + f0, f1, f2, f3 = _ld_v4(flags_addr) + c0, c1, c2, c3 = _ld_v4(flags_addr + Int64(16)) + # One arrival per warp once its lanes hold the flags: CTA (0, 0) rotates them only after + # every warp of the grid has read them. + counter = flags_addr + Int64(32) + _warp_sync(0xFFFF if THREADS % 32 else 0xFFFFFFFF, warp == WARPS - 1) + if lane == 0: + _red_add_release_gpu(counter, Int32(1)) + cur = Int64(f0) + dirty = Int64(f1) + buf_bytes = Int64(f2) & Int64(0xFFFFFFFF) + dirty_stages = f3 + col = (crank * THREADS + tidx) * ELTS + cur_off = cur * buf_bytes + + # Broadcast this rank's eight values into slot (row, rank) of every rank's buffer. + xw = _ld_v4(x_addr + Int64((row * HIDDEN + col) * 2)) + if cutlass.const_expr(self.norm): + gw = _ld_v4(gamma_addr + Int64(col * 2)) + sw = [_sanitize(w) for w in xw] + slot = Int64(((row * ws + rank) * HIDDEN + col) * 2) + _st_v4(mc_addr + cur_off + slot, sw[0], sw[1], sw[2], sw[3]) + + # Clear this rank's dirty buffer (all stages the previous launch used). + n_threads = m * CLUSTER * THREADS + gtid = (row * CLUSTER + cy) * THREADS + tidx + clear = (c0, c1, c2, c3) + for s in cutlass.range_constexpr(4): + if s < dirty_stages: + stage_bytes = buf_bytes // Int64(dirty_stages) + base = uc_addr + (dirty * Int64(dirty_stages) + Int64(s)) * stage_bytes + n_vec = (clear[s] + 15) >> 4 + for p in cutlass.range(gtid, n_vec, n_threads, unroll=1): + _st_sentinel_v4(base + Int64(p) * Int64(16)) + + # Sum all ranks' values in chunks of RANK_CHUNK, polling until no word is the sentinel. + acc = [Float32(0.0)] * ELTS + row_base = uc_addr + cur_off + Int64((row * ws * HIDDEN + col) * 2) + for rb in cutlass.range_constexpr(0, ws, self.rank_chunk): + chunk = [Float32(0.0)] * ELTS + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for i in cutlass.range_constexpr(ELTS): + chunk[i] = Float32(0.0) + for rr in cutlass.range_constexpr(self.rank_chunk): + w = _ld_vol_v4(row_base + Int64((rb + rr) * HIDDEN * 2)) + ready = ready & _clean4(w[0], w[1], w[2], w[3]) + for j in cutlass.range_constexpr(4): + lo, hi = _bf16x2_f32(w[j]) + chunk[2 * j] = _fadd(chunk[2 * j], lo) + chunk[2 * j + 1] = _fadd(chunk[2 * j + 1], hi) + for i in cutlass.range_constexpr(ELTS): + acc[i] = _fadd(acc[i], chunk[i]) + red = [_pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)] + + if cutlass.const_expr(self.norm): + # RMSNorm: this thread's BF16 squares in element order. + thread_sum = Float32(0.0) + for j in cutlass.range_constexpr(4): + lo, hi = _bf16x2_f32(_sq_bf16x2(red[j])) + thread_sum = _fadd(thread_sum, lo) + thread_sum = _fadd(thread_sum, hi) + # Warp butterflies; the half warp first adds 0 for its missing upper half. + v = thread_sum + if warp == WARPS - 1: + v = _fadd(v, Float32(0.0)) + for off in (8, 4, 2, 1): + v = _fadd(v, _shfl_bfly(v, off, 0xFFFF)) + else: + for off in (16, 8, 4, 2, 1): + v = _fadd(v, _shfl_bfly(v, off, 0xFFFFFFFF)) + if lane == 0: + _sts_f32(warp_sum + warp * 4, v) + _bar_sync() + if warp == 0: + b = Float32(0.0) + if lane < WARPS: + b = _lds_f32(warp_sum + lane * 4) + for off in (16, 8, 4, 2, 1): + b = _fadd(b, _shfl_bfly(b, off, 0xFFFFFFFF)) + if lane == 0: + _sts_f32(block_sum, b) + _bar_sync() + bsum = _lds_f32(block_sum) + # Row sum: CTA c's sum into every peer's slot c, then the four slots in CTA order. + if tidx < CLUSTER: + _st_cluster_f32(cta_sum + crank * 4, tidx, bsum) + _cluster_sync_release_acquire() + total = Float32(0.0) + for c in cutlass.range_constexpr(CLUSTER): + total = _fadd(total, _lds_f32(cta_sum + c * 4)) + scale = _rms_scale(total, eps) + + ow = [] + for j in cutlass.range_constexpr(4): + x_lo, x_hi = _bf16x2_f32(red[j]) + g_lo, g_hi = _bf16x2_f32(gw[j]) + ow.append( + _pack_bf16x2(_fmul(_fmul(x_lo, scale), g_lo), _fmul(_fmul(x_hi, scale), g_hi)) + ) + _st_v4(out_addr + Int64((row * HIDDEN + col) * 2), ow[0], ow[1], ow[2], ow[3]) + if cutlass.const_expr(self.mxfp8): + # MXFP8 form of the stored row, as the low-M FP8 GEMM quantizes its BF16 + # activation: four consecutive threads hold one 1x32 group (the half warp's + # quads use its 16-lane mask). + amax = _amax_bf16x8(ow[0], ow[1], ow[2], ow[3]) + if warp == WARPS - 1: + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 1, 0xFFFF)) + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 2, 0xFFFF)) + else: + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 1, 0xFFFFFFFF)) + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 2, 0xFFFFFFFF)) + e = _ue8m0_of_amax(amax) + q0, q1 = _quant_bf16x8_to_e4m3(ow[0], ow[1], ow[2], ow[3], _bf16x2_inv_scale(e)) + _st_v2(xq_addr + Int64(row * HIDDEN + col), q0, q1) + if (tidx & 3) == 0: + _st_u8(xs_addr + Int64(row * (HIDDEN // 32) + (col >> 5)), e) + else: + _st_v4(out_addr + Int64((row * HIDDEN + col) * 2), red[0], red[1], red[2], red[3]) + + # CTA (0, 0): rotate the buffers once every warp of the grid has read the flags. + if crank == 0: + if tidx == 0: + if row == 0: + arrived = _ld_acquire_gpu_u32(counter) + while arrived < m * (CLUSTER * WARPS): + arrived = _ld_acquire_gpu_u32(counter) + nxt = f0 + 1 + if nxt == NUM_LAMPORT_BUFFERS: + nxt = Int32(0) + _st_v4(flags_addr, nxt, f0, f2, Int32(1)) + _st_v4( + flags_addr + Int64(16), m * (HIDDEN * ws * 2), Int32(0), Int32(0), Int32(0) + ) + _st_u32(counter, Int32(0)) + + +@functools.cache +def _compiled(world_size: int, use_pdl: bool, norm: bool, mxfp8: bool): + kernel = K3LatentArNormKernel(world_size, use_pdl, norm, mxfp8) + return cute.compile( + kernel, + Int64(0), + Int64(0), + Int64(0), + Int64(0), + Int64(0), + Int64(0), + Int64(0), + Int64(0), + Int32(0), + Int32(1), + Float32(1e-6), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +@functools.cache +def _is_supported_device(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3), (10, 7)) + + +def can_run(x: torch.Tensor, gamma: torch.Tensor, world_size: int) -> bool: + """Whether :func:`run_k3_latent_ar_norm` serves this all-reduce (non-raising check).""" + return ( + x.is_cuda + and _is_supported_device(x.device.index) + and world_size in WORLD_SIZES + and x.dtype == torch.bfloat16 + and x.dim() == 2 + and x.shape[1] == HIDDEN + and 1 <= x.shape[0] <= MAX_M + and x.is_contiguous() + and x.data_ptr() % 16 == 0 + and gamma.dtype == torch.bfloat16 + and tuple(gamma.shape) == (HIDDEN,) + and gamma.is_contiguous() + and gamma.data_ptr() % 16 == 0 + ) + + +def run_k3_latent_ar_norm( + x: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + *, + uc_addr: int, + mc_addr: int, + flags: torch.Tensor, + rank: int, + world_size: int, + use_pdl: bool = True, + norm: bool = True, + mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> None: + """Launch on the current stream: ``out`` (``[M, 3584]`` BF16) = RMSNorm(``gamma``, ``eps``) of + the sum of every rank's ``x`` over the MNNVL Lamport workspace (local unicast base + ``uc_addr``, multicast base ``mc_addr``, state ``flags``), bit-identical to the C++ one-shot + kernel (``norm=False``: the BF16 sum alone). Every rank of the group launches it with its + ``rank``. ``mxfp8_out = (xq, xs)`` also stores the MXFP8 form of ``out`` (E4M3 ``xq [M, 3584]``, + UE8M0 ``xs [M, 112]``), bitwise the low-M FP8 GEMM's in-kernel quantization of ``out``.""" + mxfp8 = mxfp8_out is not None + _compiled(world_size, use_pdl, norm, mxfp8)( + Int64(x.data_ptr()), + Int64(gamma.data_ptr()), + Int64(out.data_ptr()), + Int64(mxfp8_out[0].data_ptr() if mxfp8 else 0), + Int64(mxfp8_out[1].data_ptr() if mxfp8 else 0), + Int64(uc_addr), + Int64(mc_addr), + Int64(flags.data_ptr()), + Int32(rank), + Int32(x.shape[0]), + Float32(eps), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + + +# --------------------------------------------------------------------------- +# One CTA per row +# --------------------------------------------------------------------------- + +#: Row kernel: four groups of 128 threads per CTA; threads ``t < 112`` of group ``g`` play thread +#: ``t`` of cluster CTA ``g`` of :class:`K3LatentArNormKernel`, the others shadow lane ``t - 16`` of +#: their (half) warp and never store. +ROW_GROUPS = CLUSTER +ROW_GROUP_THREADS = 128 +ROW_THREADS = ROW_GROUPS * ROW_GROUP_THREADS +ROW_WARPS = ROW_THREADS // 32 +#: Largest row count of the row kernel (one CTA per row). c240: 128 (was 64), for the helix8 b16 +#: mtp5 verify step (96 rows); the kernel itself has no row limit. +ROW_MAX_M = 128 +#: World sizes of the row kernel (one bit per rank of the arrival mask). +ROW_WORLD_SIZES = (2, 4, 8, 16) + +assert ROW_GROUP_THREADS - THREADS == 16 + + +@dsl_user_op +def _ld_vol_v4_pred(addr, pred, o0, o1, o2, o3, *, loc=None, ip=None): + """16 bytes of the Lamport buffer (read past L1) if ``pred`` != 0, else ``(o0, o1, o2, o3)`` + (a predicated load: no branch).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(pred, loc=loc, ip=ip)] + + [_iv(v, loc=loc, ip=ip) for v in (o0, o1, o2, o3)] + + [_lv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.b32 p, $4, 0;\nmov.b32 $0, $5;\nmov.b32 $1, $6;\n" + "mov.b32 $2, $7;\nmov.b32 $3, $8;\n" + "@p ld.volatile.global.v4.b32 {$0, $1, $2, $3}, [$9];\n}", + "=&r,=&r,=&r,=&r,r,r,r,r,r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +class K3LatentArRowKernel: + """:class:`K3LatentArNormKernel` with one 512-thread CTA per row (same result, bit for bit). + + * Group ``g`` (threads ``128 g ...``) does what cluster CTA ``g`` does, with the same warp + boundaries (its fourth warp is the half warp). After one CTA barrier every thread forms the + four group sums from the sixteen warp sums as the cluster kernel's warp butterfly does, + ``(w0 + w2) + (w1 + w3)``, and adds them in group order (no cluster barrier, no second CTA + barrier). + * Every round polls all ranks' slots of the thread's eight values at once. The rank's own slot + is not read back: the sanitized values the thread stored through the multicast address are + summed from its registers, so the last rank to arrive does not wait for its own data to + come back from the switch. After the first round only the slots that have not arrived are + loaded again (predicated loads). + * The previous call's buffer is cleared after the push, like the C++ kernels. + * CTA 0 rotates the buffers right after its first polling round when every warp of the grid + has counted itself by then (as a rule), so that the acquire load and the state stores are + not on the grid's tail; otherwise it waits for the count at the end, like the cluster kernel. + * ``early_state`` (ULL2 coll): the rotation state read, the warps' arrival count and the + re-arm of the previous call's buffer (nobody polls it any more) run before + ``griddepcontrol.wait``, while the producer of ``x`` still runs: the previous call on the + workspace completed before that producer started (stream order), so the state is final. + The producer may read the state (``prepushed``): the rotation still follows the wait. Same + bits. + """ + + def __init__( + self, + world_size: int, + use_pdl: bool, + norm: bool = True, + mxfp8: bool = False, + sf_blocks: int = 0, + prepushed: bool = False, + early_state: bool = False, + ): + # prepushed: the producer already stored this rank's sanitized rows into slot (row, rank) of + # every rank's current buffer (an in-kernel push); the kernel only polls and sums + # (``x`` still supplies the own slot's values, so the result is bitwise the same). + self.prepushed = prepushed + self.early_state = early_state and use_pdl + if world_size not in ROW_WORLD_SIZES: + raise ValueError(f"unsupported world size {world_size}") + if mxfp8 and not norm: + raise ValueError("the MXFP8 output is the normalized output's") + if sf_blocks and not mxfp8: + raise ValueError("sf_blocks sets the layout of the MXFP8 output's scales") + self.world_size = world_size + self.use_pdl = use_pdl + self.norm = norm + self.mxfp8 = mxfp8 + # c240: 0 stores the UE8M0 exponents row-major [M, 3584 / 32] (the low-M MXFP8 GEMM's + # operand); > 0 in the mid-M GEMM's tcgen05 atoms [3584 / 128][sf_blocks][512 B] + # (k3_fp8w_lowm_gemm.K3MxFp8QuantKernel's layout), rows M .. 128 sf_blocks - 1 neutral. + self.sf_blocks = sf_blocks + self.rank_chunk = min(RANK_CHUNK, world_size) + + @cute.jit + def __call__( + self, + x_addr: Int64, + gamma_addr: Int64, + out_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + uc_addr: Int64, + mc_addr: Int64, + flags_addr: Int64, + rank: Int32, + m: Int32, + eps: Float32, + stream: cuda.CUstream, + ): + @cute.struct + class SharedStorage: + warp_sum: cute.struct.MemRange[cutlass.Float32, ROW_WARPS] + + self.shared_storage = SharedStorage + self.kernel( + x_addr, + gamma_addr, + out_addr, + xq_addr, + xs_addr, + uc_addr, + mc_addr, + flags_addr, + rank, + m, + eps, + ).launch( + grid=(m, 1, 1), + block=(ROW_THREADS, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + x_addr: Int64, + gamma_addr: Int64, + out_addr: Int64, + xq_addr: Int64, + xs_addr: Int64, + uc_addr: Int64, + mc_addr: Int64, + flags_addr: Int64, + rank: Int32, + m: Int32, + eps: Float32, + ): + ws = self.world_size + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + grp = tidx // ROW_GROUP_THREADS + t = tidx - grp * ROW_GROUP_THREADS + active = t < THREADS + # The group's fourth warp is the half warp (lanes 0-15 active). + half = (warp % WARPS) == WARPS - 1 + storage = utils.SmemAllocator().allocate(self.shared_storage) + warp_sum = storage.warp_sum.data_ptr().toint() + + if cutlass.const_expr(self.use_pdl and not self.early_state): + _griddepcontrol_wait() + _griddepcontrol_launch_dependents() + + f0, f1, f2, f3 = _ld_v4(flags_addr) + c0, c1, c2, c3 = _ld_v4(flags_addr + Int64(16)) + counter = flags_addr + Int64(32) + cute.arch.sync_warp() + if lane == 0: + _red_add_release_gpu(counter, Int32(1)) + cur = Int64(f0) + dirty = Int64(f1) + buf_bytes = Int64(f2) & Int64(0xFFFFFFFF) + dirty_stages = f3 + # An inactive thread loads and polls the columns of lane t - 16 of its half warp. + tc = t + if t >= THREADS: + tc = t - (ROW_GROUP_THREADS - THREADS) + col = (grp * THREADS + tc) * ELTS + cur_off = cur * buf_bytes + + if cutlass.const_expr(self.early_state): + # early_state: the re-arm of the previous call's buffer also before the wait, then the + # wait. + if cutlass.const_expr(self.norm): + gw = _ld_v4(gamma_addr + Int64(col * 2)) + n_threads = m * ROW_THREADS + gtid = row * ROW_THREADS + tidx + clear = (c0, c1, c2, c3) + for s in cutlass.range_constexpr(4): + if s < dirty_stages: + stage_bytes = buf_bytes // Int64(dirty_stages) + base = uc_addr + (dirty * Int64(dirty_stages) + Int64(s)) * stage_bytes + n_vec = (clear[s] + 15) >> 4 + for p in cutlass.range(gtid, n_vec, n_threads, unroll=1): + _st_sentinel_v4(base + Int64(p) * Int64(16)) + _griddepcontrol_wait() + _griddepcontrol_launch_dependents() + + # Broadcast this rank's eight values into slot (row, rank) of every rank's buffer. + xw = _ld_v4(x_addr + Int64((row * HIDDEN + col) * 2)) + if cutlass.const_expr(self.norm and not self.early_state): + gw = _ld_v4(gamma_addr + Int64(col * 2)) + sw = [_sanitize(w) for w in xw] + if cutlass.const_expr(not self.prepushed): + if active: + slot = Int64(((row * ws + rank) * HIDDEN + col) * 2) + _st_v4(mc_addr + cur_off + slot, sw[0], sw[1], sw[2], sw[3]) + + if cutlass.const_expr(not self.early_state): + # Clear the previous call's buffer. + n_threads = m * ROW_THREADS + gtid = row * ROW_THREADS + tidx + clear = (c0, c1, c2, c3) + for s in cutlass.range_constexpr(4): + if s < dirty_stages: + stage_bytes = buf_bytes // Int64(dirty_stages) + base = uc_addr + (dirty * Int64(dirty_stages) + Int64(s)) * stage_bytes + n_vec = (clear[s] + 15) >> 4 + for p in cutlass.range(gtid, n_vec, n_threads, unroll=1): + _st_sentinel_v4(base + Int64(p) * Int64(16)) + + # Poll every slot; the own slot comes from the registers. Bit r of `have`: slot r arrived. + row_base = uc_addr + cur_off + Int64((row * ws * HIDDEN + col) * 2) + words = cute.make_rmem_tensor((ws * 4,), Int32) + have = Int32(0) + for rr in cutlass.range_constexpr(ws): + w = _ld_vol_v4(row_base + Int64(rr * HIDDEN * 2)) + for e in cutlass.range_constexpr(4): + words[rr * 4 + e] = w[e] + if rr == rank: + for e in cutlass.range_constexpr(4): + words[rr * 4 + e] = sw[e] + have = have | ( + _clean4(words[rr * 4], words[rr * 4 + 1], words[rr * 4 + 2], words[rr * 4 + 3]) + << rr + ) + # CTA 0 rotates the buffers once every warp of the grid has read the flags: here, while it + # waits for the other ranks, if they all have by the end of its first polling round (as a + # rule), else at the end. + rotated = Int32(0) + if tidx == 0: + if row == 0: + if _ld_acquire_gpu_u32(counter) >= m * ROW_WARPS: + nxt0 = f0 + 1 + if nxt0 == NUM_LAMPORT_BUFFERS: + nxt0 = Int32(0) + _st_v4(flags_addr, nxt0, f0, f2, Int32(1)) + _st_v4( + flags_addr + Int64(16), m * (HIDDEN * ws * 2), Int32(0), Int32(0), Int32(0) + ) + _st_u32(counter, Int32(0)) + rotated = Int32(1) + full = Int32((1 << ws) - 1) + while have != full: + missing = have + for rr in cutlass.range_constexpr(ws): + w = _ld_vol_v4_pred( + row_base + Int64(rr * HIDDEN * 2), + ((missing >> rr) & Int32(1)) ^ Int32(1), + words[rr * 4], + words[rr * 4 + 1], + words[rr * 4 + 2], + words[rr * 4 + 3], + ) + for e in cutlass.range_constexpr(4): + words[rr * 4 + e] = w[e] + for rr in cutlass.range_constexpr(ws): + have = have | ( + _clean4(words[rr * 4], words[rr * 4 + 1], words[rr * 4 + 2], words[rr * 4 + 3]) + << rr + ) + # Sum in chunks of RANK_CHUNK ranks (each chunk from 0, chunks in rank order). + acc = [Float32(0.0)] * ELTS + for rb in cutlass.range_constexpr(0, ws, self.rank_chunk): + chunk = [Float32(0.0)] * ELTS + for rr in cutlass.range_constexpr(self.rank_chunk): + for j in cutlass.range_constexpr(4): + lo, hi = _bf16x2_f32(words[(rb + rr) * 4 + j]) + chunk[2 * j] = _fadd(chunk[2 * j], lo) + chunk[2 * j + 1] = _fadd(chunk[2 * j + 1], hi) + for i in cutlass.range_constexpr(ELTS): + acc[i] = _fadd(acc[i], chunk[i]) + red = [_pack_bf16x2(acc[2 * j], acc[2 * j + 1]) for j in range(4)] + + if cutlass.const_expr(self.norm): + thread_sum = Float32(0.0) + for j in cutlass.range_constexpr(4): + lo, hi = _bf16x2_f32(_sq_bf16x2(red[j])) + thread_sum = _fadd(thread_sum, lo) + thread_sum = _fadd(thread_sum, hi) + v = thread_sum + if half: + if lane < 16: + v = _fadd(v, Float32(0.0)) + for off in (8, 4, 2, 1): + v = _fadd(v, _shfl_bfly(v, off, 0xFFFF)) + else: + for off in (16, 8, 4, 2, 1): + v = _fadd(v, _shfl_bfly(v, off, 0xFFFFFFFF)) + if lane == 0: + _sts_f32(warp_sum + warp * 4, v) + _bar_sync() + # The cluster kernel's butterfly over a group's four warp sums (lanes past them add 0) + # yields (w0 + w2) + (w1 + w3); every thread forms it for each group and adds the + # group sums in order, so one CTA barrier is enough. + total = Float32(0.0) + for c in cutlass.range_constexpr(ROW_GROUPS): + w4 = [_lds_f32(warp_sum + (c * WARPS + k) * 4) for k in range(WARPS)] + total = _fadd(total, _fadd(_fadd(w4[0], w4[2]), _fadd(w4[1], w4[3]))) + scale = _rms_scale(total, eps) + if active: + ow = [] + for j in cutlass.range_constexpr(4): + x_lo, x_hi = _bf16x2_f32(red[j]) + g_lo, g_hi = _bf16x2_f32(gw[j]) + ow.append( + _pack_bf16x2( + _fmul(_fmul(x_lo, scale), g_lo), _fmul(_fmul(x_hi, scale), g_hi) + ) + ) + _st_v4(out_addr + Int64((row * HIDDEN + col) * 2), ow[0], ow[1], ow[2], ow[3]) + if cutlass.const_expr(self.mxfp8): + amax = _amax_bf16x8(ow[0], ow[1], ow[2], ow[3]) + if half: + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 1, 0xFFFF)) + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 2, 0xFFFF)) + else: + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 1, 0xFFFFFFFF)) + amax = _amax_pair(amax, _shfl_bfly_b32(amax, 2, 0xFFFFFFFF)) + e = _ue8m0_of_amax(amax) + q0, q1 = _quant_bf16x8_to_e4m3(ow[0], ow[1], ow[2], ow[3], _bf16x2_inv_scale(e)) + _st_v2(xq_addr + Int64(row * HIDDEN + col), q0, q1) + if cutlass.const_expr(self.sf_blocks == 0): + if (t & 3) == 0: + _st_u8(xs_addr + Int64(row * (HIDDEN // 32) + (col >> 5)), e) + else: + if (t & 3) == 0: + # tcgen05 scale atom of (row, 32-group sf_g), as K3MxFp8QuantKernel. + sf_g = col >> 5 + sf_r = row & 127 + sf_off = ((sf_g >> 2) * self.sf_blocks + (row >> 7)) * 512 + sf_off = sf_off + (sf_r & 31) * 16 + (sf_r >> 5) * 4 + (sf_g & 3) + _st_u8(xs_addr + Int64(sf_off), e) + # The padding rows M .. 128 sf_blocks - 1 have no CTA: neutral 127. + for pad in cutlass.range(row + m, self.sf_blocks * 128, m, unroll=1): + pad_r = pad & 127 + pad_off = ((sf_g >> 2) * self.sf_blocks + (pad >> 7)) * 512 + pad_off = ( + pad_off + (pad_r & 31) * 16 + (pad_r >> 5) * 4 + (sf_g & 3) + ) + _st_u8(xs_addr + Int64(pad_off), Int32(127)) + else: + if active: + _st_v4(out_addr + Int64((row * HIDDEN + col) * 2), red[0], red[1], red[2], red[3]) + + # CTA 0: rotate the buffers now if it could not after its first polling round. + if tidx == 0: + if row == 0: + if rotated == 0: + arrived = _ld_acquire_gpu_u32(counter) + while arrived < m * ROW_WARPS: + arrived = _ld_acquire_gpu_u32(counter) + nxt = f0 + 1 + if nxt == NUM_LAMPORT_BUFFERS: + nxt = Int32(0) + _st_v4(flags_addr, nxt, f0, f2, Int32(1)) + _st_v4( + flags_addr + Int64(16), m * (HIDDEN * ws * 2), Int32(0), Int32(0), Int32(0) + ) + _st_u32(counter, Int32(0)) + + +@functools.cache +def _compiled_row( + world_size: int, + use_pdl: bool, + norm: bool, + mxfp8: bool, + sf_blocks: int = 0, + prepushed: bool = False, + early_state: bool = False, +): + kernel = K3LatentArRowKernel( + world_size, + use_pdl, + norm, + mxfp8, + sf_blocks, + prepushed=prepushed, + early_state=early_state, + ) + a = Int64(0) + return cute.compile( + kernel, + a, + a, + a, + a, + a, + a, + a, + a, + Int32(0), + Int32(1), + Float32(1e-6), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def can_run_row( + x: torch.Tensor, gamma: torch.Tensor, world_size: int, buffer_bytes: int | None = None +) -> bool: + """Whether :func:`run_k3_latent_ar_norm_row` serves this all-reduce (non-raising check); with + ``buffer_bytes``, also that one Lamport buffer of that size holds every rank's rows.""" + if buffer_bytes is not None and x.dim() == 2 and x.numel() * world_size * 2 > buffer_bytes: + return False + return ( + x.is_cuda + and _is_supported_device(x.device.index) + and world_size in ROW_WORLD_SIZES + and x.dtype == torch.bfloat16 + and x.dim() == 2 + and x.shape[1] == HIDDEN + and 1 <= x.shape[0] <= ROW_MAX_M + and x.is_contiguous() + and x.data_ptr() % 16 == 0 + and gamma.dtype == torch.bfloat16 + and tuple(gamma.shape) == (HIDDEN,) + and gamma.is_contiguous() + and gamma.data_ptr() % 16 == 0 + ) + + +def run_k3_latent_ar_norm_row( + x: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + *, + uc_addr: int, + mc_addr: int, + flags: torch.Tensor, + rank: int, + world_size: int, + use_pdl: bool = True, + norm: bool = True, + mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, + sf_blocks: int = 0, + prepushed: bool = False, + early_state: bool = False, +) -> None: + """:func:`run_k3_latent_ar_norm` with :class:`K3LatentArRowKernel` (same result, bit for bit, + for up to :data:`ROW_MAX_M` rows). Each Lamport buffer of the workspace must hold the rows of + every rank (``M x world_size x 3584 x 2`` bytes, as the one-shot sizing of the MNNVL + workspace guarantees for the one-shot sizes; see :func:`can_run_row`). ``sf_blocks`` > 0 + (c240) stores ``mxfp8_out``'s scales ``[(3584 / 128) * sf_blocks * 512]`` in the mid-M GEMM's + layout (``k3_fp8w_lowm_gemm.mid_prequant_plan``) instead of row-major ``[M, 112]``. + ``prepushed``: the producer of ``x`` already stored it into this rank's slot of every rank's + current buffer (an in-kernel all-reduce push); the kernel skips its push. + ``early_state`` (ULL2 coll): the state read, the arrival count and the re-arm of the previous + buffer run before the kernel's ``griddepcontrol.wait`` (see :class:`K3LatentArRowKernel`).""" + mxfp8 = mxfp8_out is not None + if sf_blocks and not (mxfp8 and x.shape[0] <= sf_blocks * 128): + raise ValueError(f"sf_blocks={sf_blocks} needs the MXFP8 output and M <= {sf_blocks * 128}") + _compiled_row(world_size, use_pdl, norm, mxfp8, sf_blocks, prepushed, early_state)( + Int64(x.data_ptr()), + Int64(gamma.data_ptr()), + Int64(out.data_ptr()), + Int64(mxfp8_out[0].data_ptr() if mxfp8 else 0), + Int64(mxfp8_out[1].data_ptr() if mxfp8 else 0), + Int64(uc_addr), + Int64(mc_addr), + Int64(flags.data_ptr()), + Int32(rank), + Int32(x.shape[0]), + Float32(eps), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + + +__all__ = [ + "HIDDEN", + "MAX_M", + "ROW_MAX_M", + "ROW_WORLD_SIZES", + "WORLD_SIZES", + "K3LatentArNormKernel", + "K3LatentArRowKernel", + "can_run", + "can_run_row", + "run_k3_latent_ar_norm", + "run_k3_latent_ar_norm_row", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_latent_fp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_latent_fp8_gemm.py new file mode 100644 index 000000000000..9f8a9b2df201 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_latent_fp8_gemm.py @@ -0,0 +1,1655 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 latent MoE projection as an FP8 GEMM with a fused epilogue (SM100/SM103/SM107). + +The projection uses an FP8 E4M3 weight ``W [N, K]`` with one power-of-two (UE8M0) scale per +128x128 block (the resmoothed weight and DeepGEMM-packed scale that ``fp8_swap_ab_gemm`` call +sites hold; on SM107 the scale is rebuilt from the CuTe MXFP8 layout, see +``k3_fp8w_lowm_gemm.pack_block_scale_from_mxfp8_layout``) and a BF16 activation ``x [M, K]``: + +* :func:`k3_latent_down_fp4`: ``q, sf = nvfp4_quant(bf16(x @ dequant(W)^T), global_sf)``, + the routed-expert GEMM1 input: E2M1 codes ``[M, N / 2]`` (two per byte, low nibble first) and + one E4M3 scale per 16 outputs ``[M, N / 16]`` in the linear layout. Replaces the activation + quantization [+ UE8M0 packing], DeepGEMM and ``fp4_quantize``. + +The op is two kernels chained with PDL (on SM107, up to :func:`lowm_route_max_m` tokens the +op instead runs the split-K GEMM of :mod:`k3_fp8w_lowm_gemm` with the same epilogue, see +below): + +* :class:`K3Fp8Quant1x128Kernel` quantizes the activation to E4M3 with one UE8M0 scale per + 1x128 block, using the exact ``2^ceil(log2(amax / 448))`` scale of the Triton and CUDA + ``fp8_quantize_1x128`` kernels, and writes the scale replicated over the four K32 slots of the + tcgen05 scale-factor atom (``[K/128][ceil(M/128)][512 B]``, atom byte + ``(r % 32) * 16 + (r // 32) * 4 + k32``) so the GEMM moves one contiguous 512-byte block per + K tile. +* :class:`K3Fp8LatentGemmKernel` is a persistent MXFP8 GEMM (``tcgen05.mma.kind::mxf8f6f4`` with + 1x32 scale vectors, i.e. the 1x128 scales replicated) over 128 x ``bn`` output tiles with no + split-K, so every output is the same FP32 accumulation chain over K as DeepGEMM's. Warp 0 + loads the activation and weight tiles with TMA plus the activation scale block with a bulk + copy (the first weight stages do not wait on the previous grid), warp 1 issues the MMAs into + a TMEM accumulator that is double-buffered when two fit, warp 2 writes the weight scale atoms + (the 128x128 block exponents, read from the DeepGEMM-packed scale) and warps 4-7 run the + epilogue on 32-column TMEM blocks: BF16 rounding, then either NVFP4 quantization with exactly + the arithmetic of ``fp4_quantize`` (``quantize_with_block_size``) or the residual adds + ``residual + bf16(bf16(acc) + shared)``, and 16-byte global stores. + +The persistent GEMM has one 128 x ``bn`` tile per CTA and no split-K, so a decode batch +(``M <= 128``) runs only ``N / bn`` CTAs (28 for latent_down), each limited by the weight bytes +its ring keeps in flight. On SM107 the op therefore runs the split-K cluster GEMM of +``k3_fp8w_lowm_gemm`` (1x32 MXFP8 activation scales instead of 1x128) with its NVFP4 epilogue +for ``M <= lowm_route_max_m(N, K)``, at the tactic of :func:`route_tactic`; the results equal +that GEMM at the same tactic followed by ``fp4_quantize``. +""" + +from __future__ import annotations + +import functools + +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 +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass import Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +#: K extent of one pipeline stage; also the activation / weight scale block length. +CTA_K = 128 +#: Output rows per CTA tile (UMMA M). +BM = 128 +#: tcgen05 MXF8 scale vector length. +SF_VEC = 32 +#: NVFP4 scale vector length of the output. +FP4_VEC = 16 +#: Supported output-column tiles (UMMA N). +SUPPORTED_BN = (128, 192, 256) +#: Largest M the op accepts. +MAX_M = 4096 + +_TMA_WARP = 0 +_MMA_WARP = 1 +_SFB_WARP = 2 +_EPI_WARP0 = 4 +_THREADS = 256 +_SFA_TMEM_COLS = 16 + +_SMEM_CAPACITY = utils.get_smem_capacity_in_bytes("sm_100") + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm_pure(res, operands, text, constraints, *, loc=None, ip=None): + """Inline PTX without side effects: the compiler may schedule or drop it freely.""" + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +def _fv(value, *, loc=None, ip=None): + return cutlass.Float32(value).ir_value(loc=loc, ip=ip) + + +def _struct(res, n, ty, *, loc=None, ip=None): + return tuple(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip) for i in range(n)) + + +@dsl_user_op +def _ldg_v4_b32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return tuple(Int32(v) for v in _struct(res, 4, T.i32(), loc=loc, ip=ip)) + + +@dsl_user_op +def _ldg_f32(addr, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.f32 $0, [$1];", + "=f,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _stg_v4_b32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u8(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u8 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u16(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u16 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_v4_b32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_iv(x, loc=loc, ip=ip) for x in (addr, v0, v1, v2, v3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_load_g2s(dst_smem, src_gmem, nbytes, mbar, *, loc=None, ip=None): + """Bulk-copy ``nbytes`` from global memory into local SMEM, crediting ``mbar``.""" + _asm( + None, + [ + _iv(dst_smem, loc=loc, ip=ip), + Int64(src_gmem).ir_value(loc=loc, ip=ip), + _iv(nbytes, loc=loc, ip=ip), + _iv(mbar, loc=loc, ip=ip), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];", + "r,l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.wait;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.launch_dependents;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _min_s32(a, b, *, loc=None, ip=None): + return Int32( + _asm_pure( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "min.s32 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _abs_max_bf16x16(words, *, loc=None, ip=None): + """|x| max of 16 BF16 values (8 packed words) as two packed BF16 magnitudes (sign garbage).""" + body = "".join(f"max.xorsign.abs.bf16x2 m{i}, ${2 * i + 1}, ${2 * i + 2};\n" for i in range(4)) + body += ( + "max.xorsign.abs.bf16x2 m0, m0, m1;\n" + "max.xorsign.abs.bf16x2 m2, m2, m3;\n" + "max.xorsign.abs.bf16x2 $0, m0, m2;\n" + ) + return Int32( + _asm_pure( + T.i32(), + [_iv(w, loc=loc, ip=ip) for w in words], + "{\n.reg .b32 m0, m1, m2, m3;\n" + body + "}", + "=r," + ",".join(["r"] * 8), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _max_bf16_mag(a, b, *, loc=None, ip=None): + """Packed BF16 magnitude max of two packed words (sign garbage, masked by the consumer).""" + return Int32( + _asm_pure( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ue8m0_of_amax(packed_amax, *, loc=None, ip=None): + """Biased exponent of ``2^ceil(log2(max(amax, 1e-10) / 448))`` from two packed BF16 magnitudes. + + For ``amax = 2^E * (1 + f / 128)`` (a BF16 value): ``amax / 448 <= 2^(E - 8)`` iff + ``1 + f / 128 <= 1.75``, so the biased exponent is ``E - 8`` plus one when the 7-bit + mantissa exceeds ``0x60``. Magnitudes below 1e-10 (only zero for real activations) use the + exponent of 1e-10, ``2^-42`` (biased 85). + """ + return Int32( + _asm_pure( + T.i32(), + [_iv(packed_amax, loc=loc, ip=ip)], + "{\n" + ".reg .b32 lo, hi, a, e, mant;\n" + ".reg .pred p, z;\n" + "and.b32 lo, $1, 0x7fff;\n" + "shr.u32 hi, $1, 16;\n" + "and.b32 hi, hi, 0x7fff;\n" + "max.u32 a, lo, hi;\n" + "shr.u32 e, a, 7;\n" + "and.b32 mant, a, 0x7f;\n" + "setp.gt.u32 p, mant, 0x60;\n" + "sub.u32 e, e, 8;\n" + "@p add.u32 e, e, 1;\n" + # BF16 1e-10 rounds to 0x2edc; anything at or below the FP32 1e-10 floor takes + # its scale. 0x2edb..0x2edc straddle 1e-10, so compare against the exact bits. + "setp.lt.u32 z, a, 0x2edc;\n" + "@z mov.b32 e, 85;\n" + "mov.b32 $0, e;\n" + "}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bf16x2_inv_scale(ue8m0, *, loc=None, ip=None): + """Packed BF16x2 ``2^(127 - ue8m0)``.""" + return Int32( + _asm_pure( + T.i32(), + [_iv(ue8m0, loc=loc, ip=ip)], + "{\n" + ".reg .b32 s, t;\n" + "sub.u32 s, 254, $1;\n" + "shl.b32 s, s, 7;\n" + "shl.b32 t, s, 16;\n" + "or.b32 $0, s, t;\n" + "}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _quant_bf16x8_to_e4m3(w0, w1, w2, w3, inv_scale, *, loc=None, ip=None): + """Scale eight BF16 values by the packed power-of-two ``inv_scale`` and round to E4M3. + + BF16 has the FP32 exponent range, so the BF16 multiply by a power of two is exact (no + result is subnormal for scales from real activations) and the only rounding is the + satfinite round-to-nearest-even conversion, as in ``clamp(x / s).to(float8_e4m3fn)``. + """ + res = _asm_pure( + llvm.StructType.get_literal([T.i32()] * 2), + [_iv(x, loc=loc, ip=ip) for x in (w0, w1, w2, w3, inv_scale)], + "{\n" + ".reg .b32 p0, p1, p2, p3;\n" + ".reg .b16 h0, h1, h2, h3;\n" + "mul.rn.bf16x2 p0, $2, $6;\n" + "mul.rn.bf16x2 p1, $3, $6;\n" + "mul.rn.bf16x2 p2, $4, $6;\n" + "mul.rn.bf16x2 p3, $5, $6;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h0, p0;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h1, p1;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h2, p2;\n" + "cvt.rn.satfinite.e4m3x2.bf16x2 h3, p3;\n" + "mov.b32 $0, {h0, h1};\n" + "mov.b32 $1, {h2, h3};\n" + "}", + "=r,=r,r,r,r,r,r", + loc=loc, + ip=ip, + ) + return tuple(Int32(v) for v in _struct(res, 2, T.i32(), loc=loc, ip=ip)) + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm_pure( + T.f32(), [_fv(x, loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _bf16_round_x2(a, b, *, loc=None, ip=None): + """Round two FP32 values to BF16 (RN) and return them widened back to FP32.""" + res = _asm_pure( + llvm.StructType.get_literal([T.f32()] * 2), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "{\n" + ".reg .b32 p;\n" + ".reg .b16 l, h;\n" + "cvt.rn.bf16x2.f32 p, $3, $2;\n" + "mov.b32 {l, h}, p;\n" + "cvt.f32.bf16 $0, l;\n" + "cvt.f32.bf16 $1, h;\n" + "}", + "=f,=f,f,f", + loc=loc, + ip=ip, + ) + return tuple(cutlass.Float32(v) for v in _struct(res, 2, T.f32(), loc=loc, ip=ip)) + + +@dsl_user_op +def _abs_max_f32x16(vals, *, loc=None, ip=None): + body = "".join(f"abs.f32 t{i}, ${i + 1};\n" for i in range(16)) + for step in (8, 4, 2, 1): + body += "".join(f"max.f32 t{i}, t{i}, t{i + step};\n" for i in range(step)) + body += "mov.f32 $0, t0;\n" + return cutlass.Float32( + _asm_pure( + T.f32(), + [_fv(v, loc=loc, ip=ip) for v in vals], + "{\n.reg .f32 " + ", ".join(f"t{i}" for i in range(16)) + ";\n" + body + "}", + "=f," + ",".join(["f"] * 16), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _nvfp4_scales_x2(amax0, amax1, gsf, rgsf, rcp6, *, loc=None, ip=None): + """NVFP4 scales of two 16-element groups, as in ``cvt_warp_fp16_to_fp4`` (E4M3 scale). + + ``sf = e4m3(gsf * (amax * rcp(6)))`` and ``out_scale = amax != 0 ? rcp(f32(sf) * rcp(gsf)) : 0`` + with ``rcp`` = ``rcp.approx.ftz.f32``. Returns the two scale bytes packed in 16 bits (group + 0 in the low byte) and the two output scales. + """ + res = _asm_pure( + llvm.StructType.get_literal([T.i32(), T.f32(), T.f32()]), + [_fv(x, loc=loc, ip=ip) for x in (amax0, amax1, gsf, rgsf, rcp6)], + "{\n" + ".reg .f32 s0, s1, f0, f1;\n" + ".reg .b16 h, l16, h16;\n" + ".reg .b32 hw;\n" + ".reg .pred z0, z1;\n" + "mul.rn.f32 s0, $3, $7;\n" + "mul.rn.f32 s1, $4, $7;\n" + "mul.rn.f32 s0, $5, s0;\n" + "mul.rn.f32 s1, $5, s1;\n" + "cvt.rn.satfinite.e4m3x2.f32 h, s1, s0;\n" + "cvt.rn.f16x2.e4m3x2 hw, h;\n" + "mov.b32 {l16, h16}, hw;\n" + "cvt.f32.f16 f0, l16;\n" + "cvt.f32.f16 f1, h16;\n" + "mul.rn.f32 f0, f0, $6;\n" + "mul.rn.f32 f1, f1, $6;\n" + "rcp.approx.ftz.f32 f0, f0;\n" + "rcp.approx.ftz.f32 f1, f1;\n" + "setp.eq.f32 z0, $3, 0f00000000;\n" + "setp.eq.f32 z1, $4, 0f00000000;\n" + "selp.f32 $1, 0f00000000, f0, z0;\n" + "selp.f32 $2, 0f00000000, f1, z1;\n" + "cvt.u32.u16 $0, h;\n" + "}", + "=r,=f,=f,f,f,f,f,f", + loc=loc, + ip=ip, + ) + sf = llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip) + o0 = llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip) + o1 = llvm.extractvalue(T.f32(), res, [2], loc=loc, ip=ip) + return Int32(sf), cutlass.Float32(o0), cutlass.Float32(o1) + + +@dsl_user_op +def _e2m1x8(vals, scale, *, loc=None, ip=None): + """Scale eight FP32 values (``mul.rn``) and pack them as E2M1 into 32 bits, element 0 lowest.""" + body = "".join(f"mul.rn.f32 t{i}, ${i + 1}, $9;\n" for i in range(8)) + body += "".join( + f"cvt.rn.satfinite.e2m1x2.f32 b{j}, t{2 * j + 1}, t{2 * j};\n" for j in range(4) + ) + body += "mov.b32 $0, {b0, b1, b2, b3};\n" + return Int32( + _asm_pure( + T.i32(), + [_fv(v, loc=loc, ip=ip) for v in vals] + [_fv(scale, loc=loc, ip=ip)], + "{\n.reg .f32 t0, t1, t2, t3, t4, t5, t6, t7;\n.reg .b8 b0, b1, b2, b3;\n" + body + "}", + "=r," + ",".join(["f"] * 9), + loc=loc, + ip=ip, + ) + ) + + +# --------------------------------------------------------------------------- +# Activation quantization +# --------------------------------------------------------------------------- + + +class K3Fp8Quant1x128Kernel: + """``x`` BF16 ``[M, K]`` -> E4M3 ``[M, K]`` + UE8M0 1x128 scales in the UMMA scale atom layout. + + A work item is one (row, 32-element group); the four groups of a 128-element block are + handled by four consecutive lanes, which share the block's absolute maximum through two + shuffles. Each lane writes its own (replicated) scale byte. Rows ``M .. 128 * ceil(M/128)`` + get the neutral scale 1.0 (their FP8 rows are TMA zero-fill in the GEMM). With + ``sf_m_major`` the scale atoms are ordered ``[ceil(M/128)][K/128][512 B]`` (the CUTLASS + block-scaled layout) instead of ``[K/128][ceil(M/128)][512 B]``. + """ + + def __init__( + self, k: int, use_pdl: bool, threads: int = 128, items: int = 2, sf_m_major: bool = False + ): + if k % CTA_K: + raise ValueError(f"K={k} must be a multiple of {CTA_K}") + self.k = k + self.groups = k // SF_VEC + self.use_pdl = use_pdl + self.sf_m_major = sf_m_major + self._THREADS = threads + self._ITEMS = items + + @cute.jit + def __call__( + self, + x_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ): + items = ((m + BM - 1) // BM) * BM * self.groups + per_cta = self._THREADS * self._ITEMS + self.kernel(x_ptr, xq_ptr, xs_ptr, m).launch( + grid=((items + per_cta - 1) // per_cta, 1, 1), + block=(self._THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel(self, x_ptr: cute.Pointer, xq_ptr: cute.Pointer, xs_ptr: cute.Pointer, m: Int32): + if cutlass.const_expr(self.use_pdl): + # The GEMM reads this grid's outputs only after its own griddepcontrol.wait; let it + # launch now so its weight stream overlaps this kernel. + _griddepcontrol_launch_dependents() + _griddepcontrol_wait() + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + m_tiles = (m + BM - 1) // BM + items = m_tiles * BM * self.groups + stride = gdim * self._THREADS + ipt = self._ITEMS + rX = cute.make_rmem_tensor((ipt * 16,), Int32) + # Issue every load first; padded or out-of-range items re-read a valid row. + for it in cutlass.range_constexpr(ipt): + item = _min_s32(bidx * self._THREADS + tidx + it * stride, items - 1) + r = _min_s32(item // self.groups, m - 1) + g = item - (item // self.groups) * self.groups + src = x_ptr.toint() + Int64(r * self.k + g * SF_VEC) * 2 + for v in cutlass.range_constexpr(4): + w0, w1, w2, w3 = _ldg_v4_b32(src + v * 16) + rX[it * 16 + v * 4] = w0 + rX[it * 16 + v * 4 + 1] = w1 + rX[it * 16 + v * 4 + 2] = w2 + rX[it * 16 + v * 4 + 3] = w3 + for it in cutlass.range_constexpr(ipt): + item = bidx * self._THREADS + tidx + it * stride + # Lanes 4i..4i+3 always hold the four groups of one 128-element block (the item + # count and the per-pass stride are multiples of 4), so the shuffles stay inside + # the block even for out-of-range quads. + ws = [rX[it * 16 + i] for i in range(16)] + amax = _max_bf16_mag(_abs_max_bf16x16(ws[:8]), _abs_max_bf16x16(ws[8:])) + amax = _max_bf16_mag(amax, cute.arch.shuffle_sync_bfly(amax, offset=1)) + amax = _max_bf16_mag(amax, cute.arch.shuffle_sync_bfly(amax, offset=2)) + if item < items: + r = item // self.groups + g = item - r * self.groups + rr = r & (BM - 1) + if cutlass.const_expr(self.sf_m_major): + atom = (r >> 7) * (self.k // CTA_K) + (g >> 2) + else: + atom = (g >> 2) * m_tiles + (r >> 7) + sf_off = atom * 512 + (rr & 31) * 16 + (rr >> 5) * 4 + sf_addr = xs_ptr.toint() + Int64(sf_off + (g & 3)) + if r < m: + e_a = _ue8m0_of_amax(amax) + inv = _bf16x2_inv_scale(e_a) + packed = [] + for v in cutlass.range_constexpr(4): + p0, p1 = _quant_bf16x8_to_e4m3(*ws[4 * v : 4 * v + 4], inv) + packed.extend((p0, p1)) + dst = xq_ptr.toint() + Int64(r * self.k + g * SF_VEC) + _stg_v4_b32(dst, *packed[:4]) + _stg_v4_b32(dst + 16, *packed[4:]) + _stg_u8(sf_addr, e_a) + else: + _stg_u8(sf_addr, Int32(127)) + + +# --------------------------------------------------------------------------- +# GEMM with fused epilogues +# --------------------------------------------------------------------------- + +#: Residual epilogue: per-warp SMEM staging of a 32-row x 32-column BF16 block, rows padded to +#: 80 bytes so the row-per-thread 16-byte stores are bank-conflict free. +_RES_ROW_BYTES = 80 +_RES_WARP_BYTES = 32 * _RES_ROW_BYTES +#: Residual epilogue: blocks of shared / residual words loaded ahead of their use (deeper +#: prefetch measured no faster on GB300). +_RES_PREFETCH = 1 + +#: NVFP4-quantized output (codes + linear E4M3 scales) for the routed-expert GEMM1. +EPI_NVFP4 = "nvfp4" +#: BF16 output plus the MoE residual update ``residual + bf16(bf16(acc) + shared)``. +EPI_RESIDUAL = "residual" + + +@dsl_user_op +def _residual_from_bf16x2(t, shared, residual, *, loc=None, ip=None): + """``bf16(r + bf16(t + s))`` for two BF16 columns, each add in FP32 as PyTorch does.""" + return Int32( + _asm_pure( + T.i32(), + [_iv(x, loc=loc, ip=ip) for x in (t, shared, residual)], + "{\n" + ".reg .b32 u;\n" + ".reg .b16 l, h;\n" + ".reg .f32 a0, a1, b0, b1;\n" + "mov.b32 {l, h}, $1;\n" + "cvt.f32.bf16 a0, l;\n" + "cvt.f32.bf16 a1, h;\n" + "mov.b32 {l, h}, $2;\n" + "cvt.f32.bf16 b0, l;\n" + "cvt.f32.bf16 b1, h;\n" + "add.rn.f32 a0, a0, b0;\n" + "add.rn.f32 a1, a1, b1;\n" + "cvt.rn.bf16x2.f32 u, a1, a0;\n" + "mov.b32 {l, h}, u;\n" + "cvt.f32.bf16 a0, l;\n" + "cvt.f32.bf16 a1, h;\n" + "mov.b32 {l, h}, $3;\n" + "cvt.f32.bf16 b0, l;\n" + "cvt.f32.bf16 b1, h;\n" + "add.rn.f32 a0, b0, a0;\n" + "add.rn.f32 a1, b1, a1;\n" + "cvt.rn.bf16x2.f32 $0, a1, a0;\n" + "}", + "=r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _f32x2_to_bf16x2(a, b, *, loc=None, ip=None): + """Round two FP32 values to BF16 (RN), packed with ``a`` in the low half.""" + return Int32( + _asm_pure( + T.i32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _lds_v4_b32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return tuple(Int32(v) for v in _struct(res, 4, T.i32(), loc=loc, ip=ip)) + + +def smem_bytes(bn: int, stages: int, epilogue: str = EPI_NVFP4) -> int: + """Conservative host-side mirror of the GEMM's shared-memory footprint.""" + stage = BM * CTA_K + bn * CTA_K + 512 + _ceil_div(bn, BM) * 512 + staging = 4 * _RES_WARP_BYTES if epilogue == EPI_RESIDUAL else 0 + return 1024 + stages * stage + staging + 2048 + + +def max_stages(bn: int, epilogue: str = EPI_NVFP4) -> int: + s = 1 + while smem_bytes(bn, s + 1, epilogue) <= _SMEM_CAPACITY: + s += 1 + return s + + +class K3Fp8LatentGemmKernel: + """MXFP8 GEMM ``[M, K] x [N, K]^T`` over a static persistent tile schedule, fused epilogue. + + CTA ``c`` runs tiles ``c, c + grid, ...``; tile ``t`` is output rows ``128 * (t % m_tiles)`` + and columns ``bn * (t // m_tiles)``, so the CTAs sharing a weight tile run together. The + activation (E4M3, from :class:`K3Fp8Quant1x128Kernel`) and weight stages share one ring; + the accumulator is double-buffered in TMEM when two fit, so the epilogue of one tile + overlaps the mainloop of the next. + """ + + def __init__( + self, + n: int, + k: int, + bn: int, + stages: int, + epilogue: str, + use_pdl: bool, + ): + if bn not in SUPPORTED_BN: + raise ValueError(f"bn must be one of {SUPPORTED_BN}, got {bn}") + if n % BM or k % CTA_K: + raise ValueError(f"N={n} and K={k} must be multiples of {BM}") + if epilogue not in (EPI_NVFP4, EPI_RESIDUAL): + raise ValueError(f"unknown epilogue {epilogue!r}") + if smem_bytes(bn, stages, epilogue) > _SMEM_CAPACITY: + raise ValueError(f"{stages} stages of bn={bn} exceed shared memory") + self.n = n + self.k = k + self.bn = bn + self.stages = stages + self.epilogue_kind = epilogue + self.use_pdl = use_pdl + self.kt_total = k // CTA_K + # The last column tile may be partial (bn = 192); its missing weight rows are TMA + # zero fill and its missing columns are not stored. + self.n_tiles = _ceil_div(n, bn) + # Scale-factor atoms (128 rows) per stage of the weight operand. + self.nb = _ceil_div(bn, BM) + # A tile spans at most this many 128-row weight scale blocks. + self.wblocks = bn // BM + (1 if bn % BM else 0) + self.mma_tiler = (BM, bn, CTA_K) + self.mma_tiler_sfb = (BM, self.nb * BM, CTA_K) + self.a_stage_bytes = BM * CTA_K + self.b_stage_bytes = bn * CTA_K + self.sfb_stage_bytes = self.nb * 512 + sf_cols = _SFA_TMEM_COLS + self.nb * 16 + self.acc_bufs = 2 if 2 * bn + sf_cols <= 512 else 1 + cols = self.acc_bufs * bn + sf_cols + self.tmem_cols = 1 << (cols - 1).bit_length() + self.sfa_tmem_off = self.acc_bufs * bn + self.sfb_tmem_off = self.sfa_tmem_off + _SFA_TMEM_COLS + # Weight block exponents per lane: K tiles lane, lane + 32, ... + self.exp_regs = _ceil_div(self.kt_total, 32) + # Prefetched weight stages issued before waiting on the previous grid. + self.prefetch = min(stages, self.kt_total) + self.staging_bytes = 4 * _RES_WARP_BYTES if epilogue == EPI_RESIDUAL else 16 + # Epilogue warps (warps 4 .. 4 + epi_warps - 1); subclasses may use two per TMEM lane quarter. + self.epi_warps = 4 + + @cute.jit + def __call__( + self, + xq_ptr: cute.Pointer, + xs_ptr: cute.Pointer, + w_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + out0_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + num_ctas: Int32, + stream: cuda.CUstream, + ): + """``nvfp4``: out0 = codes, out1 = scales, aux0 = global scale (aux1 unused). + ``residual``: out0 = BF16 output, aux0 = shared-expert output, aux1 = residual.""" + n, k = self.n, self.k + mA = cute.make_tensor(xq_ptr, cute.make_layout((m, k, 1), stride=(k, 1, m * k))) + mB = cute.make_tensor(w_ptr, cute.make_layout((n, k, 1), stride=(k, 1, n * k))) + # DeepGEMM packed UE8M0 weight scale: int32 [N, ceil(K/512)], MN-major; byte + # (kb % 4) of word (n, kb // 4) holds the exponent of block (n // 128, kb). + mWS = cute.make_tensor( + ws_ptr, + cute.make_layout((n // BM, (4, _ceil_div(k // CTA_K, 4))), stride=(BM * 4, (1, 4 * n))), + ) + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float8E8M0FNU, + SF_VEC, + tcgen05.CtaGroup.ONE, + (BM, self.bn), + ) + a_layout = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.stages + ) + b_layout = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, cutlass.Float8E4M3FN, self.stages + ) + sfa_layout = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, self.mma_tiler, SF_VEC, self.stages + ) + # The weight scale operand covers whole 128-row atoms (bn rounded up to 128). + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float8E8M0FNU, + SF_VEC, + tcgen05.CtaGroup.ONE, + (BM, self.nb * BM), + ) + sfb_layout = blockscaled_utils.make_smem_layout_sfb( + tiled_mma_sfb, self.mma_tiler_sfb, SF_VEC, self.stages + ) + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) + ) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mA, + cute.slice_(a_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + mB, + cute.slice_(b_layout, (None, None, None, 0)), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + + @cute.struct + class SharedStorage: + full: cute.struct.MemRange[cutlass.Int64, self.stages] + empty: cute.struct.MemRange[cutlass.Int64, self.stages] + acc_full: cute.struct.MemRange[cutlass.Int64, self.acc_bufs] + acc_empty: cute.struct.MemRange[cutlass.Int64, self.acc_bufs] + tmem_ready: cutlass.Int64 + tmem_free: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E4M3FN, cute.cosize(a_layout.outer)], 1024 + ] + sB: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E4M3FN, cute.cosize(b_layout.outer)], 1024 + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfa_layout)], 128 + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float8E8M0FNU, cute.cosize(sfb_layout)], 128 + ] + sOut: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, self.staging_bytes], 128] + + self.shared_storage = SharedStorage + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + xs_ptr, + mWS, + out0_ptr, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + ).launch( + grid=(num_ctas, 1, 1), + block=(32 * (_EPI_WARP0 + self.epi_warps), 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB: cute.Tensor, + xs_ptr: cute.Pointer, + mWS: cute.Tensor, + out0_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + m_tiles = (m + BM - 1) // BM + num_tiles = m_tiles * self.n_tiles + + storage = utils.SmemAllocator().allocate(self.shared_storage) + full = storage.full.data_ptr() + empty = storage.empty.data_ptr() + + if warp_idx == _TMA_WARP: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(self.stages): + # Arrivals: the TMA producer (with the transaction bytes) and the SFB writer. + cute.arch.mbarrier_init(full + i, 2) + cute.arch.mbarrier_init(empty + i, 1) + for i in cutlass.range_constexpr(self.acc_bufs): + cute.arch.mbarrier_init(storage.acc_full.data_ptr() + i, 1) + cute.arch.mbarrier_init(storage.acc_empty.data_ptr() + i, self.epi_warps) + cute.arch.mbarrier_init(storage.tmem_ready, 1) + cute.arch.mbarrier_init(storage.tmem_free, self.epi_warps) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + if warp_idx == _TMA_WARP: + self.producer( + tma_atom_a, + mA, + tma_atom_b, + mB, + tiled_mma, + storage, + xs_ptr, + m_tiles, + a_layout, + b_layout, + ) + elif warp_idx == _SFB_WARP: + self.sfb_writer(storage, mWS, m_tiles) + elif warp_idx == _MMA_WARP: + self.mma_warp( + tiled_mma, + tiled_mma_sfb, + storage, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + num_tiles, + ) + elif warp_idx >= _EPI_WARP0: + self.epilogue( + tiled_mma, + storage, + out0_ptr, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + m_tiles, + tidx - 32 * _EPI_WARP0, + ) + + @cute.jit + def tile_coords(self, t: Int32, m_tiles: Int32): + """``(n_tile, m_tile)`` of linear tile ``t``: column-major, so the CTAs running together + share a weight tile. Every warp role maps tiles through this method.""" + n_tile = t // m_tiles + return n_tile, t - n_tile * m_tiles + + @cute.jit + def producer( + self, + tma_atom_a, + mA, + tma_atom_b, + mB, + tiled_mma, + storage, + xs_ptr, + m_tiles, + a_layout, + b_layout, + ): + """TMA warp: weight + activation tiles and the activation scale block per stage.""" + full = storage.full.data_ptr() + empty = storage.empty.data_ptr() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + num_tiles = m_tiles * self.n_tiles + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + gA = cute.local_tile(mA, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gB = cute.local_tile(mB, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + thr_mma = tiled_mma.get_slice(0) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(thr_mma.partition_A(gA), 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(thr_mma.partition_B(gB), 0, 3), + ) + sfa_base = storage.sSFA.data_ptr().toint() + stage_bytes = self.a_stage_bytes + self.b_stage_bytes + 512 + # Weights do not depend on the previous grid: fill the ring with the first tile's + # weight stages, then wait for the activation. + n_tile0, _ = self.tile_coords(bidx, m_tiles) + if bidx < num_tiles: + for kt in cutlass.range(self.prefetch, unroll=1): + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(full + kt, stage_bytes) + cute.copy( + tma_atom_b, + tBgB[(None, n_tile0, kt, 0)], + tBsB[(None, kt)], + tma_bar_ptr=full + kt, + ) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + it = Int32(0) + for t in cutlass.range(bidx, num_tiles, gdim, unroll=1): + n_tile, m_tile = self.tile_coords(t, m_tiles) + for kt in cutlass.range(self.kt_total, unroll=1): + stage = it % self.stages + if it >= self.stages: + cute.arch.mbarrier_wait(empty + stage, ((it // self.stages) & 1) ^ 1) + if it >= self.prefetch: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(full + stage, stage_bytes) + cute.copy( + tma_atom_b, + tBgB[(None, n_tile, kt, 0)], + tBsB[(None, stage)], + tma_bar_ptr=full + stage, + ) + cute.copy( + tma_atom_a, + tAgA[(None, m_tile, kt, 0)], + tAsA[(None, stage)], + tma_bar_ptr=full + stage, + ) + with cute.arch.elect_one(): + _bulk_load_g2s( + sfa_base + stage * 512, + xs_ptr.toint() + Int64((kt * m_tiles + m_tile) * 512), + 512, + (full + stage).toint(), + ) + it = it + 1 + + @cute.jit + def sfb_writer(self, storage, mWS: cute.Tensor, m_tiles: Int32): + """Write each stage's weight scale atoms from the 128x128 block exponents. + + Atom ``b`` covers tile rows ``128 b .. 128 b + 127``; its byte + ``(r % 32) * 16 + (r // 32) * 4 + k32`` belongs to tile row ``128 b + r``. A 32-row group + never straddles a weight block, so lane ``l`` writes one 4-byte word per group ``r // 32``. + Rows past ``N`` reuse the last block's exponent (their weight rows are zero). + """ + full = storage.full.data_ptr() + empty = storage.empty.data_ptr() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + num_tiles = m_tiles * self.n_tiles + lane = cute.arch.lane_idx() + nblk_w = self.n // BM + sfb_base = storage.sSFB.data_ptr().toint() + it = Int32(0) + for t in cutlass.range(bidx, num_tiles, gdim, unroll=1): + n_tile, _ = self.tile_coords(t, m_tiles) + row0 = n_tile * self.bn + blk0 = row0 // BM + exps = [] + for wb in cutlass.range_constexpr(self.wblocks): + blk = _min_s32(blk0 + wb, nblk_w - 1) + for j in cutlass.range_constexpr(self.exp_regs): + e = Int32(127) + if lane + 32 * j < self.kt_total: + e = Int32(mWS[blk, lane + 32 * j]) + exps.append(e) + # Weight block (relative to blk0) of each (atom, 32-row group); warp-uniform. + sel = [] + for b in cutlass.range_constexpr(self.nb): + for g in cutlass.range_constexpr(4): + sel.append((row0 + b * BM + g * 32) // BM - blk0) + for kt in cutlass.range(self.kt_total, unroll=1): + stage = it % self.stages + if it >= self.stages: + cute.arch.mbarrier_wait(empty + stage, ((it // self.stages) & 1) ^ 1) + ek = [] + for wb in cutlass.range_constexpr(self.wblocks): + e = cute.arch.shuffle_sync(exps[wb * self.exp_regs], kt & 31) + for j in cutlass.range_constexpr(1, self.exp_regs): + ej = cute.arch.shuffle_sync(exps[wb * self.exp_regs + j], kt & 31) + if (kt >> 5) == j: + e = ej + ek.append(e) + for b in cutlass.range_constexpr(self.nb): + words = [] + for g in cutlass.range_constexpr(4): + e = ek[0] + for wb in cutlass.range_constexpr(1, self.wblocks): + if sel[b * 4 + g] == wb: + e = ek[wb] + words.append(e * Int32(0x01010101)) + _sts_v4_b32( + sfb_base + stage * self.sfb_stage_bytes + b * 512 + lane * 16, *words + ) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(full + stage) + it = it + 1 + + @cute.jit + def mma_warp( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + storage, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + num_tiles: Int32, + ): + full = storage.full.data_ptr() + empty = storage.empty.data_ptr() + acc_full = storage.acc_full.data_ptr() + acc_empty = storage.acc_empty.data_ptr() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding_buf) + cute.arch.relinquish_tmem_alloc_permit() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(storage.tmem_ready) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + acc_layout = tiled_mma.make_fragment_C(acc_shape).layout + tCtSFA = cute.make_tensor( + cute.recast_ptr(acc_tmem_ptr + self.sfa_tmem_off, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, self.mma_tiler, SF_VEC, cute.slice_(sfa_layout, (None, None, None, 0)) + ), + ) + tCtSFB = cute.make_tensor( + cute.recast_ptr(acc_tmem_ptr + self.sfb_tmem_off, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.make_tmem_layout_sfb( + tiled_mma_sfb, + self.mma_tiler_sfb, + SF_VEC, + cute.slice_(sfb_layout, (None, None, None, 0)), + ), + ) + s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = _s2t_partition_sf( + storage.sSFA.get_tensor(sfa_layout), tCtSFA + ) + s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = _s2t_partition_sf( + storage.sSFB.get_tensor(sfb_layout), tCtSFB + ) + sA = storage.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sB = storage.sB.get_tensor(b_layout.outer, swizzle=b_layout.inner) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + num_kphases = cute.size(tCrA, mode=[2]) + + it = Int32(0) + ti = Int32(0) + for _t in cutlass.range(bidx, num_tiles, gdim, unroll=1): + buf = ti % self.acc_bufs + if ti >= self.acc_bufs: + cute.arch.mbarrier_wait(acc_empty + buf, ((ti // self.acc_bufs) & 1) ^ 1) + tCtAcc = cute.make_tensor(acc_tmem_ptr + buf * self.bn, acc_layout) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for _kt in cutlass.range(self.kt_total, unroll=1): + stage = it % self.stages + cute.arch.mbarrier_wait(full + stage, (it // self.stages) & 1) + cute.copy(s2t_sfa, tCsSFA_s2t[(None, None, None, None, stage)], tCtSFA_s2t) + cute.copy(s2t_sfb, tCsSFB_s2t[(None, None, None, None, stage)], tCtSFB_s2t) + for kp in cutlass.range_constexpr(num_kphases): + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator) + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[(None, None, kp, stage)], + tCrB[(None, None, kp, stage)], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + with cute.arch.elect_one(): + tcgen05.commit(empty + stage) + it = it + 1 + with cute.arch.elect_one(): + tcgen05.commit(acc_full + buf) + ti = ti + 1 + + cute.arch.mbarrier_wait(storage.tmem_free, 0) + cute.arch.dealloc_tmem(acc_tmem_ptr, self.tmem_cols) + + @cute.jit + def epilogue( + self, + tiled_mma: cute.TiledMma, + storage, + out0_ptr: cute.Pointer, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + m_tiles: Int32, + q: Int32, + ): + """Thread ``q`` owns output row ``128 * m_tile + q`` of every tile; 32-column TMEM blocks.""" + acc_full = storage.acc_full.data_ptr() + acc_empty = storage.acc_empty.data_ptr() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + num_tiles = m_tiles * self.n_tiles + gsf = cutlass.Float32(0.0) + rgsf = cutlass.Float32(0.0) + rcp6 = cutlass.Float32(0.0) + if cutlass.const_expr(self.epilogue_kind == EPI_NVFP4): + gsf = _ldg_f32(aux0_ptr.toint()) + rgsf = _rcp_approx_ftz(gsf) + rcp6 = _rcp_approx_ftz(cutlass.Float32(6.0)) + cute.arch.mbarrier_wait(storage.tmem_ready, 0) + acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + acc_layout = tiled_mma.make_fragment_C(acc_shape).layout + tAcc0 = cute.make_tensor(acc_tmem_ptr, acc_layout)[((None, None), 0, 0)] + cAcc = cute.make_identity_tensor((BM, self.bn)) + blk = 32 + tiled_t2r = tcgen05.make_tmem_copy( + cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(blk)), cutlass.Float32), + cute.local_tile(tAcc0, (BM, blk), (0, 0)), + ) + thr_t2r = tiled_t2r.get_slice(q) + nblk = self.bn // blk + tTR_cAcc0 = thr_t2r.partition_D(cute.local_tile(cAcc, (BM, blk), (0, 0))) + rBuf = [cute.make_rmem_tensor(tTR_cAcc0.shape, cutlass.Float32) for _ in range(2)] + + ti = Int32(0) + for t in cutlass.range(bidx, num_tiles, gdim, unroll=1): + n_tile, m_tile = self.tile_coords(t, m_tiles) + buf = ti % self.acc_bufs + tAcc = cute.make_tensor(acc_tmem_ptr + buf * self.bn, acc_layout)[((None, None), 0, 0)] + row = m_tile * BM + q + col0 = n_tile * self.bn + # Residual epilogue: the shared-expert and residual words of block j are loaded + # _RES_PREFETCH blocks ahead (the first ones while the mainloop still runs) in the + # coalesced layout. + pf = {} + if cutlass.const_expr(self.epilogue_kind == EPI_RESIDUAL): + for j in cutlass.range_constexpr(min(_RES_PREFETCH, nblk)): + pf[j] = self._load_residual_block( + aux0_ptr, aux1_ptr, m_tile, col0 + j * blk, m, q + ) + cute.arch.mbarrier_wait(acc_full + buf, (ti // self.acc_bufs) & 1) + cute.copy( + tiled_t2r, thr_t2r.partition_S(cute.local_tile(tAcc, (BM, blk), (0, 0))), rBuf[0] + ) + for j in cutlass.range_constexpr(nblk): + cute.arch.fence_view_async_tmem_load() + rAcc = rBuf[j % 2] + if cutlass.const_expr(j + 1 < nblk): + cute.copy( + tiled_t2r, + thr_t2r.partition_S(cute.local_tile(tAcc, (BM, blk), (0, j + 1))), + rBuf[(j + 1) % 2], + ) + if cutlass.const_expr(self.epilogue_kind == EPI_NVFP4): + if row < m and col0 + j * blk < self.n: + self._store_nvfp4( + rAcc, out0_ptr, out1_ptr, row, col0 + j * blk, gsf, rgsf, rcp6 + ) + else: + if cutlass.const_expr(j + _RES_PREFETCH < nblk): + pf[j + _RES_PREFETCH] = self._load_residual_block( + aux0_ptr, aux1_ptr, m_tile, col0 + (j + _RES_PREFETCH) * blk, m, q + ) + self._store_residual( + rAcc, storage, out0_ptr, pf[j], m_tile, col0 + j * blk, m, q + ) + # Every TMEM load of this buffer has completed (fence above); release it. + cute.arch.fence_view_async_tmem_load() + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(acc_empty + buf) + ti = ti + 1 + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(storage.tmem_free) + + @cute.jit + def _store_nvfp4(self, rAcc, q_ptr, sf_ptr, row, col, gsf, rgsf, rcp6): + """32 accumulators of one row -> BF16 -> 16 bytes of E2M1 codes + two E4M3 scales.""" + vals = [] + for i in cutlass.range_constexpr(16): + a, b = _bf16_round_x2(rAcc[2 * i], rAcc[2 * i + 1]) + vals.extend((a, b)) + amax0 = _abs_max_f32x16(vals[:16]) + amax1 = _abs_max_f32x16(vals[16:]) + sf2, os0, os1 = _nvfp4_scales_x2(amax0, amax1, gsf, rgsf, rcp6) + w0 = _e2m1x8(vals[0:8], os0) + w1 = _e2m1x8(vals[8:16], os0) + w2 = _e2m1x8(vals[16:24], os1) + w3 = _e2m1x8(vals[24:32], os1) + _stg_v4_b32(q_ptr.toint() + Int64(row) * (self.n // 2) + Int64(col // 2), w0, w1, w2, w3) + _stg_u16(sf_ptr.toint() + Int64(row) * (self.n // FP4_VEC) + Int64(col // FP4_VEC), sf2) + + @cute.jit + def _load_residual_block(self, shared_ptr, residual_ptr, m_tile, col, m, q): + """Shared / residual words of one 32-column block in the coalesced store layout. + + Warp ``q // 32`` handles rows ``32 (q // 32) + 8 i + lane // 4`` (``i < 4``), lane + ``l`` the 16 bytes at column ``col + 8 (l % 4)``; rows / columns that are not stored + read a clamped in-bounds address. + """ + lane = q % 32 + r0 = m_tile * BM + (q - lane) + lane // 4 + c = _min_s32(col + (lane % 4) * 8, self.n - 8) + words = [] + for i in cutlass.range_constexpr(4): + r = _min_s32(r0 + 8 * i, m - 1) + off = (Int64(r) * self.n + Int64(c)) * 2 + words.append( + (_ldg_v4_b32(shared_ptr.toint() + off), _ldg_v4_b32(residual_ptr.toint() + off)) + ) + return words + + @cute.jit + def _store_residual(self, rAcc, storage, out_ptr, block, m_tile, col, m, q): + """One 32-column block: ``residual + bf16(bf16(acc) + shared)`` with coalesced stores. + + Thread ``q`` rounds its row to BF16 into the warp's SMEM staging rows; after a warp + sync, every 4 lanes combine and store 64 contiguous bytes of one row. + """ + lane = q % 32 + base = storage.sOut.data_ptr().toint() + (q // 32) * _RES_WARP_BYTES + words = [] + for i in cutlass.range_constexpr(16): + words.append(_f32x2_to_bf16x2(rAcc[2 * i], rAcc[2 * i + 1])) + for v in cutlass.range_constexpr(4): + _sts_v4_b32(base + lane * _RES_ROW_BYTES + v * 16, *words[4 * v : 4 * v + 4]) + cute.arch.sync_warp() + r0 = m_tile * BM + (q - lane) + lane // 4 + c = col + (lane % 4) * 8 + for i in cutlass.range_constexpr(4): + t = _lds_v4_b32(base + (lane // 4 + 8 * i) * _RES_ROW_BYTES + (lane % 4) * 16) + sh, rs = block[i] + o = [_residual_from_bf16x2(t[v], sh[v], rs[v]) for v in range(4)] + r = r0 + 8 * i + if r < m and c < self.n: + _stg_v4_b32(out_ptr.toint() + (Int64(r) * self.n + Int64(c)) * 2, *o) + # The staging rows are rewritten by the next block. + cute.arch.sync_warp() + + +def _s2t_partition_sf(sSF: cute.Tensor, tSF: cute.Tensor): + """S2T (SMEM -> TMEM) copy of a scale-factor operand, partitioned for one thread.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), cutlass.Float8E8M0FNU + ) + tiled_copy = tcgen05.make_s2t_copy(copy_atom, tCtSF_compact) + thr_copy = tiled_copy.get_slice(0) + tCsSF_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy, thr_copy.partition_S(tCsSF_compact)) + tCtSF_s2t = thr_copy.partition_D(tCtSF_compact) + return tiled_copy, tCsSF_s2t, tCtSF_s2t + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +@functools.cache +def _num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def bn_for(m: int, n: int) -> int: + """Column tile minimizing the per-CTA work of the persistent schedule (ties: wider tile).""" + m_tiles = _ceil_div(m, BM) + best, best_cost = SUPPORTED_BN[-1], None + for bn in SUPPORTED_BN: + cost = _ceil_div(m_tiles * _ceil_div(n, bn), _num_sms()) * bn + if best_cost is None or cost <= best_cost: + best, best_cost = bn, cost + return best + + +@functools.cache +def _compiled_quant(k: int, use_pdl: bool, sf_m_major: bool = False): + kernel = K3Fp8Quant1x128Kernel(k, use_pdl, sf_m_major=sf_m_major) + dummy = 1 << 20 + x_ptr = make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16) + xq_ptr = make_ptr(cutlass.Float8E4M3FN, dummy, cute.AddressSpace.gmem, assumed_align=16) + xs_ptr = make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile(kernel, x_ptr, xq_ptr, xs_ptr, Int32(1), stream) + + +@functools.cache +def _compiled_gemm(n: int, k: int, bn: int, stages: int, epilogue: str, use_pdl: bool): + kernel = K3Fp8LatentGemmKernel(n, k, bn, stages, epilogue, use_pdl) + dummy = 1 << 20 + ptrs = [ + make_ptr(dtype, dummy, cute.AddressSpace.gmem, assumed_align=16) + for dtype in (cutlass.Float8E4M3FN, cutlass.Uint8, cutlass.Float8E4M3FN) + + (cutlass.Uint8,) * 5 + ] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile(kernel, *ptrs, Int32(1), Int32(1), stream) + + +def _gmem_ptr(dtype, tensor: torch.Tensor): + return make_ptr(dtype, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def quantize_act_1x128( + x: torch.Tensor, + xq: torch.Tensor, + xs: torch.Tensor, + *, + use_pdl: bool = True, + sf_m_major: bool = False, +) -> None: + """Launch :class:`K3Fp8Quant1x128Kernel` on the current stream.""" + m, k = x.shape + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + _compiled_quant(k, use_pdl, sf_m_major)( + _gmem_ptr(cutlass.BFloat16, x), + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + Int32(m), + stream, + ) + + +#: Per compute capability: (N, K) -> largest M for which the fused ops run the split-K GEMM of +#: ``k3_fp8w_lowm_gemm`` with the fused epilogue instead of the persistent kernel (measured on +#: Rubin, CUDA graph, cold weights: faster than the persistent kernel and than the unfused +#: CuTe MXFP8 / low-M GEMM + ``fp4_quantize`` / adds at every M up to the limit). +_LOWM_ROUTE_MAX_M = { + (10, 7): {(3584, 7168): 256}, +} +#: Per compute capability: (N, K) -> ``(min M, max M, kernel, split_k)`` ranges where the route +#: uses ``k3_fp8w_lowm_gemm.split_tactic`` (kernel "low" or "mid") instead of that module's +#: default tactic. latent_down runs on a side stream next to the routing chain: its standalone +#: optimum (up to 224 CTAs) delays the router GEMM and routing kernels, so it uses narrower +#: grids (84-112 CTAs) that measured faster per MoE layer with the routing beside them. +_LOWM_ROUTE_TACTICS = { + (10, 7): { + (3584, 7168): ((1, 16, "low", 3), (17, 96, "mid", 4)), + }, +} + + +@functools.cache +def _device_capability(device_index: int) -> tuple: + return torch.cuda.get_device_capability(device_index) + + +def lowm_route_max_m(n: int, k: int) -> int: + """Largest M for which the ops run the ``k3_fp8w_lowm_gemm`` split-K GEMM (0: never).""" + table = _LOWM_ROUTE_MAX_M.get(_device_capability(torch.cuda.current_device()), {}) + return table.get((n, k), 0) + + +def route_tactic(m: int, n: int, k: int): + """The ``k3_fp8w_lowm_gemm`` tactic of the route for M tokens (``None``: its default).""" + from . import k3_fp8w_lowm_gemm + + table = _LOWM_ROUTE_TACTICS.get(_device_capability(torch.cuda.current_device()), {}) + for lo, hi, kernel, split_k in table.get((n, k), ()): + if lo <= m <= hi: + return k3_fp8w_lowm_gemm.split_tactic(m, n, k, split_k, mid=kernel == "mid") + return None + + +def _run(x, w, w_scale, epilogue, out0, out1, aux0, aux1, bn, stages, use_pdl) -> None: + m, k = x.shape + n = w.shape[0] + if bn is None and stages is None and m <= lowm_route_max_m(n, k): + from . import k3_fp8w_lowm_gemm + + k3_fp8w_lowm_gemm.run_k3_fp8w_lowm_gemm_fused( + x, + w, + w_scale, + epilogue, + out0, + out1, + aux0, + aux1, + tactic=route_tactic(m, n, k), + use_pdl=use_pdl, + ) + return + if bn is None: + bn = bn_for(m, n) + if stages is None: + stages = max_stages(bn, epilogue) + m_tiles = _ceil_div(m, BM) + num_ctas = min(m_tiles * _ceil_div(n, bn), _num_sms()) + xq = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device) + xs = torch.empty((k // CTA_K) * m_tiles * 512, dtype=torch.uint8, device=x.device) + quantize_act_1x128(x, xq, xs, use_pdl=use_pdl) + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + _compiled_gemm(n, k, bn, stages, epilogue, use_pdl)( + _gmem_ptr(cutlass.Float8E4M3FN, xq), + _gmem_ptr(cutlass.Uint8, xs), + _gmem_ptr(cutlass.Float8E4M3FN, w), + *(_gmem_ptr(cutlass.Uint8, t) for t in (w_scale, out0, out1, aux0, aux1)), + Int32(m), + Int32(num_ctas), + stream, + ) + + +def run_k3_latent_down_fp4( + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + global_sf: torch.Tensor, + q_out: torch.Tensor, + sf_out: torch.Tensor, + *, + bn: int | None = None, + stages: int | None = None, + use_pdl: bool = True, +) -> None: + """Quantize ``x`` and launch the NVFP4-epilogue GEMM into ``q_out`` / ``sf_out``. + + ``bn`` / ``stages`` select a tile of the persistent kernel; with neither, M up to + :func:`lowm_route_max_m` runs the ``k3_fp8w_lowm_gemm`` split-K GEMM instead.""" + _run(x, w, w_scale, EPI_NVFP4, q_out, sf_out, global_sf, global_sf, bn, stages, use_pdl) + + +def _check_operands(x, w, w_scale) -> None: + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous(): + raise ValueError("x must be a contiguous 2D bfloat16 tensor") + if w.dtype != torch.float8_e4m3fn or w.dim() != 2 or not w.is_contiguous(): + raise ValueError("w must be a contiguous 2D float8_e4m3fn tensor") + n, k = w.shape + if x.shape[1] != k: + raise ValueError(f"x has K={x.shape[1]}, weight has K={k}") + if n % BM or k % CTA_K: + raise ValueError(f"N={n} and K={k} must be multiples of {BM}") + words = _ceil_div(k // CTA_K, 4) + if ( + w_scale.dtype != torch.int32 + or tuple(w_scale.shape) != (n, words) + or w_scale.stride(0) != 1 + or (words > 1 and w_scale.stride(1) != n) + ): + raise ValueError( + f"w_scale must be DeepGEMM-packed UE8M0 int32 [{n}, {words}] with strides (1, {n})" + ) + if x.data_ptr() % 16 or w.data_ptr() % 16: + raise ValueError("x and w must be 16-byte aligned") + + +def _check_global_sf(global_sf) -> None: + if ( + not isinstance(global_sf, torch.Tensor) + or global_sf.dtype != torch.float32 + or global_sf.numel() != 1 + or not global_sf.is_cuda + ): + raise ValueError("global_sf must be a one-element float32 CUDA tensor") + + +@functools.cache +def _is_supported_device(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3), (10, 7)) + + +def _can_run(x, w, checks) -> bool: + if x.dim() != 2 or w.dim() != 2 or not x.is_cuda: + return False + if not 1 <= x.shape[0] <= MAX_M: + return False + if not _is_supported_device(x.device.index): + return False + try: + for check in checks: + check() + except ValueError: + return False + return True + + +def can_run(x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, global_sf) -> bool: + """Whether :func:`k3_latent_down_fp4` accepts these inputs (non-raising check).""" + return _can_run( + x, w, (lambda: _check_operands(x, w, w_scale), lambda: _check_global_sf(global_sf)) + ) + + +def k3_latent_down_fp4( + x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, global_sf: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """NVFP4 ``(codes uint8 [M, N/2], scales uint8 [M, N/16])`` of ``bf16(x @ dequant(w)^T)``. + + ``w`` is the UE8M0-resmoothed FP8 weight and ``w_scale`` the DeepGEMM-packed UE8M0 scale + (as held by ``fp8_swap_ab_gemm`` call sites); ``global_sf`` is the NVFP4 global scale of the + consumer (``fp4_quantize``'s ``globalScale``). Scales use the linear layout. The activation + is quantized with 1x128 UE8M0 scales, or with 1x32 ones where the op runs the split-K GEMM + of ``k3_fp8w_lowm_gemm`` (``M <= lowm_route_max_m(N, K)``). + """ + _check_operands(x, w, w_scale) + _check_global_sf(global_sf) + m = x.shape[0] + n = w.shape[0] + q_out = torch.empty((m, n // 2), dtype=torch.uint8, device=x.device) + sf_out = torch.empty((m, n // FP4_VEC), dtype=torch.uint8, device=x.device) + run_k3_latent_down_fp4(x, w, w_scale, global_sf, q_out, sf_out) + return q_out, sf_out + + +__all__ = [ + "K3Fp8LatentGemmKernel", + "K3Fp8Quant1x128Kernel", + "MAX_M", + "bn_for", + "can_run", + "k3_latent_down_fp4", + "lowm_route_max_m", + "quantize_act_1x128", + "route_tactic", + "run_k3_latent_down_fp4", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_launch_attrs.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_launch_attrs.py new file mode 100644 index 000000000000..2da7d32c52c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_launch_attrs.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Driver function attributes of compiled CuTe DSL kernels (SM100/SM103/SM107). + +An SM splits its unified L1 / shared memory when a CTA lands on it while it is idle, and keeps +that split while any CTA runs on it. The split follows the preferred shared-memory carveout of the +kernel (``CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT``, percent of the maximum shared +memory; ``-1``, no preference, gives a kernel with little shared memory a small split). A CTA of +another kernel is placed on a busy SM if its kernel has no preference or prefers a split at least +as large as the SM's current one (measured on SM107). So a small kernel that runs first and should +leave room for a large-shared-memory kernel launched beside it (a programmatic dependent, or a +kernel on another stream) needs carveout 100; with ``-1`` the large kernel's CTAs wait for it to +exit. + +The CuTe DSL passes a launch's ``preferred_smem_carveout`` to the driver. A launch with +``min_blocks_per_mp > 1`` and no carveout gets one derived by the DSL from the kernel's own shared +memory, just enough for that many of its CTAs. + +The helpers here read and write these attributes on the kernels of a compiled CuTe DSL function +(``cute.compile`` result) through the driver API, without recompiling: + +* :func:`set_shared_carveout` writes the preferred carveout (and optionally the dynamic + shared-memory limit) on every kernel of a compiled function; +* :func:`apply_shared_carveout` does it only under ``TRTLLM_K3_CUTE_CARVEOUT_FIX=1`` and returns + the compiled function, so a compile site can wrap its ``cute.compile`` call; +* :func:`kernel_attributes` reports the attributes of every kernel (tests and probes). + +The attributes stick until the function's CUDA library is unloaded; set them before graph capture. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +import cuda.bindings.driver as cuda +import torch + +from ..features import active as _features + +#: Preferred carveout of the maximum shared-memory split (percent). +MAX_SHARED_CARVEOUT = 100 + +_ATTR = cuda.CUfunction_attribute +_CARVEOUT_ATTR = _ATTR.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT +_MAX_DYNAMIC_SMEM_ATTR = _ATTR.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES +_REPORTED_ATTRS = { + "carveout": _CARVEOUT_ATTR, + "max_dynamic_smem": _MAX_DYNAMIC_SMEM_ATTR, + "static_smem": _ATTR.CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, + "registers": _ATTR.CU_FUNC_ATTRIBUTE_NUM_REGS, + "max_threads": _ATTR.CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, +} + + +def carveout_fix_enabled() -> bool: + """``features.cute_carveout_fix``: K3 kernels that run ahead of large-shared-memory kernels (the + attn_res kernels, L2 prefetch tables built without a carveout) get carveout 100 through the + driver API (:func:`apply_shared_carveout`).""" + return _features().cute_carveout_fix + + +def _checked(result: tuple, what: str) -> Any: + status, *values = result + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"{what} failed: {status}") + return values[0] if len(values) == 1 else tuple(values) + + +def _device_index(device: Optional[int | torch.device]) -> int: + if device is None: + return torch.cuda.current_device() + if isinstance(device, torch.device): + return device.index if device.index is not None else torch.cuda.current_device() + return int(device) + + +def _kernels(compiled: Any, device_index: int) -> List[Tuple[str, Any]]: + """``(name, CUkernel)`` of every kernel of ``compiled`` (loaded on ``device_index`` first).""" + with torch.cuda.device(device_index): + compiled.to(device_index) + jit_module = getattr(compiled, "jit_module", None) + libraries = list(getattr(jit_module, "cuda_library", None) or []) + if not libraries: + raise RuntimeError(f"{type(compiled).__name__} has no loaded CUDA library") + kernels = [] + for library in libraries: + handle = cuda.CUlibrary(int(library)) + count = _checked(cuda.cuLibraryGetKernelCount(handle), "cuLibraryGetKernelCount") + for kernel in _checked( + cuda.cuLibraryEnumerateKernels(count, handle), "cuLibraryEnumerateKernels" + ): + name = _checked(cuda.cuKernelGetName(kernel), "cuKernelGetName") + kernels.append((name.decode() if isinstance(name, bytes) else str(name), kernel)) + return kernels + + +def set_shared_carveout( + compiled: Any, + carveout: int, + device: Optional[int | torch.device] = None, + max_dynamic_smem: Optional[int] = None, +) -> None: + """Write the preferred shared-memory carveout on every kernel of ``compiled``. + + Args: + compiled: a ``cute.compile`` result (CUDA dialect executor). + carveout: preferred carveout in percent of the maximum shared memory (0-100), or -1 for + the driver's choice. + device: device whose kernels are updated (default: the current device). + max_dynamic_smem: if given, also ``CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES``. + """ + if not -1 <= carveout <= 100: + raise ValueError(f"carveout={carveout} must be a percentage or -1") + index = _device_index(device) + cu_device = _checked(cuda.cuDeviceGet(index), "cuDeviceGet") + for _, kernel in _kernels(compiled, index): + if max_dynamic_smem is not None: + _checked( + cuda.cuKernelSetAttribute( + _MAX_DYNAMIC_SMEM_ATTR, max_dynamic_smem, kernel, cu_device + ), + "cuKernelSetAttribute(MAX_DYNAMIC_SHARED_SIZE_BYTES)", + ) + _checked( + cuda.cuKernelSetAttribute(_CARVEOUT_ATTR, carveout, kernel, cu_device), + "cuKernelSetAttribute(PREFERRED_SHARED_MEMORY_CARVEOUT)", + ) + + +def apply_shared_carveout( + compiled: Any, carveout: Optional[int], device: Optional[int | torch.device] = None +) -> Any: + """``compiled``, with ``carveout`` written by :func:`set_shared_carveout` when + ``TRTLLM_K3_CUTE_CARVEOUT_FIX=1`` and ``carveout`` is not ``None``.""" + if carveout is not None and carveout_fix_enabled(): + set_shared_carveout(compiled, carveout, device) + return compiled + + +def kernel_attributes( + compiled: Any, device: Optional[int | torch.device] = None +) -> List[Dict[str, Any]]: + """Name and driver attributes (:data:`_REPORTED_ATTRS`) of every kernel of ``compiled`` as the + current context sees them.""" + index = _device_index(device) + out = [] + with torch.cuda.device(index): + for name, kernel in _kernels(compiled, index): + function = _checked(cuda.cuKernelGetFunction(kernel), "cuKernelGetFunction") + entry: Dict[str, Any] = {"name": name} + for label, attribute in _REPORTED_ATTRS.items(): + entry[label] = _checked( + cuda.cuFuncGetAttribute(attribute, function), f"cuFuncGetAttribute({label})" + ) + out.append(entry) + return out + + +__all__ = [ + "MAX_SHARED_CARVEOUT", + "apply_shared_carveout", + "carveout_fix_enabled", + "kernel_attributes", + "set_shared_carveout", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_mla_q_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_mla_q_fused.py new file mode 100644 index 000000000000..5af39c2b52d4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_mla_q_fused.py @@ -0,0 +1,1763 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 FP4 MLA decode query in one kernel: q_a RMSNorm, q_b projection, W_UK absorption, +RoPE and the FP4 query quantization of the packed FP4 MLA decode (SM107). + +For ``T <= 16`` decode tokens and the 96 MLA heads of a Helix rank, :class:`K3MlaQFusedKernel` +turns the q_a latent ``x [T, 1536]`` (BF16, raw or already normalized) into the prequantized +query the packed FP4 MLA decode reads (``_fp4_mla_decode.run_decode``): per (token, +head) row 512 absorbed latent values and 64 RoPE values with their residual, E2M1 codes with one +E4M3 scale per 16 values. It replaces, with the same arithmetic up to the summation order of the +q_b GEMM (the MMA hardware's internal sum of a 32-value group may differ): + +0. (``fuse_norm``) the q_a RMSNorm: FlashInfer's ``RMSNormKernel`` operations in its own + per-lane layout (:meth:`K3MlaQFusedKernel.row_rstd`), bitwise its output. +1. ``q = q_b_proj(x)``: the low-M FP8 GEMM (``k3_fp8w_lowm_gemm``) with its in-kernel MXFP8 + activation quantization (E4M3 codes, one UE8M0 exponent per 32 values, bitwise the same + helpers) and the FP8 weight's 128x128 UE8M0 block scales; the product of a 32-value group is + scaled by ``2^(clamp(e_act + e_w - 127, 0, 254) - 127)`` as the GEMM's block-scaled MMA does, + accumulated in FP32 in K order per K half, the halves added (the GEMM's split-K partials), + rounded to BF16. +2. ``q_nope @ W_UK[h]^T`` with the BF16 absorption weight: ``mma.sync`` K16 steps in K order, + FP32 accumulation, rounded to BF16 (bitwise the BF16 BMM in the tests). +3. The fused FP4 generation cache update's query half + (``_fp4_mla_generation_fused_qk_rope_cache_update_kernel``): GPT-J RoPE of ``q_pe`` from the + rotary table at the token's position, rounded to BF16, then the global-scaled FP4 + quantization of the 512 absorbed values and the main + residual FP4 quantization of the 64 + RoPE values, in the packed per-request row layout with swizzled scales. This part repeats + the Triton kernel's operations (``div.full`` / ``rcp.approx`` / the contracted residual), so + given the same BF16 values it writes the same bytes. + +Work split: a cluster of two CTAs per head (192 CTAs, 12 warps each). CTA ``r`` of head ``h`` +owns q_b rows ``h * 192 + 96 r .. + 95`` (CTA 0: q_nope dims 0..95; CTA 1: q_nope dims 96..127 +and the 64 RoPE dims) and absorbed columns ``256 r .. 256 r + 255``: + +* Before ``griddepcontrol.wait`` (nothing read there depends on the previous grid): the CTA's + q_b rows (144 KB) are issued as six TMA boxes of 16 whole rows (contiguous in global memory) + that land as ``[k block][row][128 bytes]`` with the 128-byte swizzle (conflict-free + ``ldmatrix``). +* After the wait warp ``w`` loads rows ``w, w + 12, ..`` of ``x`` (32 lanes per row), normalizes + them (``fuse_norm``) and quantizes them to MXFP8 in shared memory. Warp ``w`` then runs K half + ``w // 6`` of the 16-row q_b tile ``w % 6`` with ``mma.sync`` m16n8k32 E4M3 (block scales + applied per 32-value group with one FMA); the second half's sums go to the first half's warp, + which adds them and writes q_nope (BF16) into both CTAs' shared memory (DSMEM) and q_pe + locally. CTA 1's last warp then prefetches the positions and rotary values; the CTA's W_UK rows + are prefetched into L2 after the activation phase, and once tiles 0..2 are done warp 0 loads + W_UK (64 KB, one box) into their space. After a cluster barrier warps 0..7 run two 16-column + W_UK tiles each with ``mma.sync`` m16n8k16 BF16. +* The absorbed BF16 values are staged per token and quantized 16 at a time; CTA 1 also rotates + and quantizes q_pe. The FP4 rows and scales go straight to the decode's Q buffers. + +:class:`K3MlaLatentRowKernel` is the KV half's companion: the kv_a RMSNorm and the ``k_pe`` copy +into the latent row in one launch. +""" + +from __future__ import annotations + +import functools +import operator +from typing import Optional + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.arch.nvvm_wrappers import inline_ptx +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.experimental import cuda as cuda_tma + +from .helix_a2a_fused import _fv, _lv, _pack_bf16x2, _unpack +from .k3_fp8w_lowm_gemm import ( + _asm, + _bf16x2_inv_scale, + _griddepcontrol_launch_dependents, + _griddepcontrol_wait, + _iv, + _lds_v4_b32, + _lds_v4_f32, + _quant_bf16x8_to_e4m3, + _sts_v2_b32, + _sts_v4_b32, + _sts_v4_f32, + _ue8m0_of_amax, +) +from .k3_fp8w_lowm_mxfp8 import _amax_bf16x8_words, _amax_pair +from .k3_helix_tail import ( + _globaltimer, + _mbar_expect_tx, + _mbar_init, + _mbar_try_wait, + _mma_bf16, + _shfl_bfly_b32, + _st_u64, +) + +try: + from flashinfer.norm.utils import warp_reduce as _flashinfer_warp_reduce + + NORM_FUSION_AVAILABLE = True +except ImportError: # pragma: no cover - FlashInfer without the CuTe DSL norm kernels + _flashinfer_warp_reduce = None + NORM_FUSION_AVAILABLE = False + +#: MLA geometry the kernel is built for (Kimi K3). +NUM_HEADS = 96 +NOPE_DIM = 128 +ROPE_DIM = 64 +QK_HEAD_DIM = NOPE_DIM + ROPE_DIM +KV_LORA_RANK = 512 +Q_LORA_RANK = 1536 +#: Layout of the FP4 query rows the packed FP4 MLA decode reads (``fp4_mla``): 512 absorbed +#: values, then the 64 RoPE values and their 64 residuals, 16 values per E4M3 scale. +FP4_GROUP = 16 +Q_ROW_BYTES = (KV_LORA_RANK + 2 * ROPE_DIM) // 2 +Q_SF_COLS = (KV_LORA_RANK + 2 * ROPE_DIM) // FP4_GROUP +#: Largest token count of one launch (the activation rows and their scales must fit shared +#: memory beside the resident q_b weight). +MAX_TOKENS = 16 + +_THREADS = 384 +_WARPS = _THREADS // 32 +#: Warps of the W_UK absorption (two 16-column tiles each). +_ABSORB_WARPS = 8 +#: q_b weight rows per CTA (half a head) and absorbed columns per CTA. +_CTA_ROWS = QK_HEAD_DIM // 2 +_CTA_COLS = KV_LORA_RANK // 2 +#: Weights arrive by TMA as 3D boxes whose shared image is ``[k block][row][128 bytes]`` with the +#: 128-byte swizzle (16-byte chunk ``c`` of row ``r`` lands at chunk ``c ^ (r % 8)``), so the +#: eight rows an ``ldmatrix`` 8x8 matrix reads hit distinct banks, while each box covers whole +#: weight rows (a contiguous range of global memory). One box per 16-row q_b tile (24 KB, all of +#: K); W_UK of the CTA's 256 absorbed columns is one box of two 64-wide halves (64 KB). +_BOX_BYTES = 128 +_KB = Q_LORA_RANK // _BOX_BYTES +_TILE_ROWS = 16 +_NTILE = _CTA_ROWS // _TILE_ROWS +_TILE_BYTES = _TILE_ROWS * Q_LORA_RANK +_UK_HALF = _CTA_COLS * _BOX_BYTES +_UK_BYTES = 2 * _UK_HALF +#: Padded shared row strides of the buffers the kernel writes itself (16 bytes past a multiple +#: of 128 keep ``ldmatrix`` rows in distinct banks). +_A_PITCH = Q_LORA_RANK + 16 +_QN_PITCH = NOPE_DIM * 2 + 16 +_QPE_PITCH = ROPE_DIM * 2 +_OB_PITCH = _CTA_COLS * 2 + 16 +_KGROUPS = Q_LORA_RANK // 32 +#: q_b is summed as this many K slices added at the end, the association of the low-M GEMM's +#: split-K partials at decode token counts (its SM107 tactic for this shape uses split 2 at +#: M <= 16). +_QB_SPLIT_K = 2 +_KBLOCKS = Q_LORA_RANK // 128 +#: W_UK reuses the shared memory of the first q_b chunks once every MMA warp is done with them; +#: the absorbed output is staged where the last chunks were. +_UK_TILES = -(-_UK_BYTES // _TILE_BYTES) +_OB_OFFSET = _UK_TILES * _TILE_BYTES +_SMEM_LIMIT = 227 * 1024 +#: Bytes of fixed shared buffers (barriers, block exponents) ahead of the data regions. +_FIXED_BYTES = 256 + +#: ``2^-9``, the floor of the FP4 query scales (``_fp4_mla_quantize_q_groups``). +_Q_SCALE_FLOOR = 0x3B000000 + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def smem_bytes(t_pad: int) -> int: + """Shared memory of a launch over ``t_pad`` tokens.""" + return ( + _FIXED_BYTES + + 2 * _KGROUPS * t_pad * 4 + + t_pad * _A_PITCH + + _NTILE * _TILE_BYTES + + t_pad * _QN_PITCH + + t_pad * _QPE_PITCH + + Q_LORA_RANK * 2 + + t_pad * (4 * 64 + 4) + + _NTILE * 32 * t_pad // 2 * 4 + + 2 * 1024 + + 8 * 128 # region alignment + ) + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +@dsl_user_op +def _ldsm_x4(addr, *, loc=None, ip=None): + """``ldmatrix`` of four 8x8 b16 matrices (lane ``l`` addresses a row of matrix ``l // 8``).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _ldsm_x2(addr, *, loc=None, ip=None): + """``ldmatrix`` of two 8x8 b16 matrices (lanes 0..15 address the rows).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_iv(addr, loc=loc, ip=ip)], + "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {$0, $1}, [$2];", + "=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _mma_e4m3(a0, a1, a2, a3, b0, b1, *, loc=None, ip=None): + """``d[16x8] = a[16x32] @ b[32x8]`` (E4M3 inputs, FP32 result, zero accumulator).""" + zero = Float32(0.0) + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(x, loc=loc, ip=ip) for x in (a0, a1, a2, a3, b0, b1)] + + [_fv(zero, loc=loc, ip=ip) for _ in range(4)], + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _block_scale(e_act, w_exp, *, loc=None, ip=None): + """FP32 ``2^(clamp(e_act + w_exp - 127, 0, 254) - 127)``: the block-scaled MMA's scale of a + 32-value group with the weight exponent folded into the activation's, as the low-M GEMM + stores it (0 means ``2^-127``).""" + return Float32( + _asm( + T.f32(), + [_iv(e_act, loc=loc, ip=ip), _iv(w_exp, loc=loc, ip=ip)], + "{\n.reg .b32 e, b;\n.reg .pred p;\n" + "add.s32 e, $1, $2;\nsub.s32 e, e, 127;\nmax.s32 e, e, 0;\nmin.s32 e, e, 254;\n" + "shl.b32 b, e, 23;\nsetp.eq.s32 p, e, 0;\n@p mov.b32 b, 0x00400000;\n" + "mov.b32 $0, b;\n}", + "=f,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _lds_v2_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v2.f32 {$0, $1}, [$2];", + "=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(x, loc=loc, ip=ip) for x in (a, b, c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _lds_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.b32 $0, [$1];", "=r,r", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _sts_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_u8(addr, *, loc=None, ip=None): + return Int32( + _asm(T.i32(), [_iv(addr, loc=loc, ip=ip)], "ld.shared.u8 $0, [$1];", "=r,r", loc=loc, ip=ip) + ) + + +@dsl_user_op +def _sts_u8(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.u8 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_bf16(addr, value, *, loc=None, ip=None): + """Store ``value`` rounded to BF16 (nearest even) at shared ``addr``.""" + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\nst.shared.b16 [$0], h;\n}", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_bf16_both(addr, peer, value, *, loc=None, ip=None): + """Store ``value`` rounded to BF16 at shared ``addr`` of this CTA and of cluster rank + ``peer`` (same offset).""" + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(peer, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "{\n.reg .b16 h;\n.reg .b32 ra;\ncvt.rn.bf16.f32 h, $2;\nst.shared.b16 [$0], h;\n" + "mapa.shared::cluster.u32 ra, $0, $1;\nst.shared::cluster.b16 [ra], h;\n}", + "r,r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_bf16(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\nst.global.b16 [$0], h;\n}", + "l,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ldg_u8(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.u8 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ldg_s32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), [_lv(addr, loc=loc, ip=ip)], "ld.global.s32 $0, [$1];", "=r,l", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _ldg_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), [_lv(addr, loc=loc, ip=ip)], "ld.global.f32 $0, [$1];", "=f,l", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _ldg_v4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _ldg_v4_pred(addr, pred, *, loc=None, ip=None): + """16 bytes at ``addr`` if ``pred > 0``, else zeros (no load).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip), _iv(pred, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.gt.s32 p, $5, 0;\nmov.b32 $0, 0;\nmov.b32 $1, 0;\n" + "mov.b32 $2, 0;\nmov.b32 $3, 0;\n@p ld.global.v4.b32 {$0, $1, $2, $3}, [$4];\n}", + "=r,=r,=r,=r,l,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _stg_v2(addr, v0, v1, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(v0, loc=loc, ip=ip), _iv(v1, loc=loc, ip=ip)], + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_v4(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u8(addr, value, *, loc=None, ip=None): + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u8 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@cute.jit +def _tma_load_3d(dst: Int32, desc, y: Int32, bar: Int32): + """One TMA box at coordinates ``(0, y, 0)`` of the 3D ``desc`` into shared ``dst``, completing + on the shared mbarrier ``bar``.""" + inline_ptx( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [{$r0}], [{$r1}, {{$r2}, {$r3}, {$r2}}], [{$r4}];", + read_only_args=[dst, desc, Int32(0), y, bar], + ) + + +@cute.jit +def _tma_prefetch_l2_3d(desc, y: Int32): + """Prefetch the TMA box at ``(0, y, 0)`` of the 3D ``desc`` into L2.""" + inline_ptx( + "cp.async.bulk.prefetch.tensor.3d.L2.global.tile [{$r0}, {{$r1}, {$r2}, {$r1}}];", + read_only_args=[desc, Int32(0), y], + ) + + +@dsl_user_op +def _rmsnorm_bf16x8(x0, x1, x2, x3, w0, w1, w2, w3, rstd, *, loc=None, ip=None): + """``bf16((float(x) * rstd) * (float(w) + 0))`` of eight BF16 values (four packed words + each): the normalization step of FlashInfer's ``RMSNormKernel`` (weight bias 0), each + product rounded in FP32, then to BF16 (nearest even).""" + body = "" + for i in range(4): + body += ( + f"shl.b32 xl, ${i + 4}, 16;\nand.b32 xh, ${i + 4}, 0xffff0000;\n" + f"shl.b32 wl, ${i + 8}, 16;\nand.b32 wh, ${i + 8}, 0xffff0000;\n" + "mov.b32 fxl, xl;\nmov.b32 fxh, xh;\nmov.b32 fwl, wl;\nmov.b32 fwh, wh;\n" + "add.rn.f32 fwl, fwl, 0f00000000;\nadd.rn.f32 fwh, fwh, 0f00000000;\n" + "mul.rn.f32 fxl, fxl, $12;\nmul.rn.f32 fxh, fxh, $12;\n" + "mul.rn.f32 fxl, fxl, fwl;\nmul.rn.f32 fxh, fxh, fwh;\n" + f"cvt.rn.bf16x2.f32 ${i}, fxh, fxl;\n" + ) + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(v, loc=loc, ip=ip) for v in (x0, x1, x2, x3, w0, w1, w2, w3)] + + [_fv(rstd, loc=loc, ip=ip)], + "{\n.reg .b32 xl, xh, wl, wh;\n.reg .f32 fxl, fxh, fwl, fwh;\n" + body + "}", + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +def _swz(row, chunk): + """Byte offset of 16-byte chunk ``chunk`` of row ``row`` in a 128-byte-swizzled box.""" + return row * Int32(_BOX_BYTES) + ((chunk ^ (row & Int32(7))) << Int32(4)) + + +@dsl_user_op +def _mbar_arrive(addr, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)], + "mbarrier.arrive.shared::cta.b64 _, [$0];", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bar_sync(bar_id, count: int, *, loc=None, ip=None): + """``bar.sync bar_id, count`` (named CTA barrier).""" + _asm(None, [_iv(bar_id, loc=loc, ip=ip)], f"bar.sync $0, {count};", "r", loc=loc, ip=ip) + + +@dsl_user_op +def _bar_arrive(bar_id, count: int, *, loc=None, ip=None): + """``bar.arrive bar_id, count`` (named CTA barrier, no wait).""" + _asm(None, [_iv(bar_id, loc=loc, ip=ip)], f"bar.arrive $0, {count};", "r", loc=loc, ip=ip) + + +@dsl_user_op +def _cluster_sync(*, loc=None, ip=None): + """All threads of the cluster: release this CTA's shared stores, acquire the peers'.""" + _asm( + None, + [], + "barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;", + "", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cluster_ctarank(*, loc=None, ip=None): + return Int32(_asm(T.i32(), [], "mov.u32 $0, %cluster_ctarank;", "=r", loc=loc, ip=ip)) + + +@dsl_user_op +def _bf16x2_to_f32x2(word, *, loc=None, ip=None): + """The two BF16 halves of ``word`` (low half first) as FP32 (exact).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_iv(word, loc=loc, ip=ip)], + "{\n.reg .b32 lo, hi;\nshl.b32 lo, $2, 16;\nand.b32 hi, $2, 0xffff0000;\n" + "mov.b32 $0, lo;\nmov.b32 $1, hi;\n}", + "=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _rope_pair_bf16(x, y, c, s, *, loc=None, ip=None): + """GPT-J rotation of the pair ``(x, y)`` by ``(c, s)``: ``(x c - y s, y c + x s)``, each + product and sum rounded separately in FP32 (``_fp4_mla_rope_fp32``), then to BF16 and back + to FP32 (exact).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_fv(v, loc=loc, ip=ip) for v in (x, y, c, s)], + "{\n" + ".reg .f32 xc, ys, yc, xs, e, o;\n" + ".reg .b16 he, ho;\n" + "mul.rn.f32 xc, $2, $4;\n" + "mul.rn.f32 ys, $3, $5;\n" + "sub.rn.f32 e, xc, ys;\n" + "mul.rn.f32 yc, $3, $4;\n" + "mul.rn.f32 xs, $2, $5;\n" + "add.rn.f32 o, yc, xs;\n" + "cvt.rn.bf16.f32 he, e;\n" + "cvt.rn.bf16.f32 ho, o;\n" + "cvt.f32.bf16 $0, he;\n" + "cvt.f32.bf16 $1, ho;\n" + "}", + "=f,=f,f,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _fp4_scale(amax, *, loc=None, ip=None): + """The E4M3 scale of a 16-value group and the reciprocal the codes are scaled by: + ``e4m3(max(amax / 6, 2^-9))`` (byte) and ``rcp.approx.ftz(float(that))``, the operations of + ``_fp4_mla_quantize_q_groups``.""" + res = _asm( + llvm.StructType.get_literal([T.i32(), T.f32(), T.f32()]), + [_fv(amax, loc=loc, ip=ip)], + "{\n" + ".reg .f32 s, z;\n" + ".reg .b16 h8, hl, hh;\n" + ".reg .b32 h2;\n" + "div.full.f32 s, $3, 0f40C00000;\n" + f"max.f32 s, s, 0f{_Q_SCALE_FLOOR:08X};\n" + "mov.f32 z, 0f00000000;\n" + "cvt.rn.satfinite.e4m3x2.f32 h8, z, s;\n" + "cvt.rn.f16x2.e4m3x2 h2, h8;\n" + "mov.b32 {hl, hh}, h2;\n" + "cvt.f32.f16 $1, hl;\n" + "rcp.approx.ftz.f32 $2, $1;\n" + "cvt.u32.u16 $0, h8;\n" + "and.b32 $0, $0, 0xff;\n" + "}", + "=r,=f,=f,f", + loc=loc, + ip=ip, + ) + return ( + Int32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), res, [2], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _e2m1x8(v, inv, *, loc=None, ip=None): + """E2M1 codes of the eight values ``v[i] * inv`` (``cvt.rn.satfinite``), element ``2k`` in + the low nibble of byte ``k``.""" + body = "".join(f"mul.rn.f32 x{i}, ${i + 1}, $9;\n" for i in range(8)) + body += "".join( + f"cvt.rn.satfinite.e2m1x2.f32 b{k}, x{2 * k + 1}, x{2 * k};\n" for k in range(4) + ) + body += "mov.b32 $0, {b0, b1, b2, b3};\n" + return Int32( + _asm( + T.i32(), + [_fv(x, loc=loc, ip=ip) for x in v] + [_fv(inv, loc=loc, ip=ip)], + "{\n.reg .f32 x0, x1, x2, x3, x4, x5, x6, x7;\n.reg .b8 b0, b1, b2, b3;\n" + body + "}", + "=r," + ",".join(["f"] * 9), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _e2m1x8_to_f32(word, *, loc=None, ip=None): + """The eight E2M1 codes of ``word`` (element ``2k`` in the low nibble of byte ``k``) as FP32, + through ``cvt.rn.f16x2.e2m1x2`` (exact), as ``_fp4_mla_unpack_e2m1x2``.""" + body = "mov.b32 {b0, b1, b2, b3}, $8;\n" + for k in range(4): + body += f"cvt.rn.f16x2.e2m1x2 h, b{k};\nmov.b32 {{l, u}}, h;\n" + body += f"cvt.f32.f16 ${2 * k}, l;\ncvt.f32.f16 ${2 * k + 1}, u;\n" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 8), + [_iv(word, loc=loc, ip=ip)], + "{\n.reg .b8 b0, b1, b2, b3;\n.reg .b32 h;\n.reg .b16 l, u;\n" + body + "}", + ",".join(["=f"] * 8) + ",r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 8, loc=loc, ip=ip) + + +@dsl_user_op +def _residual(v, main, scale, *, loc=None, ip=None): + """``v - main * scale`` as ``_fp4_mla_quantize_q_residual_groups`` computes it.""" + return Float32( + _asm( + T.f32(), + [_fv(x, loc=loc, ip=ip) for x in (v, main, scale)], + "{\n.reg .f32 p;\nneg.f32 p, $2;\nfma.rn.f32 $0, p, $3, $1;\n}", + "=f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _abs_max16(v, *, loc=None, ip=None): + """``max |v[i]|`` of 16 FP32 values (``abs.f32`` / ``max.f32``).""" + body = "".join(f"abs.f32 a{i}, ${i + 1};\n" for i in range(16)) + body += "".join(f"max.f32 a{i}, a{i}, a{i + 8};\n" for i in range(8)) + body += "".join(f"max.f32 a{i}, a{i}, a{i + 4};\n" for i in range(4)) + body += "max.f32 a0, a0, a2;\nmax.f32 a1, a1, a3;\nmax.f32 $0, a0, a1;\n" + regs = ", ".join(f"a{i}" for i in range(16)) + return Float32( + _asm( + T.f32(), + [_fv(x, loc=loc, ip=ip) for x in v], + "{\n.reg .f32 " + regs + ";\n" + body + "}", + "=f," + ",".join(["f"] * 16), + loc=loc, + ip=ip, + ) + ) + + +def _quant_group16(v): + """FP4 codes (two words) and the E4M3 scale byte of 16 FP32 values; also the scale as FP32.""" + sf, sf_f32, inv = _fp4_scale(_abs_max16(v)) + return _e2m1x8(v[:8], inv), _e2m1x8(v[8:], inv), sf, sf_f32 + + +def _sf_offset(row, col: Int32): + """Byte offset of scale column ``col`` of query row ``row`` in the swizzled 128x4 layout + (``_fp4_mla_q_sf_offset``).""" + padded = _ceil_div(Q_SF_COLS, 4) * 4 + return ( + Int64(col % Int32(4)) + + Int64(col // Int32(4)) * Int64(4 * 128) + + Int64(row % Int32(32)) * Int64(16) + + Int64((row % Int32(128)) // Int32(32)) * Int64(4) + + Int64(row // Int32(128)) * Int64(128 * padded) + ) + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +class K3MlaQFusedKernel: + """Fused MLA query of ``t_pad`` (a multiple of 8) tokens; see the module docstring. + + Args: + t_pad: Tokens rounded up to a multiple of 8 (the MMA N tiles); rows past the launch's + token count are computed on zeros and not stored. + use_pdl: Programmatic dependent launch: weights are fetched before + ``griddepcontrol.wait``. + debug: Also store the BF16 q_b output and the BF16 absorbed / rotated query (tests). + pre_tiles: q_b tiles issued before ``griddepcontrol.wait`` (the rest right after it). + fuse_norm: ``x`` is the raw q_a latent: the kernel applies the q_a RMSNorm itself + (FlashInfer ``RMSNormKernel`` arithmetic, see :meth:`K3MlaQFusedKernel.row_rstd`). + timing: Thread 0 of every CTA stores ``%globaltimer`` at its phase boundaries into + ``ts[cta * 16 + phase]`` (microbenchmarks): 0 entry, 1 after the wait, 2 activation + quantized, 8 warp 0's q_b tile landed, 3 its MMAs done, 4 its q_b output stored, 5 + after the cluster barrier, 9 W_UK landed, 6 absorbed, 7 end. + early_trigger: (``use_pdl``) release the dependent grid at entry instead of after the + query stores, so a dependent that only reads the query after its own + ``griddepcontrol.wait`` (the FP4 decode with ``pdl_wait``) runs its prologue on the + SMs this grid's CTAs leave as they finish. + """ + + def __init__( + self, + t_pad: int, + use_pdl: bool, + debug: bool = False, + timing: bool = False, + fuse_norm: bool = False, + pre_tiles: int = _NTILE, + early_trigger: bool = False, + ): + if t_pad % 8 or not 8 <= t_pad <= MAX_TOKENS: + raise ValueError(f"t_pad={t_pad}") + self.t_pad = t_pad + self.nt = t_pad // 8 + self.use_pdl = use_pdl + self.debug = debug + self.timing = timing + self.fuse_norm = fuse_norm + if not 0 <= pre_tiles <= _NTILE: + raise ValueError(f"pre_tiles={pre_tiles}") + self.pre_tiles = pre_tiles + self.early_trigger = early_trigger and use_pdl + assert t_pad * _OB_PITCH <= (_NTILE - _UK_TILES) * _TILE_BYTES + + @cute.jit + def __call__( + self, + x_addr: Int64, + wqb_ptr: cute.Pointer, + ew_addr: Int64, + wuk_ptr: cute.Pointer, + gs_addr: Int64, + cos_sin_addr: Int64, + pos_addr: Int64, + qfp4_addr: Int64, + qsf_addr: Int64, + dbg_qb_addr: Int64, + dbg_q_addr: Int64, + ts_addr: Int64, + nw_addr: Int64, + eps: Float32, + x_ld: Int32, + num_tokens: Int32, + gen_len: Int32, + rows_per_seq: Int32, + stream: cuda.CUstream, + ): + t_pad = self.t_pad + + @cute.struct + class SharedStorage: + mbar: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 16], 128] + ew: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, 128], 128] + bsc: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, 2 * _KGROUPS * t_pad], 128] + act: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, t_pad * _A_PITCH], 128] + ring: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, _NTILE * _TILE_BYTES], 1024] + qn: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, t_pad * _QN_PITCH], 128] + qpe: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, t_pad * _QPE_PITCH], 128] + nw: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, Q_LORA_RANK * 2], 128] + cs: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, t_pad * 4 * 64], 128] + pos: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, t_pad], 128] + part: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, _NTILE * 32 * t_pad // 2], 128 + ] + + self.shared_storage = SharedStorage + # q_b weight [N, 1536] as (128 bytes, row, k block) and W_UK [H * 512, 128] as (64 values, + # row, half): a box of whole rows lands as [k block][row][128 bytes]. + w_qb = cute.make_tensor( + wqb_ptr, + cute.make_layout( + (_BOX_BYTES, NUM_HEADS * QK_HEAD_DIM, _KB), stride=(1, Q_LORA_RANK, _BOX_BYTES) + ), + ) + w_uk = cute.make_tensor( + wuk_ptr, + cute.make_layout( + (_BOX_BYTES // 2, NUM_HEADS * KV_LORA_RANK, 2), + stride=(1, NOPE_DIM, _BOX_BYTES // 2), + ), + ) + wqb_desc = cuda_tma.create_tensor_map_tiled_from_view( + w_qb, + box_dims=(_BOX_BYTES, _TILE_ROWS, _KB), + stride_order=(0, 1, 2), + swizzle=cuda_tma.TensorMapSwizzle.s128b, + ) + wuk_desc = cuda_tma.create_tensor_map_tiled_from_view( + w_uk, + box_dims=(_BOX_BYTES // 2, _CTA_COLS, 2), + stride_order=(0, 1, 2), + swizzle=cuda_tma.TensorMapSwizzle.s128b, + ) + self.kernel( + x_addr, wqb_desc, ew_addr, wuk_desc, gs_addr, cos_sin_addr, pos_addr, qfp4_addr, + qsf_addr, dbg_qb_addr, dbg_q_addr, ts_addr, nw_addr, eps, x_ld, num_tokens, gen_len, + rows_per_seq, + ).launch( + grid=(2 * NUM_HEADS, 1, 1), + block=(_THREADS, 1, 1), + cluster=(2, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) # fmt: skip + + @cute.kernel + def kernel( + self, + x_addr: Int64, + wqb_desc: cutlass.GridConstant[cuda_tma.TensorMap], + ew_addr: Int64, + wuk_desc: cutlass.GridConstant[cuda_tma.TensorMap], + gs_addr: Int64, + cos_sin_addr: Int64, + pos_addr: Int64, + qfp4_addr: Int64, + qsf_addr: Int64, + dbg_qb_addr: Int64, + dbg_q_addr: Int64, + ts_addr: Int64, + nw_addr: Int64, + eps: Float32, + x_ld: Int32, + num_tokens: Int32, + gen_len: Int32, + rows_per_seq: Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + ts = ts_addr + Int64(bidx) * 128 + self.stamp(tidx, ts, 0) + warp = tidx // 32 + lane = tidx % 32 + head = bidx // Int32(2) + rank = _cluster_ctarank() + peer = rank ^ Int32(1) + row0 = head * Int32(QK_HEAD_DIM) + rank * Int32(_CTA_ROWS) + + storage = utils.SmemAllocator().allocate(self.shared_storage) + mbar = storage.mbar.data_ptr().toint() + full_bar = mbar + uk_free_bar = mbar + Int32(8 * _NTILE) + wuk_bar = mbar + Int32(8 * (_NTILE + 1)) + ew_s = storage.ew.data_ptr().toint() + bsc_s = storage.bsc.data_ptr().toint() + act_s = storage.act.data_ptr().toint() + ring_s = storage.ring.data_ptr().toint() + qn_s = storage.qn.data_ptr().toint() + qpe_s = storage.qpe.data_ptr().toint() + nw_s = storage.nw.data_ptr().toint() + cs_s = storage.cs.data_ptr().toint() + pos_s = storage.pos.data_ptr().toint() + part_s = storage.part.data_ptr().toint() + wuk_s = ring_s + ob_s = ring_s + Int32(_OB_OFFSET) + wqb_tmap = wqb_desc.get_ptr() + wuk_tmap = wuk_desc.get_ptr() + + # ---- before the wait: barriers, block exponents, weights (none depends on the + # previous grid). + if cutlass.const_expr(self.early_trigger): + _griddepcontrol_launch_dependents() + if tidx == 0: + for s in cutlass.range_constexpr(_NTILE): + _mbar_init(full_bar + Int32(8 * s), 1) + _mbar_init(uk_free_bar, _UK_TILES * _QB_SPLIT_K) + _mbar_init(wuk_bar, 1) + if tidx < Int32(_NTILE * _KBLOCKS): + tile = tidx // Int32(_KBLOCKS) + kb = tidx - tile * Int32(_KBLOCKS) + block_row = (row0 + tile * Int32(_TILE_ROWS)) // Int32(128) + _sts_u8(ew_s + tidx, _ldg_u8(ew_addr + Int64(block_row * Int32(_KBLOCKS) + kb))) + if cutlass.const_expr(self.fuse_norm): + if tidx < Int32(Q_LORA_RANK * 2 // 16): + _sts_v4_b32(nw_s + tidx * Int32(16), *_ldg_v4(nw_addr + Int64(tidx) * 16)) + cute.arch.barrier() + if tidx == 0: + # Every 16-row q_b tile (all of K) into its own slot. + for tile in cutlass.range_constexpr(self.pre_tiles): + self.load_tile(tile, row0, wqb_tmap, ring_s, full_bar) + + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + self.stamp(tidx, ts, 1) + if cutlass.const_expr(self.pre_tiles < _NTILE): + if tidx == 0: + for tile in cutlass.range_constexpr(self.pre_tiles, _NTILE): + self.load_tile(tile, row0, wqb_tmap, ring_s, full_bar) + + # ---- Activation: warp w owns rows w, w + 12, .. ; lane l holds elements 256 b + 8 l .. + 7 + # (b = 0..5) of a row, the layout of FlashInfer's RMSNormKernel for a 1536-wide row. + rows = [] + for i in cutlass.range_constexpr(_ceil_div(self.t_pad, _WARPS)): + t = warp + Int32(i * _WARPS) + src = ( + x_addr + + ( + Int64(cutlass.max(cutlass.min(t, num_tokens - Int32(1)), Int32(0))) + * Int64(x_ld) + + Int64(lane) * 8 + ) + * 2 + ) + rows.append( + (t, [list(_ldg_v4_pred(src + Int64(512 * b), num_tokens - t)) for b in range(6)]) + ) + for t, xw in rows: + if t < Int32(self.t_pad): + self.quantize_row(t, xw, lane, eps, nw_s, ew_s, act_s, bsc_s) + cute.arch.barrier() + self.stamp(tidx, ts, 2) + if tidx == 0: + # W_UK into L2 behind the q_b weight (it is read after the q_b MMAs). + _tma_prefetch_l2_3d(wuk_tmap, head * Int32(KV_LORA_RANK) + rank * Int32(_CTA_COLS)) + + # ---- q_b: warp w runs K half w // 6 of the CTA's 16-row tile w % 6. + self.q_b_tile(warp, lane, rank, peer, head, row0, bsc_s, act_s, ring_s, qn_s, qpe_s, + part_s, full_bar, uk_free_bar, dbg_qb_addr, num_tokens, tidx, ts) # fmt: skip + self.stamp(tidx, ts, 4) + if warp == Int32(_WARPS - 1): + if rank == Int32(1): + self.prefetch_rope(lane, num_tokens, pos_addr, cos_sin_addr, pos_s, cs_s) + if warp == Int32(0): + # W_UK into the slots of the first q_b tiles once their warps are done with them. + self.wait(uk_free_bar, Int32(0)) + if lane == 0: + _mbar_expect_tx(wuk_bar, _UK_BYTES) + uk_row = head * Int32(KV_LORA_RANK) + rank * Int32(_CTA_COLS) + _tma_load_3d(wuk_s, wuk_tmap, uk_row, wuk_bar) + _cluster_sync() + self.stamp(tidx, ts, 5) + + # ---- W_UK: warp w < 8 owns absorbed columns 32 w .. 32 w + 31 of the CTA's 256. + if warp < Int32(_ABSORB_WARPS): + self.wait(wuk_bar, Int32(0)) + self.stamp(tidx, ts, 9) + self.absorb(warp, lane, wuk_s, qn_s, ob_s) + cute.arch.barrier() + self.stamp(tidx, ts, 6) + + # ---- FP4 query rows. + # Absorbed groups (16 per token) and, on CTA 1, RoPE groups (4 per token) in one pass. + gs = _ldg_f32(gs_addr) + n_prefix = Int32(self.t_pad * 16) + n_items = n_prefix + rank * Int32(self.t_pad * 4) + for item in cutlass.range(tidx, n_items, Int32(_THREADS), unroll=1): + if item < n_prefix: + t = item // Int32(16) + grp = item - t * Int32(16) + if t < num_tokens: + self.store_prefix_group(t, grp, rank, head, gs, ob_s, qfp4_addr, qsf_addr, + dbg_q_addr, gen_len, rows_per_seq) # fmt: skip + else: + t = (item - n_prefix) // Int32(4) + rg = item - n_prefix - t * Int32(4) + if t < num_tokens: + if _lds_b32(pos_s + t * Int32(4)) >= Int32(0): + self.store_rope_group(t, rg, head, gs, qpe_s, cs_s, qfp4_addr, qsf_addr, + dbg_q_addr, gen_len, rows_per_seq) # fmt: skip + if cutlass.const_expr(self.use_pdl and not self.early_trigger): + _griddepcontrol_launch_dependents() + self.stamp(tidx, ts, 7) + + @cute.jit + def quantize_row( + self, + t: Int32, + xw, + lane: Int32, + eps: Float32, + nw_s: Int32, + ew_s: Int32, + act_s: Int32, + bsc_s: Int32, + ): + """Activation row ``t``, held by the warp as ``xw[b]`` (four words of elements + ``256 b + 8 l .. + 7`` per lane ``l``): normalized (``fuse_norm``: :meth:`row_rstd`, then + the norm kernel's per-element products), then the low-M GEMM's MXFP8 quantization (the + amax of a 32-value group, lanes ``4 m .. 4 m + 3`` of a block, by shuffles; its UE8M0 + exponent and E4M3 codes) into shared memory, with the group's MMA scales.""" + if cutlass.const_expr(self.fuse_norm): + rstd = self.row_rstd(xw, eps) + for b in cutlass.range_constexpr(6): + wn = _lds_v4_b32(nw_s + (Int32(256 * b) + lane * Int32(8)) * Int32(2)) + xw[b] = list(_rmsnorm_bf16x8(*xw[b], *wn, rstd)) + # The six groups' amax shuffles are interleaved (independent chains). + amaxes = [_amax_bf16x8_words(*xw[b]) for b in range(6)] + amaxes = [_amax_pair(m, _shfl_bfly_b32(m, 1)) for m in amaxes] + amaxes = [_amax_pair(m, _shfl_bfly_b32(m, 2)) for m in amaxes] + for b in cutlass.range_constexpr(6): + e_a = _ue8m0_of_amax(amaxes[b]) + q0, q1 = _quant_bf16x8_to_e4m3(*xw[b], _bf16x2_inv_scale(e_a)) + _sts_v2_b32(act_s + t * Int32(_A_PITCH) + Int32(256 * b) + lane * Int32(8), q0, q1) + if lane % 4 == 0: + # The group's MMA scale for both 128-row weight blocks the CTA's rows touch + # (tiles 0 and 5 lie in the first and the last of them). + kg = Int32(8 * b) + lane // Int32(4) + kb = Int32(2 * b) + lane // Int32(16) + for rb in cutlass.range_constexpr(2): + w_exp = _lds_u8(ew_s + Int32(rb * (_NTILE - 1) * _KBLOCKS) + kb) + _sts_f32(bsc_s + ((Int32(rb * _KGROUPS) + kg) * Int32(self.t_pad) + t) * 4, + _block_scale(e_a, w_exp)) # fmt: skip + + @cute.jit + def load_tile(self, tile: int, row0: Int32, wqb_tmap, ring_s: Int32, full_bar: Int32): + """One thread: q_b tile ``tile`` (16 whole rows) into its slot.""" + bar = full_bar + Int32(8 * tile) + _mbar_expect_tx(bar, _TILE_BYTES) + _tma_load_3d( + ring_s + Int32(tile * _TILE_BYTES), wqb_tmap, row0 + Int32(tile * _TILE_ROWS), bar + ) + + @cute.jit + def row_rstd(self, xw, eps: Float32): + """``rsqrt(mean(x^2) + eps)`` of a 1536-wide BF16 row held as FlashInfer's + ``RMSNormKernel`` holds it (32 lanes per row, vector width 8, six vector blocks: ``xw[b]`` + are the four words of elements ``256 b + 8 l .. + 7``), with that kernel's operations: + the 48 values squared and summed as one ``TensorSSA`` (block-major), ``warp_reduce`` + (butterfly), the division by the width and the fast ``rsqrt``. The same DSL operations on + the same values in the same order give the same bits.""" + xs = cute.make_rmem_tensor((48,), Float32) + for b in cutlass.range_constexpr(6): + for k in cutlass.range_constexpr(4): + lo, hi = _bf16x2_to_f32x2(xw[b][k]) + xs[8 * b + 2 * k] = lo + xs[8 * b + 2 * k + 1] = hi + x = xs.load() + x_sq = x * x + local_val = x_sq.reduce(cute.ReductionOp.ADD, init_val=Float32(0.0), reduction_profile=0) + sum_sq = _flashinfer_warp_reduce(local_val, operator.add, width=32) + mean_sq = sum_sq / Float32(Q_LORA_RANK) + return cute.math.rsqrt(mean_sq + eps, fastmath=True) + + @cute.jit + def stamp(self, tidx: Int32, ts: Int64, phase: int): + """``timing``: thread 0 stores ``%globaltimer`` as phase ``phase`` of this CTA.""" + if cutlass.const_expr(self.timing): + if tidx == 0: + _st_u64(ts + Int64(phase * 8), _globaltimer()) + + @cute.jit + def wait(self, bar: Int32, phase: Int32): + done = Int32(0) + while done == 0: + done = _mbar_try_wait(bar, phase) + + @cute.jit + def q_b_tile( + self, + warp: Int32, + lane: Int32, + rank: Int32, + peer: Int32, + head: Int32, + row0: Int32, + bsc_s: Int32, + act_s: Int32, + ring_s: Int32, + qn_s: Int32, + qpe_s: Int32, + part_s: Int32, + full_bar: Int32, + uk_free_bar: Int32, + dbg_qb_addr: Int64, + num_tokens: Int32, + tidx: Int32, + ts: Int64, + ): + """K half ``warp // 6`` of the CTA's 16-row q_b tile ``j = warp % 6``: FP32 sums in K order; + the second half's warp hands its sums to the first half's warp (shared memory, named + barrier ``1 + j``), which adds them (the low-M GEMM's split-K association), rounds to BF16 + and stores q_nope (both CTAs) or q_pe.""" + nt = self.nt + j = warp % Int32(_NTILE) + half = warp // Int32(_NTILE) + g = lane // 4 + q = lane % 4 + acc = cute.make_rmem_tensor((4 * nt,), Float32) + for i in cutlass.range_constexpr(4 * nt): + acc[i] = Float32(0.0) + # ldmatrix row addresses: A (16 rows x 32 K bytes), B (8 tokens x 32 K bytes per n tile; + # with .x4 two n tiles). + a_row = (lane & 7) + ((lane >> 3) & 1) * 8 + a_hi = lane >> 4 + b_tok = (lane & 7) + (lane >> 4) * Int32(8) + b_off = b_tok * Int32(_A_PITCH) + ((lane >> 3) & 1) * Int32(16) + # Scales of this tile's 128-row weight block (the CTA's first or second). + rb = ((row0 + j * Int32(_TILE_ROWS)) >> 7) - (row0 >> 7) + s_base = bsc_s + (rb * Int32(_KGROUPS * self.t_pad) + q * Int32(2)) * Int32(4) + self.wait(full_bar + j * Int32(8), Int32(0)) + self.stamp(tidx, ts, 8) + tile = ring_s + j * Int32(_TILE_BYTES) + a_off = a_row * Int32(_BOX_BYTES) + kb_half = _KB // _QB_SPLIT_K + # One 128-byte K block (four K32 steps) per iteration: its fragments and scales are all + # loaded and its MMAs all issued before the scaled sums, which add the steps in K order. + for kb in cutlass.range(half * Int32(kb_half), (half + 1) * Int32(kb_half), 1, unroll=1): + plane = tile + kb * Int32(_TILE_ROWS * _BOX_BYTES) + a_off + frags = [] + for kk in cutlass.range_constexpr(4): + kg = kb * Int32(4) + Int32(kk) + a = _ldsm_x4(plane + ((Int32(kk * 2) + a_hi) ^ (a_row & Int32(7))) * 16) + bs = [] + for p in cutlass.range_constexpr(_ceil_div(nt, 2)): + addr = act_s + b_off + Int32(p * 16 * _A_PITCH) + kg * Int32(32) + if cutlass.const_expr(nt == 1): + bs.append(_ldsm_x2(addr)) + else: + r4 = _ldsm_x4(addr) + bs.append((r4[0], r4[1])) + bs.append((r4[2], r4[3])) + scales = [ + _lds_v2_f32(s_base + (kg * Int32(self.t_pad) + Int32(n * 8)) * Int32(4)) + for n in range(nt) + ] + frags.append((a, bs, scales)) + prods = [ + [_mma_e4m3(a[0], a[1], a[2], a[3], bs[n][0], bs[n][1]) for n in range(nt)] + for a, bs, _ in frags + ] + for kk in cutlass.range_constexpr(4): + for n in cutlass.range_constexpr(nt): + d = prods[kk][n] + s0, s1 = frags[kk][2][n] + acc[4 * n + 0] = _fma_rn(d[0], s0, acc[4 * n + 0]) + acc[4 * n + 1] = _fma_rn(d[1], s1, acc[4 * n + 1]) + acc[4 * n + 2] = _fma_rn(d[2], s0, acc[4 * n + 2]) + acc[4 * n + 3] = _fma_rn(d[3], s1, acc[4 * n + 3]) + self.stamp(tidx, ts, 3) + if j < Int32(_UK_TILES): + cute.arch.sync_warp() + if lane == 0: + _mbar_arrive(uk_free_bar) + part = part_s + (j * Int32(32) + lane) * Int32(4 * nt * 4) + if half == Int32(1): + for i in cutlass.range_constexpr(nt): + _sts_v4_f32(part + Int32(16 * i), acc[4 * i], acc[4 * i + 1], acc[4 * i + 2], + acc[4 * i + 3]) # fmt: skip + _bar_arrive(Int32(1) + j, 64) + else: + _bar_sync(Int32(1) + j, 64) + for i in cutlass.range_constexpr(nt): + other = _lds_v4_f32(part + Int32(16 * i)) + for c in cutlass.range_constexpr(4): + acc[4 * i + c] = _add_rn(acc[4 * i + c], other[c]) + # Rows of this tile within the head's 192 (q_nope dims, then the RoPE dims). + hr = rank * Int32(_CTA_ROWS) + j * Int32(_TILE_ROWS) + g + is_nope = rank * Int32(_CTA_ROWS) + j * Int32(_TILE_ROWS) < Int32(NOPE_DIM) + for n in cutlass.range_constexpr(nt): + for h in cutlass.range_constexpr(2): + row = hr + Int32(8 * h) + for c in cutlass.range_constexpr(2): + t = Int32(n * 8) + q * Int32(2) + Int32(c) + value = acc[4 * n + 2 * h + c] + if is_nope: + _sts_bf16_both( + qn_s + t * Int32(_QN_PITCH) + row * Int32(2), peer, value + ) + else: + _sts_bf16( + qpe_s + t * Int32(_QPE_PITCH) + (row - Int32(NOPE_DIM)) * 2, value + ) + if cutlass.const_expr(self.debug): + if t < num_tokens: + _stg_bf16( + dbg_qb_addr + + (Int64(t) * (NUM_HEADS * QK_HEAD_DIM) + + Int64(head) * QK_HEAD_DIM + Int64(row)) * 2, + value, + ) # fmt: skip + + @cute.jit + def absorb(self, warp: Int32, lane: Int32, wuk_s: Int32, qn_s: Int32, ob_s: Int32): + nt = self.nt + g = lane // 4 + q = lane % 4 + a_row = (lane & 7) + ((lane >> 3) & 1) * 8 + a_hi = lane >> 4 + b_tok = (lane & 7) + (lane >> 4) * Int32(8) + b_off = b_tok * Int32(_QN_PITCH) + ((lane >> 3) & 1) * Int32(16) + # The two tiles' MMA chains (each in K order) are interleaved. + accs = [[[Float32(0.0)] * 4 for _ in range(nt)] for _ in range(2)] + for kk in cutlass.range_constexpr(NOPE_DIM // 16): + bs = [] + for p in cutlass.range_constexpr(_ceil_div(nt, 2)): + addr = qn_s + b_off + Int32(p * 16 * _QN_PITCH + kk * 32) + if cutlass.const_expr(nt == 1): + bs.append(_ldsm_x2(addr)) + else: + r4 = _ldsm_x4(addr) + bs.append((r4[0], r4[1])) + bs.append((r4[2], r4[3])) + box = wuk_s + Int32((kk // 4) * _UK_HALF) + a_mm = [ + _ldsm_x4(box + _swz((warp * Int32(2) + Int32(mm)) * Int32(16) + a_row, + Int32((kk % 4) * 2) + a_hi)) + for mm in range(2) + ] # fmt: skip + for mm in cutlass.range_constexpr(2): + a = a_mm[mm] + for n in cutlass.range_constexpr(nt): + accs[mm][n] = list( + _mma_bf16(a[0], a[1], a[2], a[3], bs[n][0], bs[n][1], *accs[mm][n]) + ) + for mm in cutlass.range_constexpr(2): + m = warp * Int32(2) + Int32(mm) + acc = accs[mm] + col = m * Int32(16) + g + for n in cutlass.range_constexpr(nt): + for h in cutlass.range_constexpr(2): + for c in cutlass.range_constexpr(2): + t = Int32(n * 8) + q * Int32(2) + Int32(c) + _sts_bf16( + ob_s + t * Int32(_OB_PITCH) + (col + Int32(8 * h)) * Int32(2), + acc[n][2 * h + c], + ) + + @cute.jit + def q_row(self, t: Int32, head: Int32, gen_len: Int32, rows_per_seq: Int32): + seq = t // gen_len + return seq * rows_per_seq + (t - seq * gen_len) * Int32(NUM_HEADS) + head + + @cute.jit + def store_prefix_group( + self, + t: Int32, + grp: Int32, + rank: Int32, + head: Int32, + gs: Float32, + ob_s: Int32, + qfp4_addr: Int64, + qsf_addr: Int64, + dbg_q_addr: Int64, + gen_len: Int32, + rows_per_seq: Int32, + ): + """Absorbed values ``16 grp .. 16 grp + 15`` of this CTA's half of token ``t``: FP4 codes + and scale of ``float(bf16) * gs`` (``_fp4_mla_quantize_q_groups_words``).""" + src = ob_s + t * Int32(_OB_PITCH) + grp * Int32(32) + words = list(_lds_v4_b32(src)) + list(_lds_v4_b32(src + Int32(16))) + v = [] + for w in words: + lo, hi = _bf16x2_to_f32x2(w) + v.append(_mul_rn(lo, gs)) + v.append(_mul_rn(hi, gs)) + w0, w1, sf, _ = _quant_group16(v) + col = rank * Int32(16) + grp + qrow = self.q_row(t, head, gen_len, rows_per_seq) + _stg_v2(qfp4_addr + Int64(qrow) * Q_ROW_BYTES + Int64(col) * 8, w0, w1) + _stg_u8(qsf_addr + _sf_offset(qrow, col), sf) + if cutlass.const_expr(self.debug): + dst = ( + dbg_q_addr + + ( + (Int64(t) * NUM_HEADS + Int64(head)) * (KV_LORA_RANK + ROPE_DIM) + + Int64(col) * 16 + ) + * 2 + ) + _stg_v4(dst, *words[:4]) + _stg_v4(dst + Int64(16), *words[4:]) + + @cute.jit + def prefetch_rope( + self, + lane: Int32, + num_tokens: Int32, + pos_addr: Int64, + cos_sin_addr: Int64, + pos_s: Int32, + cs_s: Int32, + ): + """One warp: every token's position, and the 16 rotary ``(cos, sin)`` values of each of + its four RoPE groups at that position, into shared memory (item ``4 t + g``, 64 bytes).""" + for j in cutlass.range_constexpr(_ceil_div(self.t_pad * 4, 32)): + item = lane + Int32(32 * j) + t = item // Int32(4) + rg = item - t * Int32(4) + if t < num_tokens: + pos = _ldg_s32(pos_addr + Int64(t) * 4) + if rg == Int32(0): + _sts_b32(pos_s + t * Int32(4), pos) + if pos >= Int32(0): + table = cos_sin_addr + (Int64(pos) * (2 * ROPE_DIM) + Int64(rg) * 16) * 4 + words = [_ldg_v4(table + Int64(16 * i)) for i in range(4)] + for i in cutlass.range_constexpr(4): + _sts_v4_b32(cs_s + item * Int32(64) + Int32(16 * i), *words[i]) + + @cute.jit + def store_rope_group( + self, + t: Int32, + rg: Int32, + head: Int32, + gs: Float32, + qpe_s: Int32, + cs_s: Int32, + qfp4_addr: Int64, + qsf_addr: Int64, + dbg_q_addr: Int64, + gen_len: Int32, + rows_per_seq: Int32, + ): + """RoPE values ``16 rg .. 16 rg + 15`` of token ``t``: rotated by the prefetched table + values of its position, rounded to BF16, then the main and residual FP4 groups + (``_fp4_mla_quantize_q_residual_groups``).""" + src = qpe_s + t * Int32(_QPE_PITCH) + rg * Int32(32) + words = list(_lds_v4_b32(src)) + list(_lds_v4_b32(src + Int32(16))) + cs = [] + for i in cutlass.range_constexpr(4): + cs += list(_lds_v4_f32(cs_s + (t * Int32(4) + rg) * Int32(64) + Int32(16 * i))) + rot = [] + for k in cutlass.range_constexpr(8): + x, y = _bf16x2_to_f32x2(words[k]) + e, o = _rope_pair_bf16(x, y, cs[2 * k], cs[2 * k + 1]) + rot.append(e) + rot.append(o) + v = [_mul_rn(x, gs) for x in rot] + m0, m1, sf, sf_f32 = _quant_group16(v) + main = list(_e2m1x8_to_f32(m0)) + list(_e2m1x8_to_f32(m1)) + res = [_residual(v[i], main[i], sf_f32) for i in range(16)] + r0, r1, rsf, _ = _quant_group16(res) + col = Int32(KV_LORA_RANK // FP4_GROUP) + rg * Int32(2) + qrow = self.q_row(t, head, gen_len, rows_per_seq) + dst = qfp4_addr + Int64(qrow) * Q_ROW_BYTES + Int64(col) * 8 + _stg_v4(dst, m0, m1, r0, r1) + _stg_u8(qsf_addr + _sf_offset(qrow, col), sf) + _stg_u8(qsf_addr + _sf_offset(qrow, col + Int32(1)), rsf) + if cutlass.const_expr(self.debug): + out = ( + dbg_q_addr + + ( + (Int64(t) * NUM_HEADS + Int64(head)) * (KV_LORA_RANK + ROPE_DIM) + + KV_LORA_RANK + + Int64(rg) * 16 + ) + * 2 + ) + packed = [_pack_bf16x2(rot[2 * k], rot[2 * k + 1]) for k in range(8)] + _stg_v4(out, *packed[:4]) + _stg_v4(out + Int64(16), *packed[4:]) + + +class K3MlaLatentRowKernel: + """The latent row of the MLA KV half: ``out[t] = [rmsnorm(c_kv[t], w, eps) | k_pe[t]]`` for a + 512-wide ``c_kv`` and 64-wide ``k_pe`` (BF16), one warp per token, with the arithmetic of + FlashInfer's ``RMSNormKernel`` for a 512-wide row (lane ``l`` holds elements ``256 b + 8 l .. + + 7``, b = 0, 1; the same reduction and products as :meth:`K3MlaQFusedKernel.row_rstd`), so + ``out`` is bitwise ``torch.cat([rmsnorm(c_kv), k_pe], -1)``. One launch instead of the norm + and the ``k_pe`` copy. PDL: the weight is read before ``griddepcontrol.wait``; dependents are + released right after it (every earlier grid is then complete).""" + + def __init__(self, use_pdl: bool): + self.use_pdl = use_pdl + + @cute.jit + def __call__( + self, + ckv_addr: Int64, + kpe_addr: Int64, + w_addr: Int64, + out_addr: Int64, + eps: Float32, + ckv_ld: Int32, + kpe_ld: Int32, + out_ld: Int32, + num_tokens: Int32, + grid: Int32, + stream: cuda.CUstream, + ): + self.kernel( + ckv_addr, kpe_addr, w_addr, out_addr, eps, ckv_ld, kpe_ld, out_ld, num_tokens + ).launch(grid=(grid, 1, 1), block=(128, 1, 1), stream=stream, use_pdl=self.use_pdl) + + @cute.kernel + def kernel( + self, + ckv_addr: Int64, + kpe_addr: Int64, + w_addr: Int64, + out_addr: Int64, + eps: Float32, + ckv_ld: Int32, + kpe_ld: Int32, + out_ld: Int32, + num_tokens: Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + lane = tidx % 32 + t = bidx * Int32(4) + tidx // Int32(32) + col = lane * Int32(8) + ww = [_ldg_v4(w_addr + Int64(Int32(256 * b) + col) * 2) for b in range(2)] + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + _griddepcontrol_launch_dependents() + if t < num_tokens: + row = ckv_addr + Int64(t) * Int64(ckv_ld) * 2 + xw = [list(_ldg_v4(row + Int64(Int32(256 * b) + col) * 2)) for b in range(2)] + kpe = _ldg_v4_pred(kpe_addr + (Int64(t) * Int64(kpe_ld) + Int64(col)) * 2, + Int32(8) - lane) # fmt: skip + xs = cute.make_rmem_tensor((16,), Float32) + for b in cutlass.range_constexpr(2): + for k in cutlass.range_constexpr(4): + lo, hi = _bf16x2_to_f32x2(xw[b][k]) + xs[8 * b + 2 * k] = lo + xs[8 * b + 2 * k + 1] = hi + x = xs.load() + x_sq = x * x + local_val = x_sq.reduce( + cute.ReductionOp.ADD, init_val=Float32(0.0), reduction_profile=0 + ) + sum_sq = _flashinfer_warp_reduce(local_val, operator.add, width=32) + rstd = cute.math.rsqrt(sum_sq / Float32(KV_LORA_RANK) + eps, fastmath=True) + out = out_addr + Int64(t) * Int64(out_ld) * 2 + for b in cutlass.range_constexpr(2): + y = _rmsnorm_bf16x8(*xw[b], *ww[b], rstd) + _stg_v4(out + Int64(Int32(256 * b) + col) * 2, *y) + if lane < Int32(ROPE_DIM // 8): + _stg_v4(out + Int64(Int32(KV_LORA_RANK) + col) * 2, *kpe) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + +_COMPILED: dict = {} + + +@functools.cache +def _device_ok(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) == (10, 7) + + +def block_exponents(weight_scale: torch.Tensor, n: int, k: int) -> torch.Tensor: + """UE8M0 exponents ``uint8 [ceil(n / 128), k / 128]`` of a 128x128 block-scaled FP8 weight + from its scale as the K3 FP8 projections hold it: power-of-two FP32 block scales, or the + Rubin MXFP8 GEMM's R128c4 UE8M0 layout (``transform_k128_scales_to_cutedsl_mxfp8_layout``, + whose 512-byte atoms start with the block exponent).""" + nb, kb = _ceil_div(n, 128), k // 128 + if weight_scale.dtype == torch.uint8 and weight_scale.numel() == nb * kb * 512: + return weight_scale.reshape(nb, kb, 512)[:, :, 0].contiguous() + if weight_scale.dtype == torch.float32 and weight_scale.numel() == nb * kb: + bits = weight_scale.reshape(nb, kb).contiguous().view(torch.int32) + if bool(((bits & 0x7FFFFF) != 0).any()): + raise ValueError("block scales must be powers of two") + return ((bits >> 23) & 0xFF).to(torch.uint8) + raise ValueError( + f"unsupported q_b weight scale {tuple(weight_scale.shape)} {weight_scale.dtype}" + ) + + +def weights_supported(w_qb: torch.Tensor, w_uk: torch.Tensor) -> bool: + """Whether :func:`k3_mla_q_fused` runs with these q_b / W_UK weights on their device.""" + return bool( + w_qb.is_cuda + and _device_ok(w_qb.device.index if w_qb.device.index is not None else 0) + and w_qb.dtype == torch.float8_e4m3fn + and tuple(w_qb.shape) == (NUM_HEADS * QK_HEAD_DIM, Q_LORA_RANK) + and w_qb.is_contiguous() + and w_qb.data_ptr() % 16 == 0 + and w_uk.dtype == torch.bfloat16 + and tuple(w_uk.shape) == (NUM_HEADS, KV_LORA_RANK, NOPE_DIM) + and w_uk.is_contiguous() + and w_uk.data_ptr() % 16 == 0 + ) + + +def input_supported(x: torch.Tensor) -> bool: + """Whether :func:`k3_mla_q_fused` reads ``x`` (BF16 ``[T, 1536]``, ``T <= MAX_TOKENS``).""" + return bool( + x.is_cuda + and x.dim() == 2 + and 0 < x.shape[0] <= MAX_TOKENS + and x.shape[1] == Q_LORA_RANK + and x.dtype == torch.bfloat16 + and x.stride(1) == 1 + and x.stride(0) % 8 == 0 + and x.data_ptr() % 16 == 0 + ) + + +def can_run(x: torch.Tensor, w_qb: torch.Tensor, w_uk: torch.Tensor) -> bool: + """Whether :func:`k3_mla_q_fused` serves ``x`` with these weights.""" + return input_supported(x) and weights_supported(w_qb, w_uk) + + +def k3_mla_q_fused( + x: torch.Tensor, + w_qb: torch.Tensor, + w_qb_exp: torch.Tensor, + w_uk: torch.Tensor, + q_global_scale: torch.Tensor, + cos_sin: torch.Tensor, + positions: torch.Tensor, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + *, + gen_len: int, + rows_per_seq: int, + use_pdl: bool = True, + debug: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + timestamps: Optional[torch.Tensor] = None, + norm: Optional[tuple[torch.Tensor, float]] = None, + pre_tiles: int = _NTILE, + early_trigger: bool = False, +) -> None: + """Write the FP4 decode query of ``x`` into ``q_fp4`` / ``q_sf`` (see the module docstring). + + Args: + x: BF16 ``[T, 1536]`` normalized q_a latent (row stride a multiple of 8 elements). + w_qb: q_b weight, FP8 ``[96 * 192, 1536]``; ``w_qb_exp``: its block exponents + (:func:`block_exponents`). + w_uk: BF16 absorption weight ``[96, 512, 128]``. + q_global_scale: FP32 ``[1]`` global scale of the FP4 query. + cos_sin: FP32 rotary table, 128 values ``(cos, sin)`` per position. + positions: int32 ``[>= T]`` position of every token (a negative position skips the + token's RoPE part, as the cache update does). + q_fp4, q_sf: The decode's Q buffers: uint8 ``[rows, 320]`` and the swizzled E4M3 scales. + Token ``t`` of request ``t // gen_len`` owns row + ``(t // gen_len) * rows_per_seq + (t % gen_len) * 96 + head``. + debug: Optional BF16 ``([T, 96 * 192], [T, 96, 576])`` outputs of the q_b projection and + of the absorbed / rotated query (tests). + norm: Optional ``(weight, eps)`` of the q_a RMSNorm (BF16 ``[1536]``): ``x`` is then the + raw q_a latent and the kernel normalizes it (``fuse_norm``). + pre_tiles: q_b tiles fetched before the PDL wait (all by default; microbenchmarks). + timestamps: Optional int64 ``[192, 16]``: ``%globaltimer`` at the phase boundaries of every + CTA (:class:`K3MlaQFusedKernel` ``timing``). + early_trigger: Release the dependent grid at entry (:class:`K3MlaQFusedKernel`). + """ + num_tokens = x.shape[0] + if num_tokens == 0: + return + if num_tokens > MAX_TOKENS: + raise ValueError(f"{num_tokens} tokens exceed the fused MLA query's {MAX_TOKENS}") + t_pad = 8 if num_tokens <= 8 else 16 + if positions.dtype != torch.int32 or positions.numel() < num_tokens: + raise ValueError("positions must be int32 covering every token") + if q_fp4.dtype != torch.uint8 or q_fp4.shape[-1] != Q_ROW_BYTES or not q_fp4.is_contiguous(): + raise ValueError("q_fp4 must be a contiguous uint8 [rows, 320] tensor") + if cos_sin.dtype != torch.float32 or not cos_sin.is_contiguous(): + raise ValueError("cos_sin must be a contiguous FP32 table") + dbg = (0, 0) if debug is None else (debug[0].data_ptr(), debug[1].data_ptr()) + if norm is not None and not ( + NORM_FUSION_AVAILABLE + and norm[0].dtype == torch.bfloat16 + and norm[0].numel() == Q_LORA_RANK + and norm[0].is_contiguous() + and norm[0].data_ptr() % 16 == 0 + ): + raise ValueError("the fused q_a RMSNorm needs FlashInfer and a contiguous BF16 weight") + key = ( + t_pad, + use_pdl, + debug is not None, + timestamps is not None, + norm is not None, + pre_tiles, + early_trigger, + ) + args = ( + Int64(x.data_ptr()), + make_ptr(cutlass.Float8E4M3FN, w_qb.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int64(w_qb_exp.data_ptr()), + make_ptr(cutlass.BFloat16, w_uk.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int64(q_global_scale.data_ptr()), + Int64(cos_sin.data_ptr()), + Int64(positions.data_ptr()), + Int64(q_fp4.data_ptr()), + Int64(q_sf.data_ptr()), + Int64(dbg[0]), + Int64(dbg[1]), + Int64(0 if timestamps is None else timestamps.data_ptr()), + Int64(0 if norm is None else norm[0].data_ptr()), + Float32(1e-6 if norm is None else norm[1]), + Int32(x.stride(0)), + Int32(num_tokens), + Int32(gen_len), + Int32(rows_per_seq), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + compiled = _COMPILED.get(key) + if compiled is None: + kernel = K3MlaQFusedKernel( + t_pad, + use_pdl, + debug=debug is not None, + timing=timestamps is not None, + fuse_norm=norm is not None, + pre_tiles=pre_tiles, + early_trigger=early_trigger, + ) + compiled = cute.compile(kernel, *args) + _COMPILED[key] = compiled + compiled(*args) + + +def latent_row_supported(c_kv: torch.Tensor, k_pe: torch.Tensor, weight: torch.Tensor) -> bool: + """Whether :func:`k3_mla_latent_row` serves these operands.""" + return bool( + NORM_FUSION_AVAILABLE + and c_kv.is_cuda + and c_kv.dim() == 2 + and k_pe.dim() == 2 + and c_kv.shape[1] == KV_LORA_RANK + and k_pe.shape == (c_kv.shape[0], ROPE_DIM) + and c_kv.dtype == k_pe.dtype == weight.dtype == torch.bfloat16 + and c_kv.stride(1) == 1 + and k_pe.stride(1) == 1 + and c_kv.stride(0) % 8 == 0 + and k_pe.stride(0) % 8 == 0 + and all(t.data_ptr() % 16 == 0 for t in (c_kv, k_pe, weight)) + and weight.numel() == KV_LORA_RANK + and weight.is_contiguous() + ) + + +def k3_mla_latent_row( + c_kv: torch.Tensor, + k_pe: torch.Tensor, + weight: torch.Tensor, + eps: float, + *, + use_pdl: bool = True, +) -> torch.Tensor: + """``torch.cat([rmsnorm(c_kv, weight, eps), k_pe], -1)`` (BF16 ``[T, 576]``) in one launch, + bitwise FlashInfer's CuTe DSL RMSNorm followed by the copy (see :class:`K3MlaLatentRowKernel`; + check :func:`latent_row_supported`).""" + num_tokens = c_kv.shape[0] + out = torch.empty((num_tokens, KV_LORA_RANK + ROPE_DIM), dtype=c_kv.dtype, device=c_kv.device) + if num_tokens == 0: + return out + args = ( + Int64(c_kv.data_ptr()), + Int64(k_pe.data_ptr()), + Int64(weight.data_ptr()), + Int64(out.data_ptr()), + Float32(eps), + Int32(c_kv.stride(0)), + Int32(k_pe.stride(0)), + Int32(out.stride(0)), + Int32(num_tokens), + Int32(_ceil_div(num_tokens, 4)), + cuda.CUstream(torch.cuda.current_stream(c_kv.device).cuda_stream), + ) + key = ("latent", use_pdl) + compiled = _COMPILED.get(key) + if compiled is None: + compiled = cute.compile(K3MlaLatentRowKernel(use_pdl), *args) + _COMPILED[key] = compiled + compiled(*args) + return out + + +__all__ = [ + "MAX_TOKENS", + "k3_mla_latent_row", + "latent_row_supported", + "input_supported", + "weights_supported", + "Q_ROW_BYTES", + "Q_SF_COLS", + "block_exponents", + "can_run", + "k3_mla_q_fused", + "smem_bytes", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_moe_tail1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_moe_tail1.py new file mode 100644 index 000000000000..76deb9929078 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_moe_tail1.py @@ -0,0 +1,982 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 MoE back half in one kernel for at most 8 rows (the column-parallel MoE tail). + +After the routed experts, the MoE block of a decode step needs, per token ``t``, +``y = latent_up(rmsnorm(sum_r p_r))`` where ``p_r`` is rank ``r``'s routed latent partial +``[M, 3584]`` (BF16), and the next layer's pre-attention mix adds ``y`` and the all-reduced shared +expert to the residual. :class:`K3MoeTail1Kernel` computes this rank's ``HIDDEN / W`` columns of +``y`` (rows ``[c0, c0 + HIDDEN / W)`` of latent_up) straight from the partials the routed experts +push, and multicasts them into the all-gather rows of the tail's Lamport call, which the next +pre-attention mix finishes (``k3_attn_res`` ``finish_ag``) together with the shared expert's slots: + +* Input. Each rank's BF16 latent partial rows go into slot ``(token, rank)`` of every rank's + current buffer of the latent Lamport workspace (one-shot layout ``[M][W][3584]``, sentinel + ``0x80000000`` per 32-bit word, ``-0`` sent as ``+0``), stored by the routed experts + (an ``output_push`` of their combine); the kernel polls those slots and + waits for the experts' grid only at its end. +* Grid. ``G`` clusters of 16 CTAs x 256 threads. Cluster ``g`` owns ``RG = 112`` output rows + (:data:`ROWS_PER_CLUSTER`, ``G = (HIDDEN / W) / RG``: 4 clusters at 16 ranks); CTA ``q`` of the + cluster owns latent columns ``[224 q, 224 q + 224)``. Before the poll each lane loads its FP8 + latent_up fragments (the 128x128 UE8M0 block scales applied, BF16 exact) into registers. +* Poll and sum. Warp ``t`` polls token ``t``'s 224 columns of all ``W`` slots (lane ``l < 28``: + 16 bytes = 8 columns per slot; every lane reloads all slots each pass and the warp leaves the + loop by a vote). The slots are summed in FP32 in rank order, in chunks of eight ranks (the C++ + one-shot all-reduce's order) and rounded to BF16 (bitwise the latent all-reduce's sum). Each lane + adds its eight squares (FP32) and stores ``bf16(x * gamma)`` into shared memory. +* RMSNorm late. A token's partial sum of squares goes to its two owner CTAs (``st.async``); the up + projection runs on ``bf16(x * gamma)`` and the owners scale the summed products by + ``rsqrt(sum / 3584 + eps)`` before rounding, so the squares exchange overlaps the MMAs. +* Up projection. ``mma.sync`` m16n8k16 BF16 with FP32 accumulation; warp ``w`` takes the 16-row + m-tiles ``w, w + 8, ..`` over the CTA's 224 columns (14 K steps). The MMA's K index is permuted + (:func:`_k_col`): lane group ``t``'s four columns of a K step are contiguous and those of four + consecutive steps are one 16-byte FP8 load per weight row (steps 12-13: 8 bytes); its B fragment + of a step is the 8 bytes of the token row at the same columns. Every weight load of a lane is + issued before any conversion, every polling pass issues all its loads before any check. +* Owners. Output line (token ``t``, row half ``h``) is owned by CTA ``t + 8 h`` of the cluster: + every m-tile's 4-row partials go there by ``st.async``; the owner sums the 16 CTAs' partials in + CTA order, scales, rounds to BF16, replaces ``-0`` by ``+0`` and multicasts the row segment into + the all-gather rows of the tail call. +* Lamport state of the latent workspace (the role of the latent all-reduce): the kernel reads the + rotation state at entry (final: the previous call completed before the experts started) and + counts every CTA there; it re-arms the previous call's buffer at entry. The CTA that counts last + rotates the buffers and records this call's ``M x W`` slots as the next call's bytes to clear, + after its ``griddepcontrol.wait`` (the experts, which read the state, have completed). The tail + call's state is only read (``cur``, bytes per buffer); the pre-attention mix advances it. + +PDL: the dependent (the pre-attention mix) is released at entry. Weights, gamma, the block +exponents and both rotation states are read before any wait; the kernel stores nothing a +predecessor reads. + +Numerics (not bitwise the column-parallel tail): the latent sum is the all-reduce's; the RMS uses +FP32 squares of it; the projection's activation is ``bf16(x * gamma)`` in BF16 instead of the +MXFP8 ``bf16((x * r) * gamma)``, and ``r`` multiplies the FP32 product sum before its rounding. +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from .k3_lamport_ar import ( + _asm, + _fv, + _iv, + _lv, + _unpack, + acc_bf16x2, + add_f32, + atom_add_acq_rel_u32, + clean4, + clear_dirty_stages, + griddepcontrol_launch_dependents, + griddepcontrol_wait, + ld_volatile_v4, + multimem_st_v4, + pack_bf16x2, + sanitize_bf16x2, + st_u32, + st_v4, +) + +__all__ = [ + "HIDDEN", + "LATENT", + "MAX_M", + "ROWS_PER_CLUSTER", + "K3MoeTail1Kernel", + "can_run", + "rank_rows", + "run_k3_moe_tail1", +] + +#: Latent (latent_up's K) and hidden (latent_up's N) widths. +LATENT = 3584 +HIDDEN = 7168 +#: CTAs per cluster; each owns ``LATENT / CLUSTER`` latent columns. +CLUSTER = 16 +KS = LATENT // CLUSTER +#: 16-byte column chunks of a token's K slice (one polling lane each). +CHUNKS = KS // 8 +#: Threads per CTA (warp ``t`` polls token ``t``). +THREADS = 256 +WARPS = THREADS // 32 +#: Largest row count (one warp per token). +MAX_M = 8 +#: MMA K steps (16) over a CTA's K slice. +K_STEPS = KS // 16 +#: 16-row m-tiles a warp may hold in registers (56 registers each). +MAX_MTILES_PER_WARP = 2 +#: Output rows of one 16-CTA cluster (four clusters at 16 ranks). +ROWS_PER_CLUSTER = 112 + + +def _k_col(j: int): + """``(base, lane stride)`` of K step ``j``'s first column for lane group ``t`` (column + ``base + stride * t``, four consecutive columns per step): steps ``4 i .. 4 i + 3`` (i < 3) of + lane group ``t`` are the 16 columns ``64 i + 16 t ..``, steps 12-13 the 8 columns + ``192 + 8 t ..``; a bijection of the 224 columns onto (step, lane group, element).""" + if j < 12: + return 64 * (j // 4) + 4 * (j % 4), 16 + return 192 + 4 * (j - 12), 8 + + +#: Bytes of a token row of ``bf16(x * gamma)`` in shared memory (16 bytes of padding). +XS_PITCH = KS * 2 + 16 +#: FP32 words per token row of a warp's staged MMA partial (16 rows + padding: conflict-free). +STAGE_PITCH = 20 +STAGE_WARP_BYTES = MAX_M * STAGE_PITCH * 4 +#: Lamport buffers of a workspace and the word of its CTA arrival counter. +NUM_BUFFERS = 3 +_FLAG_ARRIVALS = 8 +#: Ranks per FP32 partial sum of the latent (the C++ one-shot all-reduce's chunking). +RANK_CHUNK = 8 + +assert KS % 8 == 0 and KS == 224 and 56 * 4 == KS + + +@dsl_user_op +def _ld_nc_v2(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.v2.b32 {$0, $1}, [$2];", + "=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _ld_nc_u8(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.u8 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) # fmt: skip + ) + + +@dsl_user_op +def _ld_nc_v4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _deq_bf16x2(word, scale, high: bool, *, loc=None, ip=None): + """Two E4M3 codes (the low or high 16 bits of ``word``, low byte first) times ``scale`` (a power + of two) as a BF16 pair, low value in the low half. Exact for latent_up's weights: E4M3 -> F16 -> + F32 are exact and the power-of-two scale keeps the 4 significant bits.""" + half = "h" if high else "l" + return Int32( + _asm( + T.i32(), + [_iv(word, loc=loc, ip=ip), _fv(scale, loc=loc, ip=ip)], + "{\n.reg .b16 l, h, p0, p1;\n.reg .b32 p;\n.reg .f32 f0, f1;\n" + "mov.b32 {l, h}, $1;\n" + f"cvt.rn.f16x2.e4m3x2 p, {half};\n" + "mov.b32 {p0, p1}, p;\n" + "cvt.f32.f16 f0, p0;\ncvt.f32.f16 f1, p1;\n" + "mul.rn.f32 f0, f0, $2;\nmul.rn.f32 f1, f1, $2;\n" + "cvt.rn.bf16x2.f32 $0, f1, f0;\n}", + "=r,r,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _exp2_of_ue8m0(e, *, loc=None, ip=None): + """``2^(e - 127)`` of a UE8M0 exponent ``1 <= e <= 254`` (the FP32 word ``e << 23``).""" + return Float32( + _asm( + T.f32(), + [_iv(e, loc=loc, ip=ip)], + "{\n.reg .b32 b;\nshl.b32 b, $1, 23;\nmov.b32 $0, b;\n}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bf16x2_f32(word, *, loc=None, ip=None): + """(low, high) BF16 of a pair as FP32.""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_iv(word, loc=loc, ip=ip)], + "{\n.reg .b32 lo, hi;\nshl.b32 lo, $2, 16;\nand.b32 hi, $2, 0xffff0000;\n" + "mov.b32 $0, lo;\nmov.b32 $1, hi;\n}", + "=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _fma(a, b, c, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip), _fv(c, loc=loc, ip=ip)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fmul(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _rms_scale(total, eps, *, loc=None, ip=None): + """``rsqrt(total / 3584 + eps)`` (IEEE division, then add, approximate reciprocal root).""" + return Float32( + _asm( + T.f32(), + [_fv(total, loc=loc, ip=ip), _fv(eps, loc=loc, ip=ip)], + "{\n.reg .f32 q;\ndiv.rn.f32 q, $1, 0f45600000;\nadd.rn.f32 q, q, $2;\n" + "rsqrt.approx.f32 $0, q;\n}", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_bfly_f32(value, offset: int, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(value, loc=loc, ip=ip)], + f"shfl.sync.bfly.b32 $0, $1, {offset}, 31, -1;", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _mma_bf16(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """``d[16x8] = a[16x16] @ b[16x8] + c`` (BF16 inputs, FP32 accumulation).""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(x, loc=loc, ip=ip) for x in (a0, a1, a2, a3, b0, b1)] + + [_fv(x, loc=loc, ip=ip) for x in (c0, c1, c2, c3)], + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _lds_v2(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v2.b32 {$0, $1}, [$2];", + "=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 2, loc=loc, ip=ip) + + +@dsl_user_op +def _lds_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _lds_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.f32 $0, [$1];", + "=f,r", + loc=loc, + ip=ip, + ) # fmt: skip + ) + + +@dsl_user_op +def _lds_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.b32 $0, [$1];", + "=r,r", + loc=loc, + ip=ip, + ) # fmt: skip + ) + + +@dsl_user_op +def _sts_b32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_v4(addr, w0, w1, w2, w3, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)] + [_iv(w, loc=loc, ip=ip) for w in (w0, w1, w2, w3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mbar_init(addr, count: int, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)], + f"mbarrier.init.shared::cta.b64 [$0], {count};", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mbar_expect_tx(addr, nbytes: int, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)], + f"mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], {nbytes};", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_mbar_init(*, loc=None, ip=None): + _asm(None, [], "fence.mbarrier_init.release.cluster;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _mbar_try_wait_cluster(addr, *, loc=None, ip=None): + """1 once phase 0 of the local mbarrier at ``addr`` (completed by peers' ``st.async``) has + completed, else 0 (cluster-scope acquire).""" + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], 0;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_async_f32(local_addr, peer, local_mbar, value, *, loc=None, ip=None): + """One FP32 into CTA ``peer``'s copy of ``local_addr``, counted by its copy of + ``local_mbar``.""" + _asm( + None, + [ + _iv(local_addr, loc=loc, ip=ip), + _iv(peer, loc=loc, ip=ip), + _iv(local_mbar, loc=loc, ip=ip), + _fv(value, loc=loc, ip=ip), + ], + "{\n.reg .b32 ra, rm;\nmapa.shared::cluster.u32 ra, $0, $1;\n" + "mapa.shared::cluster.u32 rm, $2, $1;\n" + "st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [ra], $3, [rm];\n}", + "r,r,r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_async_v4_f32(local_addr, peer, local_mbar, v0, v1, v2, v3, *, loc=None, ip=None): + """16 bytes into CTA ``peer``'s copy of ``local_addr``, counted by its copy of + ``local_mbar``.""" + _asm( + None, + [_iv(x, loc=loc, ip=ip) for x in (local_addr, peer, local_mbar)] + + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n.reg .b32 ra, rm;\nmapa.shared::cluster.u32 ra, $0, $1;\n" + "mapa.shared::cluster.u32 rm, $2, $1;\n" + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 " + "[ra], {$3, $4, $5, $6}, [rm];\n}", + "r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_g2s(dst, src, nbytes: int, mbar, *, loc=None, ip=None): + """Bulk copy of ``nbytes`` from global ``src`` to shared ``dst``, completing on ``mbar``.""" + _asm( + None, + [_iv(dst, loc=loc, ip=ip), _lv(src, loc=loc, ip=ip), _iv(mbar, loc=loc, ip=ip)], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes " + f"[$0], [$1], {nbytes}, [$2];", + "r,l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_v4(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _cluster_arrive_relaxed(*, loc=None, ip=None): + _asm(None, [], "barrier.cluster.arrive.relaxed.aligned;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _cluster_wait(*, loc=None, ip=None): + _asm(None, [], "barrier.cluster.wait.aligned;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _bar_sync(barrier_id: int, count: int, *, loc=None, ip=None): + _asm(None, [], f"bar.sync {barrier_id}, {count};", "", loc=loc, ip=ip) + + +class K3MoeTail1Kernel: + """The MoE back half of ``m <= 8`` rows on one rank of ``world`` (see the module docstring). + + :param world: Ranks of the latent / tail workspaces (``HIDDEN / world`` output rows per rank, + in ``HIDDEN / world / ROWS_PER_CLUSTER`` clusters; warp ``w`` multiplies the 16-row m-tiles + ``w, w + 8, ..`` of its cluster's rows). + :param m: Tokens of the call. + :param mc_out: Store the output through the multicast mapping (``multimem.st``); False stores + to the given address (single-GPU tests). + """ + + def __init__(self, world: int, m: int, mc_out: bool = True): + rows = HIDDEN // world + if HIDDEN % world or rows % ROWS_PER_CLUSTER: + raise ValueError( + f"{rows} rows per rank do not split into {ROWS_PER_CLUSTER}-row clusters" + ) + if not 1 <= m <= MAX_M: + raise ValueError(f"m={m} rows: 1..{MAX_M}") + self.world = world + self.m = m + self.rg = ROWS_PER_CLUSTER + self.groups = rows // ROWS_PER_CLUSTER + self.nctas = self.groups * CLUSTER + self.mtiles = ROWS_PER_CLUSTER // 16 + self.mpw = (self.mtiles + WARPS - 1) // WARPS + assert self.mpw <= MAX_MTILES_PER_WARP + # Two owners per token (row halves of the cluster's rows), 8-row lines per lane. + self.ro = ROWS_PER_CLUSTER // 2 + self.mc_out = mc_out + self.part_words = CLUSTER * self.ro + self.rank_chunk = min(RANK_CHUNK, world) + + @cute.jit + def __call__( + self, + lat_uc: Int64, + lat_flags: Int64, + w_addr: Int64, + wexp_addr: Int64, + gamma_addr: Int64, + ag_base: Int64, + ag_flags: Int64, + c0: Int32, + eps: Float32, + stream: cuda.CUstream, + ): + @cute.struct + class SharedStorage: + bars: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 2], 16] + sq: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, CLUSTER], 16] + xs: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, MAX_M * XS_PITCH], 16] + stage: cute.struct.Align[ + cute.struct.MemRange[cutlass.Uint8, WARPS * STAGE_WARP_BYTES], 16 + ] + part: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, self.part_words], 16] + + self.shared_storage = SharedStorage + self.kernel( + lat_uc, + lat_flags, + w_addr, + wexp_addr, + gamma_addr, + ag_base, + ag_flags, + c0, + eps, + ).launch( # fmt: skip + grid=(self.nctas, 1, 1), + block=(THREADS, 1, 1), + cluster=(CLUSTER, 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=True, + ) + + def rotate(self, lat_flags, cur, nbytes): + """The latent workspace's next state: buffer ``cur`` becomes the dirty one with this call's + ``m x W`` slots to re-arm; the arrival counter back at zero (one thread).""" + nxt = cutlass.select_(cur + 1 == NUM_BUFFERS, Int32(0), cur + 1) + st_v4(lat_flags, nxt, cur, nbytes, Int32(1)) + call_bytes = self.m * self.world * LATENT * 2 + st_v4(lat_flags + Int64(16), Int32(call_bytes), Int32(0), Int32(0), Int32(0)) + st_u32(lat_flags + Int64(_FLAG_ARRIVALS * 4), Int32(0)) + + @cute.kernel + def kernel( + self, + lat_uc: Int64, + lat_flags: Int64, + w_addr: Int64, + wexp_addr: Int64, + gamma_addr: Int64, + ag_base: Int64, + ag_flags: Int64, + c0: Int32, + eps: Float32, + ): + W = self.world + m = self.m + KJ = K_STEPS + RO = self.ro + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + q = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + gcl = bidx // CLUSTER + gq = lane // 4 + tq = lane % 4 + + griddepcontrol_launch_dependents() + + storage = utils.SmemAllocator().allocate(self.shared_storage) + bars = storage.bars.data_ptr().toint() + bar_sq = bars + bar_part = bars + 8 + sq_s = storage.sq.data_ptr().toint() + xs_s = storage.xs.data_ptr().toint() + stage_s = storage.stage.data_ptr().toint() + part_s = storage.part.data_ptr().toint() + k0 = q * KS + rbase = c0 + gcl * self.rg + + # ---- latent_up loads first (before any wait; their conversions follow the setup) ---- + # Per m-tile and weight row: three 16-byte loads (K steps 0-11) and one 8-byte load (steps + # 12-13) of FP8 codes, and the UE8M0 exponents of the K slice's (up to) three 128-column + # scale blocks. + b0 = k0 // 128 + raw = [] + exps = [] + for s in cutlass.range_constexpr(self.mpw): + mt = cutlass.min(warp + WARPS * s, self.mtiles - 1) + r0 = rbase + mt * 16 + gq + for h in cutlass.range_constexpr(2): + w_row = w_addr + Int64(r0 + 8 * h) * LATENT + Int64(k0) + for i in cutlass.range_constexpr(3): + raw += list(_ld_nc_v4(w_row + Int64(64 * i + 16 * tq))) + raw += list(_ld_nc_v2(w_row + Int64(192 + 8 * tq))) + e_row = wexp_addr + Int64(((r0 + 8 * h) // 128) * (LATENT // 128)) + for bi in cutlass.range_constexpr(3): + blk = cutlass.min(b0 + bi, LATENT // 128 - 1) + exps.append(_ld_nc_u8(e_row + Int64(blk))) + lat_lane = (warp < m) & (lane < CHUNKS) + gcol = k0 + cutlass.select_(lane < CHUNKS, lane, Int32(0)) * 8 + gw = _ld_nc_v4(gamma_addr + Int64(gcol) * 2) + + # ---- setup: barriers, rotation states (every thread reads them), the arrival count ---- + # Owner role: line (token q % 8, row half q // 8) of the cluster's rows. + own_t = q % MAX_M + own_h = q // MAX_M + # Rotation states: final, the previous calls on both workspaces completed before the + # experts (whose launch released this grid) passed their own wait. + f = ld_volatile_v4(lat_flags) + fc = ld_volatile_v4(lat_flags + Int64(16)) + a = ld_volatile_v4(ag_flags) + cur, dirty, nbytes, dstages = f[0], f[1], f[2], f[3] + ag_cur, ag_bytes = a[0], a[2] + if tidx == 0: + _mbar_init(bar_sq, 1) + _mbar_init(bar_part, 1) + if own_t < m: + _mbar_expect_tx(bar_sq, CLUSTER * 4) + _mbar_expect_tx(bar_part, CLUSTER * RO * 4) + _fence_mbar_init() + # Every thread of the CTA has read the state before its count. + _bar_sync(1, THREADS) + # The CTAs of a cluster start together; the wait (every peer's barriers initialized) comes + # before the first remote store. + _cluster_arrive_relaxed() + arr_old = Int32(0) + if tidx == 0: + # Count this CTA: the last one rotates the buffers once the experts (which read the + # state) completed. + arr_old = atom_add_acq_rel_u32(lat_flags + Int64(_FLAG_ARRIVALS * 4), Int32(1)) + + # ---- latent_up fragments: BF16 of the FP8 codes times the block scale ---- + afr = cute.make_rmem_tensor((self.mpw * KJ * 4,), Int32) + scales = [_exp2_of_ue8m0(e) for e in exps] + for s in cutlass.range_constexpr(self.mpw): + for j in cutlass.range_constexpr(KJ): + cb, cs = _k_col(j) + bi = (k0 + cb + cs * tq) // 128 - b0 + sc = [] + for h in cutlass.range_constexpr(2): + e3 = (s * 2 + h) * 3 + hi = cutlass.select_(bi == 1, scales[e3 + 1], scales[e3 + 2]) + sc.append(cutlass.select_(bi == 0, scales[e3], hi)) + w0 = raw[(s * 2) * KJ + j] + w1 = raw[(s * 2 + 1) * KJ + j] + base = (s * KJ + j) * 4 + afr[base + 0] = _deq_bf16x2(w0, sc[0], False) + afr[base + 1] = _deq_bf16x2(w1, sc[1], False) + afr[base + 2] = _deq_bf16x2(w0, sc[0], True) + afr[base + 3] = _deq_bf16x2(w1, sc[1], True) + + # The previous call's buffer (complete before this grid launched; nobody reads it). + clear_dirty_stages( + lat_uc, dirty, nbytes, dstages, fc[0], fc[1], fc[2], fc[3], + bidx * THREADS + tidx, self.nctas * THREADS, + ) # fmt: skip + + # ---- poll token `warp`'s slots of every rank (lanes < 28: 8 columns each) ---- + words = cute.make_rmem_tensor((W * 4,), Int32) + row0 = lat_uc + Int64(cur) * Int64(nbytes) + Int64((warp * W) * LATENT + gcol) * 2 + pending = Int32(0) + for e in cutlass.range_constexpr(W * 4): + words[e] = Int32(0) + # Every pass issues all W loads before it checks any (they are in flight together). + if lat_lane: + loaded = [] + for s in cutlass.range_constexpr(W): + loaded.append(list(ld_volatile_v4(row0 + Int64(s * LATENT * 2)))) + for s in cutlass.range_constexpr(W): + for e in cutlass.range_constexpr(4): + words[s * 4 + e] = loaded[s][e] + clean = clean4(loaded[s][0], loaded[s][1], loaded[s][2], loaded[s][3]) + pending = pending | ((Int32(1) - clean) << s) + while cute.arch.vote_any_sync(pending != 0): + if pending != 0: + loaded = [] + for s in cutlass.range_constexpr(W): + loaded.append(list(ld_volatile_v4(row0 + Int64(s * LATENT * 2)))) + pending = Int32(0) + for s in cutlass.range_constexpr(W): + for e in cutlass.range_constexpr(4): + words[s * 4 + e] = loaded[s][e] + clean = clean4(loaded[s][0], loaded[s][1], loaded[s][2], loaded[s][3]) + pending = pending | ((Int32(1) - clean) << s) + + # Rank-order sum (chunks of RANK_CHUNK ranks), BF16; squares; bf16(x * gamma) to smem. + acc = [Float32(0.0)] * 8 + for cb in cutlass.range_constexpr(0, W, self.rank_chunk): + part = [Float32(0.0)] * 8 + for s in cutlass.range_constexpr(cb, cb + self.rank_chunk): + for j in cutlass.range_constexpr(4): + part[2 * j], part[2 * j + 1] = acc_bf16x2( + words[s * 4 + j], part[2 * j], part[2 * j + 1] + ) + for i in cutlass.range_constexpr(8): + acc[i] = add_f32(acc[i], part[i]) + s2 = Float32(0.0) + xg = [] + for j in cutlass.range_constexpr(4): + red = pack_bf16x2(acc[2 * j], acc[2 * j + 1]) + x_lo, x_hi = _bf16x2_f32(red) + g_lo, g_hi = _bf16x2_f32(gw[j]) + s2 = _fma(x_lo, x_lo, s2) + s2 = _fma(x_hi, x_hi, s2) + xg.append(pack_bf16x2(_fmul(x_lo, g_lo), _fmul(x_hi, g_hi))) + if lane < CHUNKS: + if warp < m: + _sts_v4(xs_s + warp * XS_PITCH + lane * 16, xg[0], xg[1], xg[2], xg[3]) + else: + _sts_v4(xs_s + warp * XS_PITCH + lane * 16, Int32(0), Int32(0), Int32(0), Int32(0)) + s2 = cutlass.select_(lat_lane, s2, Float32(0.0)) + for off in (16, 8, 4, 2, 1): + s2 = s2 + _shfl_bfly_f32(s2, off) + _cluster_wait() + if warp < m: + if lane < 2: + # This CTA's sum of squares of token `warp` to its owners (row halves 0 and 1). + _st_async_f32(sq_s + q * 4, warp + lane * MAX_M, bar_sq, s2) + _bar_sync(1, THREADS) + + # ---- up projection: m-tiles w, w + 8, ..; 4-row partials to the line owners ---- + xrow = xs_s + gq * XS_PITCH + bfr = [] + for j in cutlass.range_constexpr(KJ): + cb, cs = _k_col(j) + bfr += list(_lds_v2(xrow + (cb + cs * tq) * 2)) + stg = stage_s + warp * STAGE_WARP_BYTES + for s in cutlass.range_constexpr(self.mpw): + mt = warp + WARPS * s + if mt < self.mtiles: + # Two accumulator chains (even / odd K steps), added at the end. + acc2 = [[Float32(0.0)] * 4, [Float32(0.0)] * 4] + for j in cutlass.range_constexpr(KJ): + base = (s * KJ + j) * 4 + ch = acc2[j % 2] + acc2[j % 2] = list(_mma_bf16( + afr[base], afr[base + 1], afr[base + 2], afr[base + 3], + bfr[2 * j], bfr[2 * j + 1], ch[0], ch[1], ch[2], ch[3], + )) # fmt: skip + d0 = add_f32(acc2[0][0], acc2[1][0]) + d1 = add_f32(acc2[0][1], acc2[1][1]) + d2 = add_f32(acc2[0][2], acc2[1][2]) + d3 = add_f32(acc2[0][3], acc2[1][3]) + # d0 (row gq, token 2 tq), d1 (row gq, token 2 tq + 1), d2 / d3 (row gq + 8). + _sts_f32(stg + ((2 * tq) * STAGE_PITCH + gq) * 4, d0) + _sts_f32(stg + ((2 * tq + 1) * STAGE_PITCH + gq) * 4, d1) + _sts_f32(stg + ((2 * tq) * STAGE_PITCH + gq + 8) * 4, d2) + _sts_f32(stg + ((2 * tq + 1) * STAGE_PITCH + gq + 8) * 4, d3) + cute.arch.sync_warp() + tk = lane // 4 + r4 = (lane % 4) * 4 + if tk < m: + v0, v1, v2, v3 = _lds_v4_f32(stg + (tk * STAGE_PITCH + r4) * 4) + rloc = mt * 16 + r4 + h = rloc // RO + dst = part_s + (q * RO + rloc - h * RO) * 4 + _st_async_v4_f32(dst, tk + h * MAX_M, bar_part, v0, v1, v2, v3) + cute.arch.sync_warp() + + # ---- owner: sum the CTAs' partials, scale by the token's RMS reciprocal, multicast ---- + if own_t < m: + if warp == 0: + done = _mbar_try_wait_cluster(bar_sq) + while done == 0: + done = _mbar_try_wait_cluster(bar_sq) + done = _mbar_try_wait_cluster(bar_part) + while done == 0: + done = _mbar_try_wait_cluster(bar_part) + if lane < RO // 8: + total = Float32(0.0) + for r in cutlass.range_constexpr(CLUSTER): + total = add_f32(total, _lds_f32(sq_s + r * 4)) + inv = _rms_scale(total, eps) + parts = [] + for sl in cutlass.range_constexpr(CLUSTER): + parts.append(_lds_v4_f32(part_s + (sl * RO + lane * 8) * 4)) + parts.append(_lds_v4_f32(part_s + (sl * RO + lane * 8 + 4) * 4)) + y = [Float32(0.0)] * 8 + for sl in cutlass.range_constexpr(CLUSTER): + for i in cutlass.range_constexpr(4): + y[i] = add_f32(y[i], parts[2 * sl][i]) + y[4 + i] = add_f32(y[4 + i], parts[2 * sl + 1][i]) + ow = [ + sanitize_bf16x2(pack_bf16x2(_fmul(y[2 * j], inv), _fmul(y[2 * j + 1], inv))) + for j in range(4) + ] + col = c0 + gcl * self.rg + own_h * RO + lane * 8 + dst = ( + ag_base + + Int64(ag_cur) * Int64(ag_bytes) + + (Int64((m * W + own_t) * HIDDEN) + Int64(col)) * 2 + ) + if cutlass.const_expr(self.mc_out): + multimem_st_v4(dst, ow[0], ow[1], ow[2], ow[3]) + else: + st_v4(dst, ow[0], ow[1], ow[2], ow[3]) + + # ---- rotate once the experts (which read the state) completed ---- + griddepcontrol_wait() + if tidx == 0: + if arr_old == self.nctas - 1: + self.rotate(lat_flags, cur, nbytes) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def rank_rows(world: int, rank: int): + """``(first row, rows)`` of latent_up's output this rank computes: ``HIDDEN / world`` each.""" + n = HIDDEN // world + return rank * n, n + + +def can_run(m: int, world: int) -> bool: + """Whether the kernel serves ``m`` rows on ``world`` ranks.""" + return 1 <= m <= MAX_M and HIDDEN % world == 0 and (HIDDEN // world) % ROWS_PER_CLUSTER == 0 + + +def _set_kernel_attributes(compiled, device_index: int) -> None: + """Clusters of 16 (non-portable size) and the maximum shared-memory carveout (the next + kernels' CTAs may land beside this grid's) on every kernel of ``compiled``.""" + from . import k3_launch_attrs as attrs + + status, cu_device = cuda.cuDeviceGet(device_index) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuDeviceGet: {status}") + attr = cuda.CUfunction_attribute + for _, kernel in attrs._kernels(compiled, device_index): + for a, v in ( + (attr.CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1), + (attr.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, attrs.MAX_SHARED_CARVEOUT), + ): + (status,) = cuda.cuKernelSetAttribute(a, v, kernel, cu_device) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuKernelSetAttribute({a}): {status}") + + +@functools.cache +def _compiled(world: int, m: int, mc_out: bool, device_index: int): + kernel = K3MoeTail1Kernel(world, m, mc_out) + z = Int64(0) + with torch.cuda.device(device_index): + compiled = cute.compile( + kernel, z, z, z, z, z, z, z, Int32(0), Float32(1e-6), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) # fmt: skip + _set_kernel_attributes(compiled, device_index) + return compiled + + +def run_k3_moe_tail1( + *, + m: int, + world: int, + lat_uc: int, + lat_flags: torch.Tensor, + weight: torch.Tensor, + wexp: torch.Tensor, + gamma: torch.Tensor, + eps: float, + ag_base: int, + ag_flags: torch.Tensor, + c0: int, + mc_out: bool = True, +) -> None: + """Launch :class:`K3MoeTail1Kernel` on the current stream. + + Args: + m: Tokens (1..8). + world: Ranks of both workspaces. + lat_uc / lat_flags: The latent Lamport workspace (local unicast base, uint32 [9] rotation + state) the routed experts push their partials into. + weight: latent_up ``[HIDDEN, LATENT]`` FP8 E4M3 (all rows; the kernel reads rows + ``[c0, c0 + HIDDEN / world)``). + wexp: latent_up's UE8M0 block exponents ``uint8 [HIDDEN / 128, LATENT / 128]`` (1..254). + gamma: The routed latent RMSNorm weight ``[LATENT]`` BF16. + ag_base / ag_flags: The tail call's workspace (multicast base, or with ``mc_out=False`` + a local base) and its rotation state; the kernel writes the call's all-gather rows. + c0: This rank's first output row (:func:`rank_rows`). + """ + if not can_run(m, world): + raise ValueError(f"k3_moe_tail1 does not serve m={m} world={world}") + if weight.dtype != torch.float8_e4m3fn or tuple(weight.shape) != (HIDDEN, LATENT): + raise ValueError("weight must be the [7168, 3584] FP8 E4M3 latent_up") + if wexp.dtype != torch.uint8 or tuple(wexp.shape) != (HIDDEN // 128, LATENT // 128): + raise ValueError("wexp must be uint8 [56, 28]") + if gamma.dtype != torch.bfloat16 or gamma.numel() != LATENT: + raise ValueError("gamma must be BF16 [3584]") + device = weight.device + index = device.index if device.index is not None else torch.cuda.current_device() + _compiled(world, m, mc_out, index)( + Int64(lat_uc), + Int64(lat_flags.data_ptr()), + Int64(weight.data_ptr()), + Int64(wexp.data_ptr()), + Int64(gamma.data_ptr()), + Int64(ag_base), + Int64(ag_flags.data_ptr()), + Int32(c0), + Float32(eps), + cuda.CUstream(torch.cuda.current_stream(device).cuda_stream), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_route_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_route_tail.py new file mode 100644 index 000000000000..d7d2abf9877e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_route_tail.py @@ -0,0 +1,2268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 MoE routing for decode batches (SM100/SM103/SM107): top-16 + ``moe_sort`` in one launch. + +``k3_route_tail`` replaces ``noaux_tc_op`` -> ``moe_sort`` (and, given split-K partial logits, +the split-K reduction before them) by one launch on a cluster of eight small CTAs, so the +routing occupies few SMs next to the side-stream kernels it runs with. + +* Input: FP32 router logits ``[M, 896]`` (``1 <= M <= 2048``), or ``S`` split-K partials + ``[S, R, 896]`` (``R >= M``) that are summed in split order. +* Top-k, bit-identical to ``noaux_tc_op`` on the same logits (selection, order, weights). A + team of ``G`` warps routes one token, lane ``g`` of the team owning experts + ``j * 32 G + g``. Each lane forms the order-preserving bits of ``sigmoid(logit) + bias`` (the + high word of ``noaux_tc_op``'s 64-bit key, whose low word ``65535 - expert`` breaks ties + toward smaller experts). In each warp the 16th largest lane maximum is a lower bound of the + token's 16th largest score; the team keeps the largest of these bounds. Every expert at or + above it (usually about 20 per token) is a candidate: warps compact theirs into shared + memory and rank them by their full keys. If a warp has more than 32 (64 for one-warp teams) + candidates (heavy ties), its team's first warp instead runs 16 exact arg-max rounds over all + 896 keys. The weights are the ``sigmoid`` scores renormalized like ``noaux_tc_op``. +* Sort (``moe_sort`` metadata of one expert-parallel rank): CTA ``c`` of the cluster owns a + contiguous block of at most 32 tokens and marks, per local expert, a token bitmask in shared + memory, then stores it into every CTA of the cluster (distributed shared memory). After a + cluster barrier each CTA scans the per-expert tile counts (``ceil(count / 128)`` tiles in + local-expert order) and writes the rows of its own tokens: rows of an expert follow the token + order, which is what ``moe_sort``'s single-CTA kernels produce (its cluster kernel orders rows + by atomics). CTA 0 writes the tile arrays. +* ``256 < M <= 2048`` (:class:`K3RouteWideKernel`): independent CTAs of 8-32 tokens run the + same top-k; each publishes per-expert token masks (and the top-k slot of each masked token) + to a persistent workspace, and the last CTA to arrive (arrival counter) writes the same + ``moe_sort`` outputs, rows again in ascending token order. + +The cluster kernels keep no state between launches; the multi-CTA kernel resets its arrival +counter before it exits. Either way a captured CUDA graph replays them as is. +""" + +from __future__ import annotations + +import functools +import struct + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from .k3_fused_router import ( + NUM_EXPERTS, + TOP_K, + _asm, + _atom_add_relaxed_gpu, + _div_rn_f32, + _fence_acq_rel_gpu, + _fv, + _iv, + _lds_b32, + _lds_f32, + _lv, + _popc, + _red_shared_or, + _renormalize, + _st_global_b32, + _st_global_f32, + _st_relaxed_b32, + _sts_b32, + _sts_v4_f32, + _unit_scale_fast_path, + _unpack, + _warp_inclusive_sum, + _warp_sum_f32, + sort_buffer_sizes, +) + +#: Largest token count served by the tail. +MAX_M = 256 +#: The cluster kernel picks the candidates on ``tanh.approx`` scores and evaluates ``noaux_tc``'s +#: exact sigmoid only for them (the accurate ``tanhf`` of all 896 experts dominates the top-k). +#: The selection stays bit-identical: see :meth:`K3RouteTailKernel.select_token`. +#: Score margin of that approximate candidate pass, at least twice the largest difference between +#: an approximate and an exact score: ``tanh.approx.f32`` is within 2^-10.987 of tanh, i.e. the +#: sigmoid within 2.5e-4, plus a few FP32 roundings of values below 2. +APPROX_MARGIN = 6.0e-4 +#: Largest cluster (portable cluster size). +MAX_CTAS = 8 +#: Tokens per CTA are capped by the 32-bit token mask of the sort. +_MAX_TOKENS_PER_CTA = 32 +#: Expert slots per lane: 7 blocks of 128 experts x 4 consecutive experts. +_BLOCKS = NUM_EXPERTS // 128 +_SLOTS = _BLOCKS * 4 +#: Bytes of one compacted candidate: ``{65535 - expert, key high word, sigmoid bits, 0}``. +_CAND_BYTES = 16 +#: Tile size of the CuteDSL MoE (``tile_tokens_dim``). +TILE = 128 + +assert NUM_EXPERTS % 128 == 0 + + +# --------------------------------------------------------------------------- +# Inline PTX helpers +# --------------------------------------------------------------------------- + + +@dsl_user_op +def _ld_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _ld_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), [_lv(addr, loc=loc, ip=ip)], "ld.global.f32 $0, [$1];", "=f,l", loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def _lds_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _bar_sync_named(barrier_id, num_threads: int, *, loc=None, ip=None): + _asm( + None, + [_iv(barrier_id, loc=loc, ip=ip)], + f"bar.sync $0, {num_threads};", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_nc_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _score_bits(score, *, loc=None, ip=None): + """Order-preserving bits of an FP32 score (the high word of ``noaux_tc``'s key).""" + return Int32( + _asm( + T.i32(), + [_fv(score, loc=loc, ip=ip)], + "{\n.reg .b32 b, m;\nmov.b32 b, $1;\nshr.s32 m, b, 31;\n" + "or.b32 m, m, 0x80000000;\nxor.b32 $0, b, m;\n}", + "=r,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _score_value(bits, *, loc=None, ip=None): + """Inverse of :func:`_score_bits`: the FP32 score of order-preserving bits.""" + return Float32( + _asm( + T.f32(), + [_iv(bits, loc=loc, ip=ip)], + "{\n.reg .b32 m;\nshr.s32 m, $1, 31;\nnot.b32 m, m;\n" + "or.b32 m, m, 0x80000000;\nxor.b32 m, $1, m;\nmov.b32 $0, m;\n}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _tanh_approx(x, *, loc=None, ip=None): + return Float32( + _asm(T.f32(), [_fv(x, loc=loc, ip=ip)], "tanh.approx.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _fv(value, loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _max_u32(a, b, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(a, loc=loc, ip=ip), _iv(b, loc=loc, ip=ip)], + "max.u32 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _warp_16th_largest_u32(value, lane, *, loc=None, ip=None): + """16th largest of the 32 lane values (unsigned), in every lane. + + Bitonic-sorts the lower 16 lanes descending and the upper 16 ascending; the pairwise maxima + of lanes ``i`` and ``i + 16`` are then the 16 largest values, whose minimum is the answer. + """ + body = "{\n.reg .b32 x, o, mx, mn, t, u;\n.reg .pred pj, pk, km;\nmov.b32 x, $1;\n" + k = 2 + while k <= 16: + j = k // 2 + while j >= 1: + body += f"shfl.sync.bfly.b32 o, x, {j}, 0x1f, 0xffffffff;\n" + body += "max.u32 mx, x, o;\nmin.u32 mn, x, o;\n" + # Keep the maximum iff the lane is the lower one of its pair in a descending block. + body += f"and.b32 t, $2, {j};\nsetp.eq.u32 pj, t, 0;\n" + body += f"and.b32 u, $2, {k};\nsetp.eq.u32 pk, u, 0;\nxor.pred km, pj, pk;\n" + body += "selp.b32 x, mn, mx, km;\n" + j //= 2 + k *= 2 + body += "shfl.sync.bfly.b32 o, x, 16, 0x1f, 0xffffffff;\nmax.u32 x, x, o;\n" + body += "redux.sync.min.u32 $0, x, 0xffffffff;\n}" + return Int32( + _asm( + T.i32(), + [_iv(value, loc=loc, ip=ip), _iv(lane, loc=loc, ip=ip)], + body, + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _count_ge_u32(values, threshold, *, loc=None, ip=None): + """Number of ``values`` (unsigned) at or above ``threshold``.""" + n = len(values) + body = "{\n.reg .pred p;\nmov.u32 $0, 0;\n" + for i in range(n): + body += f"setp.ge.u32 p, ${2 + i}, $1;\n@p add.u32 $0, $0, 1;\n" + body += "}" + return Int32( + _asm( + T.i32(), + [_iv(threshold, loc=loc, ip=ip)] + [_iv(v, loc=loc, ip=ip) for v in values], + body, + "=&r,r" + ",r" * n, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _store_candidates(addr, his, los, sigs, threshold, *, loc=None, ip=None): + """Append the slots with ``hi >= threshold`` as 16-byte records at ``addr``; returns the end. + + Record: ``{65535 - expert, hi, sigmoid bits, 0}`` (``ld.shared.b64`` of its first eight + bytes is ``noaux_tc``'s key). + """ + n = len(his) + body = "{\n.reg .pred p;\n.reg .b32 z;\nmov.b32 z, 0;\nmov.u32 $0, $1;\n" + for i in range(n): + h, lo, s = 3 + i, 3 + n + i, 3 + 2 * n + i + body += ( + f"setp.ge.u32 p, ${h}, $2;\n" + f"@p st.shared.v4.b32 [$0], {{${lo}, ${h}, ${s}, z}};\n" + "@p add.u32 $0, $0, 16;\n" + ) + body += "}" + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip), _iv(threshold, loc=loc, ip=ip)] + + [_iv(v, loc=loc, ip=ip) for v in his] + + [_iv(v, loc=loc, ip=ip) for v in los] + + [_fv(v, loc=loc, ip=ip) for v in sigs], + body, + "=&r,r,r" + ",r" * (2 * n) + ",f" * n, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _count_greater4(acc, addr, key, *, loc=None, ip=None): + """``acc`` plus how many of the four candidate keys at ``addr`` (16-byte records) exceed + ``key``.""" + body = "{\n.reg .b64 a, b, c, d;\n.reg .pred p, q, r, s;\n" + body += "ld.shared.b64 a, [$2];\nld.shared.b64 b, [$2+16];\n" + body += "ld.shared.b64 c, [$2+32];\nld.shared.b64 d, [$2+48];\n" + body += "setp.gt.u64 p, a, $3;\nsetp.gt.u64 q, b, $3;\nsetp.gt.u64 r, c, $3;\nsetp.gt.u64 s, d, $3;\n" + body += "mov.u32 $0, $1;\n@p add.u32 $0, $0, 1;\n@q add.u32 $0, $0, 1;\n" + body += "@r add.u32 $0, $0, 1;\n@s add.u32 $0, $0, 1;\n}" + return Int32( + _asm( + T.i32(), + [_iv(acc, loc=loc, ip=ip), _iv(addr, loc=loc, ip=ip), _lv(key, loc=loc, ip=ip)], + body, + "=&r,r,r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_dsmem_b32(addr, rank, value, *, loc=None, ip=None): + """32-bit store into the shared memory of cluster CTA ``rank`` at its offset ``addr``.""" + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(rank, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "{\n.reg .b32 r;\nmapa.shared::cluster.u32 r, $0, $1;\nst.shared::cluster.b32 [r], $2;\n}", + "r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_zero16(addr, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)], + "{\n.reg .b32 z;\nmov.b32 z, 0;\nst.shared.v4.b32 [$0], {z, z, z, z};\n}", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_b64_or_zero(valid, addr, *, loc=None, ip=None): + """``*addr`` if ``valid != 0`` else 0.""" + return Int64( + _asm( + T.i64(), + [_iv(valid, loc=loc, ip=ip), _iv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\nmov.b64 $0, 0;\n@p ld.shared.b64 $0, [$2];\n}", + "=&l,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _syncwarp(*, loc=None, ip=None): + _asm(None, [], "bar.warp.sync 0xffffffff;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _shfl_idx_b32(value, src_lane, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(value, loc=loc, ip=ip), _iv(src_lane, loc=loc, ip=ip)], + "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _shfl_idx_f32(value, src_lane, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(value, loc=loc, ip=ip), _iv(src_lane, loc=loc, ip=ip)], + "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;", + "=f,f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(value, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(value, loc=loc, ip=ip)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _redux_min_u32(value, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(value, loc=loc, ip=ip)], + "redux.sync.min.u32 $0, $1, 0xffffffff;", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _argmax_scan4( + best_hi, best_slot, best_sig, removed, his, sigs, slot0: int, *, loc=None, ip=None +): + """Fold four slots (``slot0..slot0+3``) into a lane's running best ``(hi, slot, sigmoid)``: + the largest ``hi`` whose bit in ``removed`` is clear, the first slot on ties (slots are in + expert order, so the first is the smallest expert).""" + body = "{\n.reg .pred p;\n.reg .b32 h, b;\nmov.u32 $0, $3;\nmov.u32 $1, $4;\nmov.b32 $2, $5;\n" + for i in range(4): + body += ( + f"bfe.u32 b, $6, {slot0 + i}, 1;\n" + "setp.eq.u32 p, b, 0;\n" + f"selp.b32 h, ${7 + i}, 0, p;\n" + "setp.gt.u32 p, h, $0;\n" + "@p mov.u32 $0, h;\n" + f"@p mov.u32 $1, {slot0 + i};\n" + f"@p mov.b32 $2, ${11 + i};\n" + ) + body += "}" + res = _asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.f32()]), + [ + _iv(best_hi, loc=loc, ip=ip), + _iv(best_slot, loc=loc, ip=ip), + _fv(best_sig, loc=loc, ip=ip), + _iv(removed, loc=loc, ip=ip), + ] + + [_iv(v, loc=loc, ip=ip) for v in his] + + [_fv(v, loc=loc, ip=ip) for v in sigs], + body, + "=&r,=&r,=&f,r,r,f,r,r,r,r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + return ( + Int32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)), + Int32(llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), res, [2], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _cluster_arrive_release(*, loc=None, ip=None): + _asm(None, [], "barrier.cluster.arrive.release.aligned;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _cluster_wait_acquire(*, loc=None, ip=None): + _asm(None, [], "barrier.cluster.wait.acquire.aligned;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_wait(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.wait;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _griddepcontrol_launch_dependents(*, loc=None, ip=None): + _asm(None, [], "griddepcontrol.launch_dependents;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _low_bits_mask(n, *, loc=None, ip=None): + """``(1 << n) - 1`` for ``0 <= n < 32``.""" + return Int32( + _asm( + T.i32(), + [_iv(n, loc=loc, ip=ip)], + "{\n.reg .b32 t;\nshl.b32 t, 1, $1;\nsub.u32 $0, t, 1;\n}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bit(n, *, loc=None, ip=None): + """``1 << n`` for ``0 <= n < 32``.""" + return Int32( + _asm(T.i32(), [_iv(n, loc=loc, ip=ip)], "shl.b32 $0, 1, $1;", "=r,r", loc=loc, ip=ip) + ) + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +class K3RouteTailKernel: + """Top-16 + ``moe_sort`` of up to 256 tokens on one cluster of ``num_ctas`` CTAs. + + The warps of a CTA form teams of ``warps_per_token`` warps; a team routes one token at a + time, lane ``g`` of the team (``g < 32 * warps_per_token``) owning experts + ``j * 32 * warps_per_token + g``. + + Args: + num_ctas: CTAs (cluster size), ``1..8``; host side guarantees ``ceil(M / num_ctas) + <= 32``. + num_warps: Warps per CTA. + warps_per_token: Warps routing one token (1, 2, 4 or 7). + num_splits: Split-K partial logits summed in split order (1 = plain logits). + with_sort: Also produce the ``moe_sort`` outputs. + use_pdl: Launch with programmatic dependent launch (the logits come from the + previous kernel). + local_cap: Largest number of local experts served. + """ + + def __init__( + self, + num_ctas: int, + num_warps: int, + warps_per_token: int, + num_splits: int, + with_sort: bool, + use_pdl: bool, + local_cap: int = 224, + ): + if not 1 <= num_ctas <= MAX_CTAS: + raise ValueError(f"num_ctas={num_ctas} must be in [1, {MAX_CTAS}]") + if warps_per_token not in (1, 2, 4, 7) or num_warps % warps_per_token: + raise ValueError(f"warps_per_token={warps_per_token} must be 1, 2, 4 or 7") + if num_warps // warps_per_token > 15 and warps_per_token > 1: + raise ValueError("at most 15 multi-warp teams (named barriers 1..15)") + if not 1 <= num_splits <= 64: + raise ValueError(f"num_splits={num_splits} must be in [1, 64]") + self.num_ctas = num_ctas + self.num_warps = num_warps + self.threads = 32 * num_warps + self.team_warps = warps_per_token + self.team_lanes = 32 * warps_per_token + self.teams = num_warps // warps_per_token + self.slots = NUM_EXPERTS // self.team_lanes + # Candidate records per warp (plus zero padding to a multiple of 4); more fall back to + # the exact arg-max rounds. + self.region_cap = 64 if warps_per_token == 1 else 32 + self.region_slots = self.region_cap + 4 + self.local_cap = local_cap + self.num_splits = num_splits + self.with_sort = with_sort + self.use_pdl = use_pdl + # Local experts handled per thread in the sort scan (contiguous). + self.experts_per_thread = (local_cap + self.threads - 1) // self.threads + self.max_tiles_per_expert = (MAX_M + TILE - 1) // TILE + # Skip candidate-ranking rounds in which the warp has no candidate. + self.skip_empty_rank_rounds = False + # Approximate candidate pass (cluster kernel only). + self.approx = type(self) is K3RouteTailKernel + + @cute.jit + def __call__( + self, + m: Int32, + rows: Int32, + logits_addr: Int64, + bias_addr: Int64, + ids_addr: Int64, + vals_addr: Int64, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + scale_bits: Int64, + logits_out_addr: Int64, + stream: cuda.CUstream, + ): + teams = self.teams + + @cute.struct + class SharedStorage: + bias: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, NUM_EXPERTS], 16] + mask: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.local_cap], 16] + # Every CTA's masks, broadcast over distributed shared memory. + masks: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, self.num_ctas * self.local_cap], 16 + ] + base: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.local_cap], 16] + ids: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, _MAX_TOKENS_PER_CTA * TOP_K], 16 + ] + cand: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, self.num_warps * self.region_slots * 4], 16 + ] + top: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, teams * TOP_K * 2], 16] + # Per team and warp: threshold, then candidate count. + xchg: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, teams * 16], 16] + warp_tot: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.num_warps], 16] + + self.shared_storage = SharedStorage + + kernel = self.kernel( + m, + rows, + logits_addr, + bias_addr, + ids_addr, + vals_addr, + e2p_addr, + p2e_addr, + tile_expert_addr, + tile_limit_addr, + total_addr, + num_tiles_addr, + local_offset, + local_num, + scale_bits, + logits_out_addr, + ) + if cutlass.const_expr(self.num_ctas > 1): + kernel.launch( + grid=(self.num_ctas, 1, 1), + block=(self.threads, 1, 1), + cluster=(self.num_ctas, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + else: + kernel.launch( + grid=(1, 1, 1), + block=(self.threads, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + m: Int32, + rows: Int32, + logits_addr: Int64, + bias_addr: Int64, + ids_addr: Int64, + vals_addr: Int64, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + scale_bits: Int64, + logits_out_addr: Int64, + ): + tidx, _, _ = cute.arch.thread_idx() + rank, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + storage = utils.SmemAllocator().allocate(self.shared_storage) + bias_smem = storage.bias.data_ptr().toint() + mask_smem = storage.mask.data_ptr().toint() + masks_smem = storage.masks.data_ptr().toint() + base_smem = storage.base.data_ptr().toint() + ids_smem = storage.ids.data_ptr().toint() + team = warp // self.team_warps + team_warp = warp - team * self.team_warps + cand_smem = storage.cand.data_ptr().toint() + team * ( + self.team_warps * self.region_slots * _CAND_BYTES + ) + top_smem = storage.top.data_ptr().toint() + team * (TOP_K * 8) + xchg_smem = storage.xchg.data_ptr().toint() + team * 64 + warp_tot_smem = storage.warp_tot.data_ptr().toint() + + per_cta = (m + self.num_ctas - 1) // self.num_ctas + t_begin = rank * per_cta + t_end = t_begin + per_cta + if t_end > m: + t_end = m + + if cutlass.const_expr(self.with_sort): + for j in cutlass.range(tidx, local_num, self.threads, unroll=1): + _sts_b32(mask_smem + j * 4, Int32(0)) + # The bias is a weight: stage it in shared memory before waiting on the previous grid. + for e4 in cutlass.range(tidx, NUM_EXPERTS // 4, self.threads, unroll=1): + v = _ld_nc_v4_f32(bias_addr + Int64(e4 * 16)) + _sts_v4_f32(bias_smem + e4 * 16, v[0], v[1], v[2], v[3]) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + cute.arch.barrier() + + for token in cutlass.range(t_begin + team, t_end, self.teams, unroll=1): + self.topk_token( + token, + token - t_begin, + team, + team_warp, + lane, + rows, + logits_addr, + bias_smem, + ids_addr, + vals_addr, + scale_bits, + logits_out_addr, + cand_smem, + top_smem, + xchg_smem, + ids_smem, + mask_smem, + local_offset, + local_num, + ) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.with_sort): + cute.arch.barrier() + if cutlass.const_expr(self.num_ctas > 1): + # Broadcast this CTA's token masks into every CTA of the cluster. + for j in cutlass.range(tidx, local_num, self.threads, unroll=1): + word = _lds_b32(mask_smem + j * 4) + dst = masks_smem + (rank * self.local_cap + j) * 4 + for c in cutlass.range_constexpr(self.num_ctas): + _st_dsmem_b32(dst, Int32(c), word) + _cluster_arrive_release() + _cluster_wait_acquire() + self.sort( + tidx, + warp, + lane, + rank, + t_begin, + t_end, + mask_smem, + masks_smem, + base_smem, + ids_smem, + warp_tot_smem, + e2p_addr, + p2e_addr, + tile_expert_addr, + tile_limit_addr, + total_addr, + num_tiles_addr, + local_offset, + local_num, + ) + + def load_logits(self, row: Int64, g: Int32, stride: int, n: int): + """Logit ``s * stride + g`` (``s < n``) of the first split of the token whose logits start + at ``row``. A subclass may read them differently (the ULL MoE front polls them).""" + return [_ld_f32(row + Int64((s * stride + g) * 4)) for s in range(n)] + + def load_logits4(self, row: Int64, lane: Int32, n: int): + """Logits ``128 j + 4 lane .. + 3`` (``j < n``) of the first split of the token whose + logits start at ``row``, concatenated. Not ``cute.jit`` (the DSL preprocessor does not + rewrite this method), so the loop is a plain Python loop, unrolled at trace time.""" + logit = [] + for blk in range(n): + logit.extend(_ld_v4_f32(row + Int64((blk * 128 + lane * 4) * 4))) + return logit + + @cute.jit + def team_sync(self, team: Int32): + if cutlass.const_expr(self.team_warps > 1): + _bar_sync_named(team + 1, self.team_lanes) + else: + _syncwarp() + + @cute.jit + def topk_token( + self, + token: Int32, + token_local: Int32, + team: Int32, + team_warp: Int32, + lane: Int32, + rows: Int32, + logits_addr: Int64, + bias_smem: Int32, + ids_addr: Int64, + vals_addr: Int64, + scale_bits: Int64, + logits_out_addr: Int64, + cand_smem: Int32, + top_smem: Int32, + xchg_smem: Int32, + ids_smem: Int32, + mask_smem: Int32, + local_offset: Int32, + local_num: Int32, + ): + """Top-16 of one token by one team; its first warp writes ids/weights and mask bits.""" + expert, weight = self.select_token( + token, + team, + team_warp, + lane, + rows, + logits_addr, + bias_smem, + scale_bits, + logits_out_addr, + cand_smem, + top_smem, + xchg_smem, + ) + if team_warp == 0: + if lane < TOP_K: + out = (token * TOP_K + lane) * 4 + _st_global_b32(ids_addr + Int64(out), expert) + _st_global_f32(vals_addr + Int64(out), weight) + if cutlass.const_expr(self.with_sort): + _sts_b32(ids_smem + (token_local * TOP_K + lane) * 4, expert) + local = expert - local_offset + if local >= 0 and local < local_num: + _red_shared_or(mask_smem + local * 4, _bit(token_local)) + + @cute.jit + def select_token( + self, + token: Int32, + team: Int32, + team_warp: Int32, + lane: Int32, + rows: Int32, + logits_addr: Int64, + bias_smem: Int32, + scale_bits: Int64, + logits_out_addr: Int64, + cand_smem: Int32, + top_smem: Int32, + xchg_smem: Int32, + ): + """Top-16 of one token by one team: returns ``(expert, weight)``, the ``lane``-th + selection in lanes ``0..15`` of the team's first warp.""" + n_slots = self.slots + stride = self.team_lanes + g = team_warp * 32 + lane + row = logits_addr + Int64(token) * Int64(NUM_EXPERTS * 4) + split_stride = Int64(rows) * Int64(NUM_EXPERTS * 4) + logit = self.load_logits(row, g, stride, n_slots) + for sp in cutlass.range_constexpr(1, self.num_splits): + for s in cutlass.range_constexpr(n_slots): + logit[s] = logit[s] + _ld_f32(row + split_stride * sp + Int64((s * stride + g) * 4)) + if logits_out_addr != 0: + out_row = logits_out_addr + Int64(token) * Int64(NUM_EXPERTS * 4) + for s in cutlass.range_constexpr(n_slots): + _st_global_f32(out_row + Int64((s * stride + g) * 4), logit[s]) + bias = [_lds_f32(bias_smem + (s * stride + g) * 4) for s in range(n_slots)] + if cutlass.const_expr(self.approx): + # Approximate scores a_e, |a_e - s_e| <= delta (APPROX_MARGIN = 2 delta) for the exact + # scores s_e. Records keep the logit (in the sigmoid field) for the exact pass below. + sig = [ + _tanh_approx(logit[s] * Float32(0.5)) * Float32(0.5) + Float32(0.5) + for s in range(n_slots) + ] + else: + sig = [ + cute.math.tanh(logit[s] * Float32(0.5)) * Float32(0.5) + Float32(0.5) + for s in range(n_slots) + ] + his = [_score_bits(sig[s] + bias[s]) for s in range(n_slots)] + lane_max = his[0] + for s in cutlass.range_constexpr(1, n_slots): + lane_max = _max_u32(lane_max, his[s]) + # A lower bound of the 16th largest score: the 16th largest lane maximum of a warp (16 + # of the warp's experts score at least that much); the team takes the largest. + threshold = _warp_16th_largest_u32(lane_max, lane) + if cutlass.const_expr(self.team_warps > 1): + if lane == 0: + _sts_b32(xchg_smem + team_warp * 4, threshold) + self.team_sync(team) + for w in cutlass.range_constexpr(self.team_warps): + threshold = _max_u32(threshold, _lds_b32(xchg_smem + w * 4)) + if cutlass.const_expr(self.approx): + # The threshold t is at most the 16th largest approximate score A16, and A16 <= T* + + # delta for the 16th largest exact score T* (more than 15 approximate scores above + # T* + delta would mean more than 15 exact scores above T*). Every expert of the exact + # top 16 has a_e >= s_e - delta >= T* - delta >= t - 2 delta: all are candidates, and + # ranking the candidates by their exact keys gives noaux_tc's selection and order. + threshold = _score_bits(_score_value(threshold) - Float32(APPROX_MARGIN)) + count = Int32(0) + for c4 in cutlass.range_constexpr(0, n_slots, 4): + count = count + _count_ge_u32(his[c4 : min(c4 + 4, n_slots)], threshold) + inclusive = _warp_inclusive_sum(count) + num_mine = _shfl_idx_b32(inclusive, Int32(31)) + cap = self.region_cap + region_bytes = self.region_slots * _CAND_BYTES + region = cand_smem + team_warp * region_bytes + if num_mine <= cap: + los = [Int32(65535 - s * stride) - g for s in range(n_slots)] + dst = region + (inclusive - count) * _CAND_BYTES + # Approximate pass: the record's sigmoid field carries the logit for the exact pass. + carried = sig + if cutlass.const_expr(self.approx): + carried = logit + for c4 in cutlass.range_constexpr(0, n_slots, 4): + q = slice(c4, min(c4 + 4, n_slots)) + dst = _store_candidates(dst, his[q], los[q], carried[q], threshold) + # Zero keys after the last candidate, so the ranking reads groups of four. + if lane < 4: + _sts_zero16(region + (num_mine + lane) * _CAND_BYTES) + if cutlass.const_expr(self.approx): + if num_mine <= cap: + # Exact pass over this warp's candidates: noaux_tc's sigmoid and key, in place. + _syncwarp() + for c0 in cutlass.range_constexpr(cap // 32): + c = lane + c0 * 32 + if c < num_mine: + rec = region + c * _CAND_BYTES + e = 65535 - _lds_b32(rec) + sg = cute.math.tanh(_lds_f32(rec + 8) * Float32(0.5)) * Float32( + 0.5 + ) + Float32(0.5) + _sts_b32(rec + 4, _score_bits(sg + _lds_f32(bias_smem + e * 4))) + _sts_f32(rec + 8, sg) + if cutlass.const_expr(self.team_warps > 1): + if lane == 0: + _sts_b32(xchg_smem + 32 + team_warp * 4, num_mine) + self.team_sync(team) + counts = [num_mine] + overflow = Int32(0) + if num_mine > cap: + overflow = Int32(1) + if cutlass.const_expr(self.team_warps > 1): + counts = [_lds_b32(xchg_smem + 32 + w * 4) for w in range(self.team_warps)] + for w in cutlass.range_constexpr(self.team_warps): + if counts[w] > cap: + overflow = Int32(1) + + expert = Int32(0) + score = Float32(0.0) + if overflow == 0: + # Rank this warp's candidates against all candidates of the team. + for c0 in cutlass.range_constexpr(cap // 32): + c = lane + c0 * 32 + valid = Int32(0) + if c < num_mine: + valid = Int32(1) + key = _lds_b64_or_zero(valid, region + c * _CAND_BYTES) + r = Int32(0) + for w in cutlass.range_constexpr(self.team_warps): + other = cand_smem + w * region_bytes + limit = counts[w] + if cutlass.const_expr(c0 > 0 and self.skip_empty_rank_rounds): + # No candidate of this warp in this round: skip the count. + if num_mine <= c0 * 32: + limit = Int32(0) + for o in cutlass.range(0, limit, 4, unroll=1): + r = _count_greater4(r, other + o * _CAND_BYTES, key) + if valid != 0 and r < TOP_K: + _sts_b32(top_smem + r * 8, 65535 - (Int32(key) & 0xFFFF)) + _sts_b32(top_smem + r * 8 + 4, _lds_b32(region + c * _CAND_BYTES + 8)) + self.team_sync(team) + if team_warp == 0 and lane < TOP_K: + expert = _lds_b32(top_smem + lane * 8) + score = _lds_f32(top_smem + lane * 8 + 4) + else: + if team_warp == 0: + expert, score = self.topk_rounds(row, split_stride, bias_smem, lane) + + weight = Float32(0.0) + if team_warp == 0: + total = _warp_sum_f32(score) + # Exact FP32 fast path of noaux_tc's FP64 expression (see k3_fused_router). + if _unit_scale_fast_path(scale_bits, total) != 0: + weight = _div_rn_f32(score, total) + else: + weight = _renormalize(score, total, scale_bits) + return expert, weight + + @cute.jit + def topk_rounds(self, row: Int64, split_stride: Int64, bias_smem: Int32, lane: Int32): + """Exact fallback by one warp: 16 arg-max rounds over all 896 keys of the token (lane + ``l`` owns experts ``128 j + 4 l + i``); lane ``k`` gets the ``k``-th expert.""" + logit = self.load_logits4(row, lane, _BLOCKS) + for sp in cutlass.range_constexpr(1, self.num_splits): + for blk in cutlass.range_constexpr(_BLOCKS): + part = _ld_v4_f32(row + split_stride * sp + Int64((blk * 128 + lane * 4) * 4)) + for i in cutlass.range_constexpr(4): + logit[blk * 4 + i] = logit[blk * 4 + i] + part[i] + sig = [ + cute.math.tanh(logit[s] * Float32(0.5)) * Float32(0.5) + Float32(0.5) + for s in range(_SLOTS) + ] + his = [] + for blk in cutlass.range_constexpr(_BLOCKS): + b = _lds_v4_f32(bias_smem + (blk * 128 + lane * 4) * 4) + for i in cutlass.range_constexpr(4): + his.append(_score_bits(sig[blk * 4 + i] + b[i])) + removed = Int32(0) + expert = Int32(0) + score = Float32(0.0) + for r in cutlass.range_constexpr(TOP_K): + best_hi = Int32(0) + best_slot = Int32(0) + best_sig = Float32(0.0) + for blk in cutlass.range_constexpr(_BLOCKS): + q = slice(blk * 4, blk * 4 + 4) + best_hi, best_slot, best_sig = _argmax_scan4( + best_hi, best_slot, best_sig, removed, his[q], sig[q], blk * 4 + ) + top_hi = _redux_max_u32(best_hi) + # Slot s of this lane is expert 128 (s / 4) + 4 lane + s % 4. + cand = Int32(0xFFFF) + if best_hi == top_hi: + cand = (best_slot >> 2) * 128 + lane * 4 + (best_slot & 3) + winner = _redux_min_u32(cand) + owner = (winner & 127) >> 2 + if lane == owner: + removed = removed | _bit(best_slot) + win_sig = _shfl_idx_f32(best_sig, owner) + if lane == r: + expert = winner + score = win_sig + return expert, score + + @cute.jit + def sort( + self, + tidx: Int32, + warp: Int32, + lane: Int32, + rank: Int32, + t_begin: Int32, + t_end: Int32, + mask_smem: Int32, + masks_smem: Int32, + base_smem: Int32, + ids_smem: Int32, + warp_tot_smem: Int32, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + ): + """``moe_sort`` metadata from the token masks of all CTAs of the cluster.""" + ept = self.experts_per_thread + j0 = tidx * ept + counts = [] + before = [] + for q in cutlass.range_constexpr(ept): + j = j0 + q + jj = j + if j >= local_num: + jj = Int32(0) + words = [] + for c in cutlass.range_constexpr(self.num_ctas): + if cutlass.const_expr(self.num_ctas > 1): + words.append(_lds_b32(masks_smem + (c * self.local_cap + jj) * 4)) + else: + words.append(_lds_b32(mask_smem + jj * 4)) + cnt = Int32(0) + pre = Int32(0) + for c in cutlass.range_constexpr(self.num_ctas): + n = _popc(words[c]) + cnt = cnt + n + if c < rank: + pre = pre + n + if j >= local_num: + cnt = Int32(0) + counts.append(cnt) + before.append(pre) + tiles = [(counts[q] + (TILE - 1)) >> 7 for q in range(ept)] + thread_tiles = tiles[0] + for q in cutlass.range_constexpr(1, ept): + thread_tiles = thread_tiles + tiles[q] + inclusive = _warp_inclusive_sum(thread_tiles) + if lane == 31: + _sts_b32(warp_tot_smem + warp * 4, inclusive) + cute.arch.barrier() + tile = inclusive - thread_tiles + num_tiles = Int32(0) + for w in cutlass.range_constexpr(self.num_warps): + cnt_w = _lds_b32(warp_tot_smem + w * 4) + if w < warp: + tile = tile + cnt_w + num_tiles = num_tiles + cnt_w + for q in cutlass.range_constexpr(ept): + j = j0 + q + if tiles[q] != 0: + _sts_b32(base_smem + j * 4, tile * TILE + before[q]) + if rank == 0: + for t in cutlass.range_constexpr(self.max_tiles_per_expert): + if t < tiles[q]: + limit = tile * TILE + counts[q] + if limit > (tile + t + 1) * TILE: + limit = (tile + t + 1) * TILE + _st_global_b32(tile_expert_addr + Int64((tile + t) * 4), j) + _st_global_b32(tile_limit_addr + Int64((tile + t) * 4), limit) + tile = tile + tiles[q] + if rank == 0 and tidx == 0: + _st_global_b32(total_addr, num_tiles * TILE) + _st_global_b32(num_tiles_addr, num_tiles) + cute.arch.barrier() + + n_entries = (t_end - t_begin) * TOP_K + for e in cutlass.range(tidx, n_entries, self.threads, unroll=1): + token_local = e >> 4 + expanded = t_begin * TOP_K + e + local = _lds_b32(ids_smem + e * 4) - local_offset + perm = Int32(-1) + if local >= 0 and local < local_num: + word = _lds_b32(mask_smem + local * 4) + perm = _lds_b32(base_smem + local * 4) + _popc(word & _low_bits_mask(token_local)) + _st_global_b32(p2e_addr + Int64(perm * 4), expanded) + _st_global_b32(e2p_addr + Int64(expanded * 4), perm) + + +# --------------------------------------------------------------------------- +# Multi-CTA routing (M > 256) +# --------------------------------------------------------------------------- + +#: Largest token count of the multi-CTA routing (:class:`K3RouteWideKernel`). +WIDE_MAX_M = 2048 +#: Largest number of local experts the multi-CTA sort serves (expert parallelism >= 4). +WIDE_LOCAL_CAP = 224 +#: Cells (local experts x CTAs) of the sort's shared-memory table (32 KB): the sort serves +#: ``local_num_experts * ceil(M / tokens_per_cta) <= 8192``. +WIDE_TABLE_CELLS = 8192 +_WIDE_MAX_CTAS = WIDE_MAX_M // 8 +#: Workspace of the multi-CTA sort: arrival counter, then per CTA its number of local +#: selections, the ``[CTA][local expert]`` token masks, and per CTA a list of its local +#: selections (``tokens_per_cta * 16`` slots). +_WIDE_COUNT_OFFSET = 128 +_WIDE_MASK_OFFSET = _WIDE_COUNT_OFFSET + _WIDE_MAX_CTAS * 4 +_WIDE_LIST_OFFSET = _WIDE_MASK_OFFSET + WIDE_TABLE_CELLS * 4 +_WIDE_WORKSPACE_BYTES = _WIDE_LIST_OFFSET + WIDE_MAX_M * TOP_K * 4 +#: Smem words of the top-k phase per region (aliased by the sort's table): bias, then per team. +_WIDE_BIAS_WORDS = NUM_EXPERTS +#: Cells per step of the sort's per-expert scan. +_SCAN_CHUNK = 8 +#: Shared-memory budget of the compact sort tables of 8-token CTAs: ``CTAs x L4`` mask bytes +#: (``L4``: local experts rounded up to 4), then 16-bit rows-before per 8-CTA block and expert; +#: the size of the top-k scratch of an 8-warp CTA, which they alias. +WIDE_COMPACT_TABLE_BYTES = 13824 + + +@dsl_user_op +def _ld_relaxed_b32_or_zero(valid, addr, *, loc=None, ip=None): + """``*addr`` (relaxed, GPU scope) if ``valid != 0`` else 0.""" + return Int32( + _asm( + T.i32(), + [_iv(valid, loc=loc, ip=ip), _lv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\nmov.b32 $0, 0;\n" + "@p ld.relaxed.gpu.global.b32 $0, [$2];\n}", + "=&r,r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _lds_b32_or_zero(valid, addr, *, loc=None, ip=None): + """Shared ``*addr`` if ``valid != 0`` else 0.""" + return Int32( + _asm( + T.i32(), + [_iv(valid, loc=loc, ip=ip), _iv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\nmov.b32 $0, 0;\n@p ld.shared.b32 $0, [$2];\n}", + "=&r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _lds_u8_or_zero(valid, addr, *, loc=None, ip=None): + """Shared byte ``*addr`` if ``valid != 0`` else 0.""" + return Int32( + _asm( + T.i32(), + [_iv(valid, loc=loc, ip=ip), _iv(addr, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\nmov.b32 $0, 0;\n@p ld.shared.u8 $0, [$2];\n}", + "=&r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sts_u16(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.u16.u32 h, $1;\nst.shared.u16 [$0], h;\n}", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_u16(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.u16 $0, [$1];", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sts_b32_if(valid, addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(valid, loc=loc, ip=ip), _iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $0, 0;\n@p st.shared.b32 [$1], $2;\n}", + "r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _atom_shared_add(addr, value, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "atom.shared.add.u32 $0, [$1], $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +class K3RouteWideKernel(K3RouteTailKernel): + """Top-16 (+ ``moe_sort``) of up to :data:`WIDE_MAX_M` tokens on a grid of independent CTAs. + + CTA ``c`` routes tokens ``[T c, T c + T)`` (``T = tokens_per_cta``) with the teams of + :class:`K3RouteTailKernel` (same selection code, so the same top-k). The sort is done by the + CTA that finishes last: + + * Each CTA marks, per local expert ``j``, a ``T``-bit mask of its tokens that selected ``j`` + in shared memory, appends each such selection ``(j, token, slot k)`` to its list in the + workspace and stores ``expanded_idx_to_permuted_idx = -1`` for its selections of other + ranks' experts. It then publishes the masks (``[c][j]``) and its list length and takes a + ticket from an arrival counter (release fence first). + * The CTA that draws the last ticket resets the counter and builds ``moe_sort``'s metadata: + it loads the whole mask table and the first ``2 T`` list entries of every CTA in one round + of loads (later entries, rare at EP16, are loaded as needed), then, in shared memory, + turns each expert's masks into token counts and row offsets (rows of an expert in + ascending token order), lays out the tiles like :class:`K3RouteTailKernel`, and writes + the rows of all listed selections. 16-token CTAs keep a 32-bit cell per (CTA, expert): the + expert's rows in earlier CTAs above the mask. 8-token CTAs keep byte masks plus the rows + before each block of 8 CTAs (:data:`WIDE_COMPACT_TABLE_BYTES`), aliased with the top-k + scratch, so a 256-thread CTA needs about 14 KB of shared memory and can share an SM with + a persistent GEMM CTA. + + No CTA waits on another, so the grid needs no co-residency. Launches sharing the workspace + must not overlap (stream order provides that); each launch leaves the counter at zero, so a + captured CUDA graph replays as is. + + Args: + tokens_per_cta: Tokens per CTA: 8 (compact sort tables) or 16, or 32 (top-k only). + num_warps: Warps per CTA: 8 or 16 (with the sort), or 4 (top-k only). + warps_per_token: Warps routing one token (1, 2, 4 or 7). + num_splits: Split-K partial logits summed in split order (1 = plain logits). + with_sort: Also produce the ``moe_sort`` outputs. + use_pdl: Launch with programmatic dependent launch. + local_cap: Largest number of local experts served (at most :data:`WIDE_LOCAL_CAP`). + """ + + def __init__( + self, + tokens_per_cta: int, + num_warps: int, + warps_per_token: int, + num_splits: int, + with_sort: bool, + use_pdl: bool, + local_cap: int = WIDE_LOCAL_CAP, + ): + super().__init__(1, num_warps, warps_per_token, num_splits, with_sort, use_pdl, local_cap) + if tokens_per_cta not in ((8, 16) if with_sort else (8, 16, 32)): + raise ValueError(f"tokens_per_cta={tokens_per_cta} must be 8 or 16 (32 top-k only)") + if num_warps not in ((8, 16) if with_sort else (4, 8, 16)): + raise ValueError(f"num_warps={num_warps} must be 8 or 16 (4 top-k only)") + if local_cap > WIDE_LOCAL_CAP: + raise ValueError(f"local_cap={local_cap} must be at most {WIDE_LOCAL_CAP}") + self.tokens = tokens_per_cta + self.max_ctas = WIDE_MAX_M // tokens_per_cta + # Sort: mask-table words per thread; list entries per CTA loaded up front (twice the + # tokens, about twice the expected local selections at EP16), 32 / list_slots CTAs per + # warp-wide load. + self.compact = with_sort and tokens_per_cta == 8 + self.table_words = WIDE_COMPACT_TABLE_BYTES // 4 if self.compact else WIDE_TABLE_CELLS + self.table_items = (self.table_words + self.threads - 1) // self.threads + self.list_slots = min(2 * tokens_per_cta, 32) + self.lists_per_load = 32 // self.list_slots + self.list_loads = (self.max_ctas + num_warps * self.lists_per_load - 1) // ( + num_warps * self.lists_per_load + ) + self.cand_words = num_warps * self.region_slots * 4 + self.topk_words = ( + _WIDE_BIAS_WORDS + self.cand_words + self.teams * TOP_K * 2 + self.teams * 16 + ) + self.skip_empty_rank_rounds = True + + @cute.jit + def __call__( + self, + m: Int32, + rows: Int32, + logits_addr: Int64, + bias_addr: Int64, + ids_addr: Int64, + vals_addr: Int64, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + scale_bits: Int64, + ws_addr: Int64, + stream: cuda.CUstream, + ): + scratch_words = self.topk_words + if cutlass.const_expr(self.with_sort): + scratch_words = max(scratch_words, self.table_words) + + @cute.struct + class SharedStorage: + # Top-k phase: bias, candidate regions, top-16 and exchange slots; the sort's + # [CTA][local expert] table aliases them. + scratch: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, scratch_words], 16] + mask: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.local_cap], 16] + # Sort: first tile per local expert. + tile: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.local_cap], 16] + warp_tot: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.num_warps], 16] + # [0]: last-CTA flag, [1]: number of local selections of this CTA. + flag: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, 4], 16] + + self.shared_storage = SharedStorage + + num_ctas = (m + (self.tokens - 1)) // self.tokens + self.kernel( + m, + rows, + logits_addr, + bias_addr, + ids_addr, + vals_addr, + e2p_addr, + p2e_addr, + tile_expert_addr, + tile_limit_addr, + total_addr, + num_tiles_addr, + local_offset, + local_num, + scale_bits, + ws_addr, + ).launch( + grid=(num_ctas, 1, 1), + block=(self.threads, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + m: Int32, + rows: Int32, + logits_addr: Int64, + bias_addr: Int64, + ids_addr: Int64, + vals_addr: Int64, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + scale_bits: Int64, + ws_addr: Int64, + ): + tidx, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + storage = utils.SmemAllocator().allocate(self.shared_storage) + scratch = storage.scratch.data_ptr().toint() + mask_smem = storage.mask.data_ptr().toint() + tile_smem = storage.tile.data_ptr().toint() + warp_tot_smem = storage.warp_tot.data_ptr().toint() + flag_smem = storage.flag.data_ptr().toint() + bias_smem = scratch + cand_all = bias_smem + _WIDE_BIAS_WORDS * 4 + top_all = cand_all + self.cand_words * 4 + xchg_all = top_all + self.teams * TOP_K * 8 + team = warp // self.team_warps + team_warp = warp - team * self.team_warps + cand_smem = cand_all + team * (self.team_warps * self.region_slots * _CAND_BYTES) + top_smem = top_all + team * (TOP_K * 8) + xchg_smem = xchg_all + team * 64 + + num_ctas = (m + (self.tokens - 1)) // self.tokens + t_begin = cta * self.tokens + t_end = t_begin + self.tokens + if t_end > m: + t_end = m + ws_count = ws_addr + Int64(_WIDE_COUNT_OFFSET) + ws_mask = ws_addr + Int64(_WIDE_MASK_OFFSET) + ws_list = ws_addr + Int64(_WIDE_LIST_OFFSET) + list_base = ws_list + Int64(cta * (self.tokens * TOP_K * 4)) + + if cutlass.const_expr(self.with_sort): + for j in cutlass.range(tidx, local_num, self.threads, unroll=1): + _sts_b32(mask_smem + j * 4, Int32(0)) + if tidx == 0: + _sts_b32(flag_smem + 4, Int32(0)) + # The bias is a weight: stage it in shared memory before waiting on the previous grid. + for e4 in cutlass.range(tidx, NUM_EXPERTS // 4, self.threads, unroll=1): + v = _ld_nc_v4_f32(bias_addr + Int64(e4 * 16)) + _sts_v4_f32(bias_smem + e4 * 16, v[0], v[1], v[2], v[3]) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + cute.arch.barrier() + + for token in cutlass.range(t_begin + team, t_end, self.teams, unroll=1): + expert, weight = self.select_token( + token, + team, + team_warp, + lane, + rows, + logits_addr, + bias_smem, + scale_bits, + Int64(0), + cand_smem, + top_smem, + xchg_smem, + ) + if team_warp == 0: + if lane < TOP_K: + out = (token * TOP_K + lane) * 4 + _st_global_b32(ids_addr + Int64(out), expert) + _st_global_f32(vals_addr + Int64(out), weight) + if cutlass.const_expr(self.with_sort): + local = expert - local_offset + token_local = token - t_begin + if local >= 0 and local < local_num: + _red_shared_or(mask_smem + local * 4, _bit(token_local)) + pos = _atom_shared_add(flag_smem + 4, Int32(1)) + _st_global_b32( + list_base + Int64(pos * 4), + (local << 12) | (token_local << 4) | lane, + ) + else: + _st_global_b32(e2p_addr + Int64(out), Int32(-1)) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.with_sort): + cute.arch.barrier() + if cutlass.const_expr(self.compact): + # Four 8-bit masks per word: the table is [CTA][local expert] bytes. + words = (local_num + 3) // 4 + for w in cutlass.range(tidx, words, self.threads, unroll=1): + word = Int32(0) + for b in cutlass.range_constexpr(4): + if 4 * w + b < local_num: + word = word | (_lds_b32(mask_smem + (4 * w + b) * 4) << (8 * b)) + _st_global_b32(ws_mask + Int64((cta * words + w) * 4), word) + else: + for j in cutlass.range(tidx, local_num, self.threads, unroll=1): + _st_global_b32( + ws_mask + Int64((cta * local_num + j) * 4), _lds_b32(mask_smem + j * 4) + ) + if tidx == 0: + _st_global_b32(ws_count + Int64(cta * 4), _lds_b32(flag_smem + 4)) + cute.arch.barrier() + if tidx == 0: + _fence_acq_rel_gpu() + ticket = _atom_add_relaxed_gpu(ws_addr, Int32(1)) + last = Int32(0) + if ticket == num_ctas - 1: + _fence_acq_rel_gpu() + _st_relaxed_b32(ws_addr, Int32(0)) + last = Int32(1) + _sts_b32(flag_smem, last) + cute.arch.barrier() + if _lds_b32(flag_smem) != 0: + self.wide_sort( + tidx, + warp, + lane, + num_ctas, + ws_count, + ws_mask, + ws_list, + scratch, + tile_smem, + warp_tot_smem, + e2p_addr, + p2e_addr, + tile_expert_addr, + tile_limit_addr, + total_addr, + num_tiles_addr, + local_num, + ) + + @cute.jit + def wide_sort( + self, + tidx: Int32, + warp: Int32, + lane: Int32, + num_ctas: Int32, + ws_count: Int64, + ws_mask: Int64, + ws_list: Int64, + table_smem: Int32, + tile_smem: Int32, + warp_tot_smem: Int32, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_num: Int32, + ): + """``moe_sort`` metadata from the published masks and lists of all CTAs (last CTA).""" + list_words = self.tokens * TOP_K + # Compact: byte table [CTA][row_bytes], then the 16-bit block table [CTA / 8][expert]. + row_bytes = 4 * ((local_num + 3) // 4) + block_smem = table_smem + ((num_ctas * row_bytes + 15) // 16) * 16 + cells = num_ctas * local_num + if cutlass.const_expr(self.compact): + cells = num_ctas * (row_bytes // 4) + # One round of loads: the [CTA][local expert] mask table, then the list length and the + # first list_slots entries of each CTA (lane l of a warp-wide load: CTA + # lists_per_load * (warp + q * num_warps) + l / list_slots, entry l % list_slots). + table = [] + for i in cutlass.range_constexpr(self.table_items): + idx = tidx + i * self.threads + valid = Int32(0) + if idx < cells: + valid = Int32(1) + table.append(_ld_relaxed_b32_or_zero(valid, ws_mask + Int64(idx * 4))) + slots = self.list_slots + slot = lane % slots + counts = [] + entries = [] + for q in cutlass.range_constexpr(self.list_loads): + c = (warp + q * self.num_warps) * self.lists_per_load + lane // slots + valid = Int32(0) + if c < num_ctas: + valid = Int32(1) + counts.append(_ld_relaxed_b32_or_zero(valid, ws_count + Int64(c * 4))) + entries.append( + _ld_relaxed_b32_or_zero(valid, ws_list + Int64((c * list_words + slot) * 4)) + ) + for i in cutlass.range_constexpr(self.table_items): + idx = tidx + i * self.threads + valid = Int32(0) + if idx < cells: + valid = Int32(1) + _sts_b32_if(valid, table_smem + idx * 4, table[i]) + cute.arch.barrier() + + # Thread per local expert: token count, and the expert's rows in earlier CTAs per cell + # (stored above the cell's mask, which has at most 16 bits) or, compact, per block of + # 8 CTAs. + ept = self.experts_per_thread + j0 = tidx * ept + totals = [] + for q in cutlass.range_constexpr(ept): + j = j0 + q + run = Int32(0) + if cutlass.const_expr(self.compact): + if j < local_num: + for c0 in cutlass.range(0, num_ctas, 8, unroll=1): + _sts_u16(block_smem + ((c0 >> 3) * local_num + j) * 2, run) + masks = [] + for u in cutlass.range_constexpr(8): + valid = Int32(0) + if c0 + u < num_ctas: + valid = Int32(1) + masks.append( + _lds_u8_or_zero(valid, table_smem + (c0 + u) * row_bytes + j) + ) + for u in cutlass.range_constexpr(8): + run = run + _popc(masks[u]) + else: + if j < local_num: + run = self.scan_cells(j, num_ctas, local_num, table_smem) + totals.append(run) + + # Tile layout: ceil(count / 128) tiles per local expert, in local-expert order. + tiles = [(totals[q] + (TILE - 1)) >> 7 for q in range(ept)] + thread_tiles = tiles[0] + for q in cutlass.range_constexpr(1, ept): + thread_tiles = thread_tiles + tiles[q] + inclusive = _warp_inclusive_sum(thread_tiles) + if lane == 31: + _sts_b32(warp_tot_smem + warp * 4, inclusive) + cute.arch.barrier() + tile = inclusive - thread_tiles + num_tiles = Int32(0) + for w in cutlass.range_constexpr(self.num_warps): + cnt_w = _lds_b32(warp_tot_smem + w * 4) + if w < warp: + tile = tile + cnt_w + num_tiles = num_tiles + cnt_w + for q in cutlass.range_constexpr(ept): + if j0 + q < local_num: + _sts_b32(tile_smem + (j0 + q) * 4, tile * TILE) + if tiles[q] != 0: + end = tile * TILE + totals[q] + for t in cutlass.range(tiles[q], unroll=1): + limit = (tile + t + 1) * TILE + if limit > end: + limit = end + _st_global_b32(tile_expert_addr + Int64((tile + t) * 4), j0 + q) + _st_global_b32(tile_limit_addr + Int64((tile + t) * 4), limit) + tile = tile + tiles[q] + if tidx == 0: + _st_global_b32(total_addr, num_tiles * TILE) + _st_global_b32(num_tiles_addr, num_tiles) + cute.arch.barrier() + + # Rows of the listed selections: entry (j, token t, slot k) of CTA c lands in row + # first_row(j) + rows of j in CTAs before c + tokens of the mask before t. + for q in cutlass.range_constexpr(self.list_loads): + c = (warp + q * self.num_warps) * self.lists_per_load + lane // slots + n = counts[q] + if c < num_ctas: + if slot < n: + self.place_entry( + entries[q], + c, + local_num, + row_bytes, + table_smem, + block_smem, + tile_smem, + e2p_addr, + p2e_addr, + ) + for i in cutlass.range(slot + slots, n, slots, unroll=1): + entry = _ld_relaxed_b32_or_zero( + Int32(1), ws_list + Int64((c * list_words + i) * 4) + ) + self.place_entry( + entry, + c, + local_num, + row_bytes, + table_smem, + block_smem, + tile_smem, + e2p_addr, + p2e_addr, + ) + + @cute.jit + def scan_cells(self, j: Int32, num_ctas: Int32, local_num: Int32, table_smem: Int32): + """Rewrite expert ``j``'s 32-bit cells as ``rows in earlier CTAs << 16 | mask``; + returns the expert's token count.""" + run = Int32(0) + for c0 in cutlass.range(0, num_ctas, _SCAN_CHUNK, unroll=1): + words = [] + valids = [] + for u in cutlass.range_constexpr(_SCAN_CHUNK): + valid = Int32(0) + if c0 + u < num_ctas: + valid = Int32(1) + valids.append(valid) + words.append(_lds_b32_or_zero(valid, table_smem + ((c0 + u) * local_num + j) * 4)) + for u in cutlass.range_constexpr(_SCAN_CHUNK): + _sts_b32_if( + valids[u], + table_smem + ((c0 + u) * local_num + j) * 4, + (run << 16) | words[u], + ) + run = run + _popc(words[u]) + return run + + @cute.jit + def place_entry( + self, + entry: Int32, + c: Int32, + local_num: Int32, + row_bytes: Int32, + table_smem: Int32, + block_smem: Int32, + tile_smem: Int32, + e2p_addr: Int64, + p2e_addr: Int64, + ): + """Write the row of selection ``entry`` (``j << 12 | token << 4 | slot``) of CTA ``c``.""" + j = entry >> 12 + t = (entry >> 4) & 0xFF + k = entry & 15 + rank = _lds_b32(tile_smem + j * 4) + if cutlass.const_expr(self.compact): + rank = rank + _lds_u16(block_smem + ((c >> 3) * local_num + j) * 2) + # Tokens of j in the block's earlier CTAs, then in this CTA before t. + first = c - (c & 7) + masks = [] + for u in cutlass.range_constexpr(7): + valid = Int32(0) + if first + u < c: + valid = Int32(1) + masks.append(_lds_u8_or_zero(valid, table_smem + (first + u) * row_bytes + j)) + mine = _lds_u8_or_zero(Int32(1), table_smem + c * row_bytes + j) + for u in cutlass.range_constexpr(7): + rank = rank + _popc(masks[u]) + rank = rank + _popc(mine & _low_bits_mask(t)) + else: + cell = _lds_b32(table_smem + (c * local_num + j) * 4) + rank = rank + (cell >> 16) + _popc(cell & 0xFFFF & _low_bits_mask(t)) + expanded = (c * self.tokens + t) * TOP_K + k + _st_global_b32(e2p_addr + Int64(expanded * 4), rank) + _st_global_b32(p2e_addr + Int64(rank * 4), expanded) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + +#: Local-expert capacity of the common build (expert parallelism >= 4); more use a larger one. +_SMALL_LOCAL_CAP = 224 + + +@functools.cache +def _is_supported_device(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3), (10, 7)) + + +def supports_device(device: torch.device) -> bool: + """Whether the routing kernels run on ``device`` (SM100, SM103 or SM107).""" + return device.type == "cuda" and _is_supported_device(_device_index(device)) + + +def default_config(m: int) -> tuple[int, int, int]: + """``(num_ctas, warps_per_token, num_warps)`` for ``m`` tokens (at most 32 tokens per CTA). + + Measured on GB300 (standalone, behind a producer kernel): eight 8-warp CTAs with two-warp + teams up to 32 tokens, one-warp teams up to 64; 16-warp CTAs above. + """ + if m <= 32: + return MAX_CTAS, 2, 8 + if m <= 64: + return MAX_CTAS, 1, 8 + return MAX_CTAS, 1, 16 + + +@functools.cache +def _compiled_tail( + num_ctas: int, + num_warps: int, + warps_per_token: int, + num_splits: int, + with_sort: bool, + use_pdl: bool, + local_cap: int, +): + kernel = K3RouteTailKernel( + num_ctas, num_warps, warps_per_token, num_splits, with_sort, use_pdl, local_cap + ) + dummy = Int64(1 << 20) + return cute.compile( + kernel, + Int32(num_ctas), + Int32(num_ctas), + *([dummy] * 10), + Int32(0), + Int32(NUM_EXPERTS), + Int64(0), + Int64(0), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def _check_inputs(logits: torch.Tensor, e_score_correction_bias: torch.Tensor) -> None: + if logits.dtype != torch.float32 or logits.dim() not in (2, 3): + raise ValueError("logits must be a 2D [M, 896] or 3D [S, R, 896] float32 tensor") + if logits.shape[-1] != NUM_EXPERTS or not logits.is_contiguous(): + raise ValueError(f"logits must be contiguous with {NUM_EXPERTS} columns") + if ( + e_score_correction_bias.dtype != torch.float32 + or e_score_correction_bias.numel() != NUM_EXPERTS + or not e_score_correction_bias.is_contiguous() + ): + raise ValueError(f"e_score_correction_bias must be a contiguous float32 [{NUM_EXPERTS}]") + for t in (logits, e_score_correction_bias): + if not t.is_cuda or t.device != logits.device or t.data_ptr() % 16: + raise ValueError("inputs must be 16-byte aligned tensors on the same CUDA device") + + +def can_use_k3_route_tail( + num_tokens: int, logits: torch.Tensor, local_num_experts: int | None = None +) -> bool: + """Whether the routing tail serves ``num_tokens`` tokens of these logits (non-raising). + + Above :data:`MAX_M` tokens the sort needs :func:`wide_sort_fits` (``local_num_experts=None`` + skips that check). + """ + if num_tokens > MAX_M and local_num_experts is not None: + if not wide_sort_fits(num_tokens, local_num_experts): + return False + return ( + logits.is_cuda + and _is_supported_device(logits.device.index) + and 1 <= num_tokens <= WIDE_MAX_M + and logits.dtype == torch.float32 + and logits.shape[-1] == NUM_EXPERTS + ) + + +def run_k3_route_tail( + logits: torch.Tensor, + num_tokens: int, + e_score_correction_bias: torch.Tensor, + routed_scaling_factor: float, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + sort_outputs: tuple[torch.Tensor, ...] | None, + local_expert_offset: int, + local_num_experts: int, + *, + logits_out: torch.Tensor | None = None, + num_ctas: int | None = None, + warps_per_token: int | None = None, + num_warps: int | None = None, + use_pdl: bool = True, +) -> None: + """Launch the routing tail on the current stream (inputs are not re-validated). + + Args: + logits: FP32 ``[M, 896]`` logits or ``[S, R, 896]`` split-K partials (``R >= M``). + num_tokens: ``M``, ``1 <= M <= 256``. + e_score_correction_bias: ``fp32 [896]``. + routed_scaling_factor: Multiplier of the renormalized weights. + topk_ids: Output ``int32 [M, 16]``. + topk_weights: Output ``fp32 [M, 16]``. + sort_outputs: ``moe_sort`` outputs ``(tile_idx_to_expert_idx, tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, + total_num_padded_tokens, num_non_exiting_tiles)`` (``int32``), or ``None`` for + top-k only. + local_expert_offset: First expert of this rank. + local_num_experts: Experts of this rank. + logits_out: Optional output ``fp32 [M, 896]``: the (summed) logits. + num_ctas: Cluster size (default from :func:`default_config`). + warps_per_token: Warps routing one token (default from :func:`default_config`). + num_warps: Warps per CTA (default from :func:`default_config`). + use_pdl: Launch with programmatic dependent launch. + """ + m = num_tokens + default_ctas, default_team, default_warps = default_config(m) + if num_ctas is None: + num_ctas = default_ctas + if warps_per_token is None: + warps_per_token = default_team + if num_warps is None: + num_warps = default_warps + if (m + num_ctas - 1) // num_ctas > _MAX_TOKENS_PER_CTA: + raise ValueError(f"M={m} needs more than {num_ctas} CTAs") + splits = 1 if logits.dim() == 2 else logits.shape[0] + rows = logits.shape[0] if logits.dim() == 2 else logits.shape[1] + with_sort = sort_outputs is not None + local_cap = _SMALL_LOCAL_CAP if local_num_experts <= _SMALL_LOCAL_CAP else NUM_EXPERTS + compiled = _compiled_tail( + num_ctas, num_warps, warps_per_token, splits, with_sort, use_pdl, local_cap + ) + if with_sort: + tile_expert, tile_limit, e2p, p2e, total, num_tiles = (t.data_ptr() for t in sort_outputs) + else: + tile_expert = tile_limit = e2p = p2e = total = num_tiles = 0 + scale_bits = struct.unpack("= 14 (EP16: 56). +_WIDE_SMALL_LOCAL_CAP = 64 +#: Largest token count routed with 8-token CTAs by default (see :func:`default_wide_config`). +_WIDE_COMPACT_MAX_M = 1024 + + +def wide_table_fits(m: int, local_num_experts: int, tokens_per_cta: int) -> bool: + """Whether the multi-CTA sort's shared-memory tables hold ``m`` tokens over + ``local_num_experts`` local experts in CTAs of ``tokens_per_cta`` tokens.""" + ctas = (m + tokens_per_cta - 1) // tokens_per_cta + if tokens_per_cta == 8: + row_bytes = 4 * ((local_num_experts + 3) // 4) + table = (ctas * row_bytes + 15) // 16 * 16 + return table + (ctas + 7) // 8 * local_num_experts * 2 <= WIDE_COMPACT_TABLE_BYTES + return local_num_experts * ctas <= WIDE_TABLE_CELLS + + +def wide_sort_fits(m: int, local_num_experts: int) -> bool: + """Whether the multi-CTA sort serves ``m`` tokens over ``local_num_experts`` local experts + (see :func:`wide_table_fits`; CTAs of 8 or 16 tokens).""" + return ( + 1 <= m <= WIDE_MAX_M + and 1 <= local_num_experts <= WIDE_LOCAL_CAP + and (wide_table_fits(m, local_num_experts, 8) or wide_table_fits(m, local_num_experts, 16)) + ) + + +def default_wide_config(m: int, local_num_experts: int = 56) -> tuple[int, int, int]: + """``(tokens_per_cta, num_warps, warps_per_token)`` of the multi-CTA routing for ``m`` + tokens (``MAX_M < m <= WIDE_MAX_M``). + + Measured on GB300 (EP16): up to 1024 tokens, 8-warp CTAs of 8 tokens, whose compact sort + tables let them share SMs with the persistent latent_down GEMM running beside the routing + (M = 896: 23.6 instead of 28.7 us beside it, the same alone). Above, and whenever those + tables do not fit, 16-warp CTAs of 16 tokens, one token per warp (M = 1344 alone: 9.5 + instead of 11.3 us, where the last CTA sorts the masks of 84 instead of 168 CTAs). + """ + if m <= _WIDE_COMPACT_MAX_M and wide_table_fits(m, local_num_experts, 8): + return 8, 8, 1 + return 16, 16, 1 + + +@functools.cache +def _compiled_wide( + tokens_per_cta: int, + num_warps: int, + warps_per_token: int, + num_splits: int, + with_sort: bool, + use_pdl: bool, + local_cap: int, +): + kernel = K3RouteWideKernel( + tokens_per_cta, num_warps, warps_per_token, num_splits, with_sort, use_pdl, local_cap + ) + dummy = Int64(1 << 20) + return cute.compile( + kernel, + Int32(MAX_M + 1), + Int32(MAX_M + 1), + *([dummy] * 10), + Int32(0), + Int32(local_cap), + Int64(0), + dummy, + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +_WIDE_WORKSPACES: dict[int, torch.Tensor] = {} + + +def _device_index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() + + +def wide_workspace(device: torch.device) -> torch.Tensor: + """Persistent per-device state of the multi-CTA sort (arrival counter + mask slots). + + Between launches the counter is zero; each launch restores that. It must be created + outside CUDA graph capture (a capture would allocate it from the graph's private pool). + """ + index = _device_index(device) + workspace = _WIDE_WORKSPACES.get(index) + if workspace is None: + workspace = torch.zeros( + _WIDE_WORKSPACE_BYTES // 4, dtype=torch.int32, device=torch.device("cuda", index) + ) + _WIDE_WORKSPACES[index] = workspace + return workspace + + +def wide_workspace_ready(device: torch.device) -> bool: + """Whether this device's multi-CTA sort workspace exists (see :func:`wide_workspace`).""" + return _device_index(device) in _WIDE_WORKSPACES + + +def run_k3_route_wide( + logits: torch.Tensor, + num_tokens: int, + e_score_correction_bias: torch.Tensor, + routed_scaling_factor: float, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + sort_outputs: tuple[torch.Tensor, ...] | None, + local_expert_offset: int, + local_num_experts: int, + *, + tokens_per_cta: int | None = None, + num_warps: int | None = None, + warps_per_token: int | None = None, + use_pdl: bool = True, + workspace: torch.Tensor | None = None, +) -> None: + """Launch the multi-CTA routing on the current stream (inputs are not re-validated). + + Same arguments and outputs as :func:`run_k3_route_tail` (without ``logits_out``) for + ``1 <= num_tokens <= WIDE_MAX_M`` and, with ``sort_outputs``, ``local_num_experts <= + WIDE_LOCAL_CAP`` and :func:`wide_table_fits`. The configuration defaults come from + :func:`default_wide_config`. ``workspace``: the device's :func:`wide_workspace` (default: + looked up here). + """ + m = num_tokens + default_tokens, default_warps, default_team = default_wide_config(m, local_num_experts) + if tokens_per_cta is None: + tokens_per_cta = default_tokens + if num_warps is None: + num_warps = default_warps + if warps_per_token is None: + warps_per_token = default_team + if not 1 <= m <= WIDE_MAX_M: + raise ValueError(f"num_tokens={m} must be in [1, {WIDE_MAX_M}]") + with_sort = sort_outputs is not None + if with_sort and ( + local_num_experts > WIDE_LOCAL_CAP + or not wide_table_fits(m, local_num_experts, tokens_per_cta) + ): + raise ValueError( + f"the multi-CTA sort serves at most {WIDE_LOCAL_CAP} local experts and " + "tables within wide_table_fits" + ) + splits = 1 if logits.dim() == 2 else logits.shape[0] + rows = logits.shape[0] if logits.dim() == 2 else logits.shape[1] + local_cap = ( + _WIDE_SMALL_LOCAL_CAP if local_num_experts <= _WIDE_SMALL_LOCAL_CAP else WIDE_LOCAL_CAP + ) + compiled = _compiled_wide( + tokens_per_cta, num_warps, warps_per_token, splits, with_sort, use_pdl, local_cap + ) + if with_sort: + tile_expert, tile_limit, e2p, p2e, total, num_tiles = (t.data_ptr() for t in sort_outputs) + if workspace is None: + workspace = wide_workspace(logits.device) + workspace = workspace.data_ptr() + else: + tile_expert = tile_limit = e2p = p2e = total = num_tiles = workspace = 0 + scale_bits = struct.unpack(" tuple[torch.Tensor, ...]: + """Fresh ``(topk_ids, topk_weights, *moe_sort outputs)`` for ``m`` tokens.""" + device = logits.device + max_tiles, max_permuted = sort_buffer_sizes(m, local_num_experts, tile_tokens_dim) + i32 = dict(dtype=torch.int32, device=device) + return ( + torch.empty((m, TOP_K), **i32), + torch.empty((m, TOP_K), dtype=torch.float32, device=device), + torch.empty((max_tiles,), **i32), + torch.empty((max_tiles,), **i32), + torch.empty((m, TOP_K), **i32), + torch.empty((max_permuted,), **i32), + torch.empty((1,), **i32), + torch.empty((1,), **i32), + ) + + +def k3_route_tail( + logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + routed_scaling_factor: float, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int = TILE, + num_tokens: int = -1, + workspace: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """``noaux_tc`` top-16 + ``moe_sort`` of FP32 router logits in one launch. + + Args: + logits: ``fp32 [M, 896]`` router logits, or ``[S, R, 896]`` split-K partials whose + first ``M`` rows are summed in split order. + e_score_correction_bias: ``fp32 [896]``. + routed_scaling_factor: Multiplier of the renormalized weights. + local_expert_offset: First expert of this rank. + local_num_experts: Experts of this rank. + tile_tokens_dim: Must be 128. + num_tokens: ``M`` (default: all rows of ``logits``), ``1 <= M <= 2048``; above 256 + tokens (multi-CTA kernel) see :func:`wide_sort_fits`. + workspace: Above 256 tokens, the device's persistent multi-CTA sort workspace + (:func:`wide_workspace`; default: looked up here). Unused at or below 256 tokens. + + Returns: + ``(topk_ids int32 [M, 16], topk_weights fp32 [M, 16], tile_idx_to_expert_idx, + tile_idx_to_mn_limit, expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, + total_num_padded_tokens, num_non_exiting_tiles)``. Top-k equals ``noaux_tc_op`` on the + same logits bit for bit; the sort outputs have ``moe_sort``'s shapes and dtypes, only + valid entries are written (as in ``moe_sort``), and the rows of an expert are in + ascending token order. + """ + _check_inputs(logits, e_score_correction_bias) + if tile_tokens_dim != TILE: + raise ValueError(f"tile_tokens_dim must be {TILE}") + rows = logits.shape[0] if logits.dim() == 2 else logits.shape[1] + m = rows if num_tokens < 0 else num_tokens + if not 1 <= m <= min(rows, WIDE_MAX_M): + raise ValueError(f"num_tokens={m} must be in [1, {min(rows, WIDE_MAX_M)}]") + if not ( + 0 <= local_expert_offset + and 1 <= local_num_experts + and local_expert_offset + local_num_experts <= NUM_EXPERTS + ): + raise ValueError("invalid local expert range") + if m > MAX_M and not wide_sort_fits(m, local_num_experts): + raise ValueError( + f"above {MAX_M} tokens the sort serves at most {WIDE_LOCAL_CAP} local experts and " + "tables within wide_sort_fits" + ) + outs = _route_tail_outputs( + logits, m, e_score_correction_bias, local_num_experts, tile_tokens_dim + ) + if m <= MAX_M: + run_k3_route_tail( + logits, + m, + e_score_correction_bias, + routed_scaling_factor, + outs[0], + outs[1], + outs[2:], + local_expert_offset, + local_num_experts, + ) + else: + run_k3_route_wide( + logits, + m, + e_score_correction_bias, + routed_scaling_factor, + outs[0], + outs[1], + outs[2:], + local_expert_offset, + local_num_experts, + workspace=workspace, + ) + return outs + + +__all__ = [ + "MAX_M", + "WIDE_LOCAL_CAP", + "WIDE_MAX_M", + "WIDE_COMPACT_TABLE_BYTES", + "WIDE_TABLE_CELLS", + "K3RouteTailKernel", + "K3RouteWideKernel", + "can_use_k3_route_tail", + "default_config", + "default_wide_config", + "k3_route_tail", + "run_k3_route_tail", + "run_k3_route_wide", + "supports_device", + "wide_sort_fits", + "wide_table_fits", + "wide_workspace", + "wide_workspace_ready", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_router_exact.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_router_exact.py new file mode 100644 index 000000000000..b4bd70ee19b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_router_exact.py @@ -0,0 +1,430 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 router logits (SM100/SM103/SM107, ``M <= 16`` tokens), bit-identical to +``trtllm::dsv3_router_gemm_op``, with the router weight read before ``griddepcontrol.wait``. + +``dsv3_router_gemm_op`` serves ``M <= 16`` tokens of the 896 x 7168 BF16 router with +``router_gemm_mma_kernel``: CTA ``b`` owns experts ``[8 b, 8 b + 8)``, its warp ``s`` of 14 owns +the K slice ``[512 s, 512 s + 512)`` and chains 32 ``mma.sync.m16n8k16`` from a zero accumulator +(two per 32-wide chunk: lane ``(g, q)`` holds k ``chunk + 8 q .. + 8`` of activation rows ``g`` and +``g + 8`` and of expert ``g``, fed as k pairs ``{0,1},{2,3}`` then ``{4,5},{6,7}``), and the CTA sums +the 14 slice partials in slice order starting from 0. + +:class:`K3RouterExactKernel` computes the same 14 slice partials with the same fragments, MMAs +and order, and sums them the same way, so every logit has the same bits. The work is cut +differently: a cluster of 2 CTAs owns 16 experts (two n8 tiles); CTA ``k`` computes slices +``7 k .. 7 k + 6`` (warp ``w`` = slice ``7 k + w``, one MMA chain per tile), sums its seven partials +in slice order, and CTA 0 hands its running sum to CTA 1 over distributed shared memory, which +adds slices 7..13 in order and stores the logits. A CTA's 112 KB weight slice is copied into shared +memory by TMA (128-byte swizzle, bank-conflict-free fragment loads) **before** +``griddepcontrol.wait``: with PDL the copy runs while the previous kernel finishes. Each warp loads +its slice of the activation (16 KB) straight into registers in one round trip after the wait. + +PDL: the kernel triggers its dependents right after its own ``griddepcontrol.wait``; before the +wait it reads only the router weight. +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +from .k3_fused_router import ( + _EVICT_FIRST, + HIDDEN_SIZE, + NUM_EXPERTS, + _asm, + _fv, + _iv, + _mma_bf16, + _st_global_f32, + _unpack, +) +from .k3_route_tail import _griddepcontrol_launch_dependents, _griddepcontrol_wait + +#: Largest token count (one m16 tile, like ``router_gemm_mma_kernel``). +MAX_M = 16 +#: K slices of ``router_gemm_mma_kernel`` (its warps) and their width. +NUM_SLICES = 14 +SLICE_K = HIDDEN_SIZE // NUM_SLICES +#: 32-wide K chunks per slice. +CHUNKS = SLICE_K // 32 +#: Cluster: 2 CTAs x 7 slices; 2 n8 expert tiles (16 experts) per cluster. +CLUSTER = 2 +SLICES_PER_CTA = NUM_SLICES // CLUSTER +TILES = 2 +NUM_WARPS = SLICES_PER_CTA +THREADS = 32 * NUM_WARPS +CLUSTER_EXPERTS = 8 * TILES +NUM_CTAS = (NUM_EXPERTS // CLUSTER_EXPERTS) * CLUSTER +#: Weight TMA box: 8 rows x 64 BF16 (one 1 KB 128-byte-swizzle atom). +_BOX_K = 64 +_BOXES_PER_SLICE = SLICE_K // _BOX_K +_W_BOX_BYTES = 8 * _BOX_K * 2 +#: One fragment set of the two tiles: 2 tiles x 32 lanes x 4 floats. +_FRAG_BYTES = TILES * 32 * 4 * 4 + +assert HIDDEN_SIZE % (NUM_SLICES * 32) == 0 and NUM_SLICES % CLUSTER == 0 +assert NUM_EXPERTS % CLUSTER_EXPERTS == 0 + + +@dsl_user_op +def _lds_v4_b32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _st_async_v4_f32(local_addr, v0, v1, v2, v3, local_mbar, peer, *, loc=None, ip=None): + """Store four FP32 values into a peer CTA's SMEM, crediting its mbarrier.""" + _asm( + None, + [_iv(local_addr, loc=loc, ip=ip)] + + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)] + + [_iv(local_mbar, loc=loc, ip=ip), _iv(peer, loc=loc, ip=ip)], + "{\n" + ".reg .b32 ra, rb;\n" + "mapa.shared::cluster.u32 ra, $0, $6;\n" + "mapa.shared::cluster.u32 rb, $5, $6;\n" + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [ra], {$1, $2, $3, $4}, [rb];\n" + "}", + "r,f,f,f,f,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ldg_v4_b32_if(valid, addr, *, loc=None, ip=None): + """16 bytes of global memory as four 32-bit words if ``valid != 0``, else zeros.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_iv(valid, loc=loc, ip=ip), Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n" + ".reg .pred p;\n" + "setp.ne.s32 p, $4, 0;\n" + "mov.b32 $0, 0;\nmov.b32 $1, 0;\nmov.b32 $2, 0;\nmov.b32 $3, 0;\n" + "@p ld.global.v4.b32 {$0, $1, $2, $3}, [$5];\n" + "}", + "=r,=r,=r,=r,r,l", + loc=loc, + ip=ip, + ) + return _unpack(res, T.i32(), Int32, 4, loc=loc, ip=ip) + + +@dsl_user_op +def _sts_v4_f32(addr, v0, v1, v2, v3, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)] + [_fv(v, loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", + "r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + loc=loc, + ip=ip, + ) + return _unpack(res, T.f32(), Float32, 4, loc=loc, ip=ip) + + +class K3RouterExactKernel: + """``logits[M, 896]``, bit-identical to ``router_gemm_mma_kernel`` (see the module docstring).""" + + def __init__(self, use_pdl: bool): + self.use_pdl = use_pdl + + @cute.jit + def __call__( + self, + w_ptr: cute.Pointer, + x_addr: Int64, + m: Int32, + out_addr: Int64, + stream: cuda.CUstream, + ): + mW = cute.make_tensor( + w_ptr, cute.make_layout((NUM_EXPERTS, HIDDEN_SIZE), stride=(HIDDEN_SIZE, 1)) + ) + atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + # Weight boxes [8 rows, 64 k, (slice, tile, box)]. + w_boxes = SLICES_PER_CTA * TILES * _BOXES_PER_SLICE + w_layout = cute.tile_to_shape(atom, (8, _BOX_K, w_boxes), order=(0, 1, 2)) + tma_atom_w, tma_w = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mW, + cute.tile_to_shape(atom, (8, _BOX_K), order=(0, 1)), + (8, _BOX_K), + ) + + @cute.struct + class SharedStorage: + w_full: cute.struct.MemRange[cutlass.Int64, NUM_WARPS] + mbox_full: cutlass.Int64 + sW: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, cute.cosize(w_layout.outer)], 1024 + ] + # Slice partials [warp][tile][lane][4], then the peer's running sum [tile][lane][4]. + red: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, NUM_WARPS * _FRAG_BYTES // 4], 16 + ] + mbox: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, _FRAG_BYTES // 4], 16] + + self.shared_storage = SharedStorage + self.kernel(tma_atom_w, tma_w, w_layout, x_addr, m, out_addr).launch( + grid=(NUM_CTAS, 1, 1), + block=(THREADS, 1, 1), + cluster=(CLUSTER, 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + tma_atom_w: cute.CopyAtom, + mWt: cute.Tensor, + w_layout: cute.ComposedLayout, + x_addr: Int64, + m: Int32, + out_addr: Int64, + ): + tidx, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + group = cta // CLUSTER + storage = utils.SmemAllocator().allocate(self.shared_storage) + w_full = storage.w_full.data_ptr() + mbox_full = storage.mbox_full + w_base = storage.sW.data_ptr().toint() + red_base = storage.red.data_ptr().toint() + mbox_base = storage.mbox.data_ptr().toint() + s = rank * SLICES_PER_CTA + warp + + with cute.arch.elect_one(): + cute.arch.mbarrier_init(w_full + warp, 1) + if warp == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mbox_full, 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + # The peer writes into this CTA's mailbox only after both CTAs initialized their barriers. + cute.arch.cluster_arrive_relaxed() + + # This warp's slice of the cluster's two weight tiles, before the wait. + sW = storage.sW.get_tensor(w_layout.outer, swizzle=w_layout.inner) + gW = cute.local_tile(mWt, (8, _BOX_K), (None, None)) + tWs, tWg = cpasync.tma_partition( + tma_atom_w, + 0, + cute.make_layout(1), + cute.group_modes(sW, 0, 2), + cute.group_modes(gW, 0, 2), + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + w_full + warp, TILES * _BOXES_PER_SLICE * _W_BOX_BYTES + ) + for t in cutlass.range_constexpr(TILES): + for j in cutlass.range_constexpr(_BOXES_PER_SLICE): + cute.copy( + tma_atom_w, + tWg[(None, group * TILES + t, s * _BOXES_PER_SLICE + j)], + tWs[(None, (warp * TILES + t) * _BOXES_PER_SLICE + j)], + tma_bar_ptr=w_full + warp, + cache_policy=Int64(_EVICT_FIRST), + ) + + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + # Safe trigger point: everything up to the previous grid is complete. + _griddepcontrol_launch_dependents() + + # router_gemm_mma_kernel fragments: lane (g, q) holds activation rows g and g + 8 and + # weight row g at k = chunk + 8 q .. + 8. The whole slice of the activation is loaded + # into registers at once. + g = lane >> 2 + q = lane & 3 + row0 = Int32(0) + if g < m: + row0 = Int32(1) + row1 = Int32(0) + if g + 8 < m: + row1 = Int32(1) + a0_base = x_addr + Int64(g * (HIDDEN_SIZE * 2) + (s * SLICE_K + q * 8) * 2) + a1_base = a0_base + Int64(8 * HIDDEN_SIZE * 2) + acts = [] + for cc in cutlass.range_constexpr(CHUNKS): + a_off = Int64(cc * 64) + acts.append( + (_ldg_v4_b32_if(row0, a0_base + a_off), _ldg_v4_b32_if(row1, a1_base + a_off)) + ) + w_warp = w_base + warp * (TILES * _BOXES_PER_SLICE * _W_BOX_BYTES) + cute.arch.mbarrier_wait(w_full + warp, 0) + accs = [[Float32(0.0)] * 4 for _ in range(TILES)] + nxt = [self._wfrag(w_warp, g, q, t, 0) for t in range(TILES)] + for cc in cutlass.range_constexpr(CHUNKS): + cur = nxt + if cc + 1 < CHUNKS: + nxt = [self._wfrag(w_warp, g, q, t, cc + 1) for t in range(TILES)] + a0, a1 = acts[cc] + for t in cutlass.range_constexpr(TILES): + bw = cur[t] + acc = _mma_bf16((a0[0], a1[0], a0[1], a1[1]), bw[0], bw[1], accs[t]) + accs[t] = list(_mma_bf16((a0[2], a1[2], a0[3], a1[3]), bw[2], bw[3], acc)) + for t in cutlass.range_constexpr(TILES): + acc = accs[t] + _sts_v4_f32( + red_base + warp * _FRAG_BYTES + (t * 32 + lane) * 16, acc[0], acc[1], acc[2], acc[3] + ) + cute.arch.barrier() + + # Thread t < 64 owns fragment (tile t / 32, lane t % 32): the running sum over this + # CTA's slices in order (CTA 0 starts from 0, CTA 1 from CTA 0's sum). + cute.arch.cluster_wait() + if rank == 1: + if tidx == 0: + # One arrival plus CTA 0's bytes complete the mailbox phase. + cute.arch.mbarrier_arrive_and_expect_tx(mbox_full, _FRAG_BYTES) + if tidx < TILES * 32: + off = tidx * 16 + tot = [Float32(0.0)] * 4 + if rank == 1: + cute.arch.mbarrier_wait(mbox_full, 0) + tot = list(_lds_v4_f32(mbox_base + off)) + for w in cutlass.range_constexpr(NUM_WARPS): + v = _lds_v4_f32(red_base + w * _FRAG_BYTES + off) + for i in cutlass.range_constexpr(4): + tot[i] = tot[i] + v[i] + if rank == 0: + _st_async_v4_f32( + mbox_base + off, tot[0], tot[1], tot[2], tot[3], mbox_full.toint(), Int32(1) + ) + else: + t_idx = tidx >> 5 + ln = tidx & 31 + for i in cutlass.range_constexpr(4): + row = (ln >> 2) + (i >> 1) * 8 + col = group * CLUSTER_EXPERTS + t_idx * 8 + (ln & 3) * 2 + (i & 1) + if row < m: + _st_global_f32(out_addr + Int64((row * NUM_EXPERTS + col) * 4), tot[i]) + # No CTA leaves while its peer may still write into its mailbox. + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + @cute.jit + def _wfrag(self, w_warp, g, q, t: int, cc: int): + """Weight row ``g`` of tile ``t`` at chunk ``cc`` of this warp's slice (swizzled box).""" + box = t * _BOXES_PER_SLICE + cc // 2 + unit = (4 * (cc % 2) + q) ^ g + return _lds_v4_b32(w_warp + box * _W_BOX_BYTES + g * 128 + (unit << 4)) + + +@functools.cache +def _compiled(use_pdl: bool): + kernel = K3RouterExactKernel(use_pdl) + return cute.compile( + kernel, + make_ptr(cutlass.BFloat16, 1 << 20, cute.AddressSpace.gmem, assumed_align=16), + Int64(1 << 20), + Int32(MAX_M), + Int64(1 << 20), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +@functools.cache +def _is_supported_device(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3), (10, 7)) + + +def can_run(hidden_states: torch.Tensor, weight: torch.Tensor) -> bool: + """Whether :func:`run_k3_router_exact` serves these inputs (non-raising check).""" + return ( + hidden_states.is_cuda + and _is_supported_device(hidden_states.device.index) + and hidden_states.dtype == torch.bfloat16 + and hidden_states.dim() == 2 + and hidden_states.is_contiguous() + and hidden_states.shape[1] == HIDDEN_SIZE + and 1 <= hidden_states.shape[0] <= MAX_M + and hidden_states.data_ptr() % 16 == 0 + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (NUM_EXPERTS, HIDDEN_SIZE) + and weight.is_contiguous() + and weight.data_ptr() % 16 == 0 + ) + + +def run_k3_router_exact( + hidden_states: torch.Tensor, + weight: torch.Tensor, + logits: torch.Tensor, + *, + use_pdl: bool = True, +) -> None: + """Launch on the current stream: ``logits`` (``fp32 [M, 896]``) = + ``dsv3_router_gemm_op(hidden_states, weight.t())`` bit for bit.""" + m = hidden_states.shape[0] + _compiled(use_pdl)( + make_ptr(cutlass.BFloat16, weight.data_ptr(), cute.AddressSpace.gmem, assumed_align=16), + Int64(hidden_states.data_ptr()), + Int32(m), + Int64(logits.data_ptr()), + cuda.CUstream(torch.cuda.current_stream(hidden_states.device).cuda_stream), + ) + + +def k3_router_logits(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """FP32 ``[M, 896]`` router logits of ``M <= 16`` tokens, bit-identical to + ``dsv3_router_gemm_op(hidden_states, weight.t())``; the weight is copied into shared memory + before ``griddepcontrol.wait``.""" + if not can_run(hidden_states, weight): + raise ValueError("unsupported router inputs") + logits = torch.empty( + (hidden_states.shape[0], NUM_EXPERTS), dtype=torch.float32, device=hidden_states.device + ) + run_k3_router_exact(hidden_states, weight, logits) + return logits + + +__all__ = [ + "MAX_M", + "K3RouterExactKernel", + "can_run", + "k3_router_logits", + "run_k3_router_exact", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_shared_persist.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_shared_persist.py new file mode 100644 index 000000000000..3cd331be13e2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_shared_persist.py @@ -0,0 +1,701 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 shared expert of at most 8 rows on a few SMs beside the routed experts (``N`` CTAs, +the column-parallel MoE tail's ``_K3_ULL_SHARED_PERSIST_CTAS``). + +The shared expert of a decode step is a TP-sharded gated MLP: this rank computes +``act = situ(x @ Wg^T, x @ Wu^T)`` for its ``I = 6144 / W`` intermediate columns and the +row-parallel partial ``act @ Wd^T`` (``[M, 7168]``), which the column-parallel MoE tail pushes +into slot ``(token, rank)`` of the tail's Lamport call (``k3_lamport_ar`` one-shot layout) for the +next pre-attention mix to sum. :class:`K3SharedPersistKernel` does all of it with ``N`` persistent +CTAs, one per SM the routed experts' grid leaves free (the experts run ``SMs - N`` CTAs), so the +shared expert finishes inside the experts' window instead of competing with the tail: + +* Weights. Each CTA streams its weights once, in consumption order, through a ring of shared + memory stages (``cp.async.bulk``, one producer warp, full / empty mbarriers, a proxy fence + between the compute warps' reads of a slot and its refill), from its launch: + the gate/up rows of its ``I / N`` intermediate columns (gate rows then up rows, ``2 I / N`` rows + x 7168), then the down rows of its 16-row m-tiles (m-tile ``T`` of the 448 goes to CTA + ``T mod N``). The fragments are prepared once on the host (:func:`persist_fragments`) in the + ``mma.sync`` A order, so a lane reads its 16 bytes of a fragment with one ``ld.shared.v4``. +* Gate/up. Eight compute warps split K: warp ``w`` takes K steps ``8 j + w`` (``j < 56``) of + every m-tile of the CTA's gate/up rows (BF16 ``mma.sync`` m16n8k16, FP32 accumulation; ``x`` as + the B operand in registers, tokens ``>= M`` zero). The eight warps' partials are summed in warp + order (shared memory), rounded to BF16 (the gate/up GEMM's output type) and the SiTU activation + (``beta tanh(g / beta) sigmoid(g) * lb tanh(u / lb)``, FP32, rounded to BF16) gives the CTA's + activation columns. +* Exchange. Each CTA writes its columns into a global ``[8, I]`` BF16 buffer and counts itself + (release); every CTA then waits for all ``N`` (acquire) and copies the whole activation into + shared memory; the last CTA to have copied it re-arms the two counters. +* Down. Warp ``w`` multiplies m-tile ``8 g + w`` of the CTA's list (``K = I``), rounds the 16 x M + outputs to BF16 (``-0`` as ``+0``) and stores them into slot ``(token, rank)`` of the tail call's + buffer ``cur`` (multicast to every rank), where the pre-attention mix reads them as it reads + the push of today's shared output. + +Numerics: the products are those of the BF16 GEMMs (FP32 accumulation, BF16 outputs) in a +different summation order, and the activation evaluates tanh / sigmoid with ``ex2.approx``: not +bitwise; the unit test bounds the difference against the GEMM + activation-op path. + +Ordering: launched without PDL on the shared stream after the layer's input is final; the tail +call's rotation state is read at entry (the previous call on that workspace was finished by the +previous layer's pre-attention mix, which precedes the fork). Consecutive launches on a stream +reuse the exchange buffer and counters (the stream serializes them). +""" + +from __future__ import annotations + +import functools +from typing import Optional + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int32, Int64 +from cutlass.cutlass_dsl import T, dsl_user_op + +from .k3_lamport_ar import ( + _asm, + _fv, + _iv, + _lv, + atom_add_acq_rel_u32, + ld_acquire_u32, + ld_volatile_v4, + multimem_st_v4, + pack_bf16x2, + sanitize_bf16x2, + st_u32, + st_v4, +) +from .k3_moe_tail1 import ( + _bar_sync, + _bulk_g2s, + _fence_mbar_init, + _lds_b32, + _lds_f32, + _lds_v4, + _lds_v4_f32, + _mbar_expect_tx, + _mbar_init, + _mma_bf16, + _sts_f32, + _sts_v4, +) + +__all__ = [ + "HIDDEN", + "INTER", + "MAX_M", + "K3SharedPersistKernel", + "PersistWorkspace", + "can_run", + "persist_fragments", + "run_k3_shared_persist", +] + +HIDDEN = 7168 +INTER = 6144 +MAX_M = 8 +#: Compute warps (K split of gate/up, m-tiles of down) and the producer warp. +WARPS = 8 +THREADS = (WARPS + 1) * 32 +#: Bytes of one ``mma.sync`` A fragment (32 lanes x 16 bytes): one (m-tile, K step). +FRAG = 512 +#: Gate/up K steps per warp (K = HIDDEN, eight warps). +GU_STEPS = HIDDEN // 16 // WARPS +#: 16-row m-tiles of the down projection's output. +DOWN_TILES = HIDDEN // 16 +#: Shared memory the kernel may use (bytes). +SMEM_BUDGET = 220 * 1024 +#: Words per token of the gate/up partials in shared memory (16 rows). +RED_PITCH = 16 +#: FP32 words per token row of a warp's staged down tile (16 rows + padding). +STAGE_PITCH = 20 + + +@dsl_user_op +def _fence_proxy_async(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.shared::cta;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _mbar_arrive(addr, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip)], + "mbarrier.arrive.shared::cta.b64 _, [$0];", + "r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mbar_try_wait(addr, parity, *, loc=None, ip=None): + """1 once the phase of parity ``parity`` of the mbarrier at ``addr`` has completed.""" + return Int32( + _asm( + T.i32(), + [_iv(addr, loc=loc, ip=ip), _iv(parity, loc=loc, ip=ip)], + "{\n.reg .pred p;\nmbarrier.try_wait.parity.shared::cta.b64 p, [$1], $2;\n" + "selp.u32 $0, 1, 0, p;\n}", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fence_acq_rel_gpu(*, loc=None, ip=None): + _asm(None, [], "fence.acq_rel.gpu;", "", loc=loc, ip=ip) + + +@dsl_user_op +def _ld_b32(addr, *, loc=None, ip=None): + """A 32-bit word of the layer input (complete before the launch).""" + return Int32( + _asm( + T.i32(), + [_lv(addr, loc=loc, ip=ip)], + "ld.global.nc.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) # fmt: skip + ) + + +@dsl_user_op +def _tanh(x, *, loc=None, ip=None): + """``tanh(x)`` in FP32: an odd Taylor polynomial below ``|x| = 0.25``, else + ``(1 - e) / (1 + e)`` with ``e = 2^(-2 |x| log2(e))`` (``ex2.approx``) and the sign.""" + return Float32( + _asm( + T.f32(), + [_fv(x, loc=loc, ip=ip)], + "{\n.reg .f32 a, x2, p, e, n, d, t;\n.reg .pred small, neg;\n" + "abs.f32 a, $1;\n" + "setp.lt.f32 small, a, 0f3E800000;\n" + "setp.lt.f32 neg, $1, 0f00000000;\n" + # polynomial: x (1 + x2 (-1/3 + x2 (2/15 + x2 (-17/315 + x2 62/2835)))) + "mul.rn.f32 x2, $1, $1;\n" + "fma.rn.f32 p, x2, 0f3CB327A4, 0fBD5D0DD1;\n" + "fma.rn.f32 p, p, x2, 0f3E088889;\n" + "fma.rn.f32 p, p, x2, 0fBEAAAAAB;\n" + "mul.rn.f32 p, p, x2;\n" + "fma.rn.f32 p, p, $1, $1;\n" + # exponential form + "mul.rn.f32 e, a, 0fC038AA3B;\n" + "ex2.approx.f32 e, e;\n" + "sub.rn.f32 n, 0f3F800000, e;\n" + "add.rn.f32 d, 0f3F800000, e;\n" + "div.rn.f32 t, n, d;\n" + "@neg neg.f32 t, t;\n" + "selp.f32 $0, p, t, small;\n}", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sigmoid(x, *, loc=None, ip=None): + """``1 / (1 + exp(-x))`` in FP32 (``ex2.approx``, IEEE division).""" + return Float32( + _asm( + T.f32(), + [_fv(x, loc=loc, ip=ip)], + "{\n.reg .f32 e, d;\n" + "mul.rn.f32 e, $1, 0fBFB8AA3B;\n" + "ex2.approx.f32 e, e;\n" + "add.rn.f32 d, 0f3F800000, e;\n" + "div.rn.f32 $0, 0f3F800000, d;\n}", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _div(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _mul(a, b, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_fv(a, loc=loc, ip=ip), _fv(b, loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """``x`` rounded to BF16 (nearest even), as FP32.""" + return Float32( + _asm( + T.f32(), + [_fv(x, loc=loc, ip=ip)], + "{\n.reg .b16 h;\n.reg .b32 w;\ncvt.rn.bf16.f32 h, $1;\n" + "mov.b32 w, {0, h};\nmov.b32 $0, w;\n}", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +class _Plan: + """Static sizes of one configuration (``world`` ranks, ``sms`` CTAs).""" + + def __init__(self, world: int, sms: int): + if INTER % world or (INTER // world) % 16: + raise ValueError(f"world={world}: the intermediate shard is not whole K steps") + self.world = world + self.sms = sms + self.inter = INTER // world + if self.inter % sms or (self.inter // sms) % 8: + raise ValueError(f"{self.inter} intermediate columns do not split into {sms} CTAs") + self.ic = self.inter // sms + self.gt = 2 * self.ic // 16 + self.kd = self.inter // 16 + if self.kd % self.gt: + raise ValueError(f"{self.kd} down K steps do not split into stages of {self.gt}") + self.stage_bytes = self.gt * WARPS * FRAG + self.gu_stages = GU_STEPS + self.dt = -(-DOWN_TILES // sms) + self.groups = -(-self.dt // WARPS) + self.dpg = self.kd // self.gt + self.stages = self.gu_stages + self.groups * self.dpg + self.cta_bytes = self.stages * self.stage_bytes + self.red_bytes = WARPS * self.gt * MAX_M * RED_PITCH * 4 + self.act_pitch = self.inter * 2 + 16 + self.act_bytes = MAX_M * self.act_pitch + self.out_bytes = WARPS * MAX_M * STAGE_PITCH * 4 + fixed = self.red_bytes + self.act_bytes + self.out_bytes + 1024 + self.ring = min(self.stages, (SMEM_BUDGET - fixed) // self.stage_bytes) + if self.ring < 2: + raise ValueError(f"world={world} sms={sms}: no room for a 2-stage weight ring") + if self.gt > 12: + raise ValueError(f"{self.gt} gate/up m-tiles per CTA: too many accumulators") + + +class K3SharedPersistKernel: + """The shared expert of ``m <= 8`` rows of one rank of ``world`` on ``sms`` CTAs (see the + module docstring).""" + + def __init__( + self, world: int, m: int, sms: int, mc_out: bool, beta: float, lbeta: Optional[float] + ): + if not 1 <= m <= MAX_M: + raise ValueError(f"m={m}: 1..{MAX_M}") + self.p = _Plan(world, sms) + self.world = world + self.m = m + self.sms = sms + self.mc_out = mc_out + self.beta = float(beta) + self.lbeta = None if lbeta is None else float(lbeta) + + @cute.jit + def __call__( + self, + x_addr: Int64, + w_addr: Int64, + xch_addr: Int64, + ctr_addr: Int64, + out_base: Int64, + out_flags: Int64, + rank: Int32, + stream: cuda.CUstream, + ): + p = self.p + + @cute.struct + class SharedStorage: + full: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, p.ring], 16] + empty: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, p.ring], 16] + ring: cute.struct.Align[ + cute.struct.MemRange[cutlass.Uint8, p.ring * p.stage_bytes], 128 + ] + red: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, p.red_bytes], 16] + act: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, p.act_bytes], 16] + out: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, p.out_bytes], 16] + + self.shared_storage = SharedStorage + self.kernel(x_addr, w_addr, xch_addr, ctr_addr, out_base, out_flags, rank).launch( + grid=(self.sms, 1, 1), + block=(THREADS, 1, 1), + smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x_addr: Int64, + w_addr: Int64, + xch_addr: Int64, + ctr_addr: Int64, + out_base: Int64, + out_flags: Int64, + rank: Int32, + ): + p = self.p + m = self.m + W = self.world + tidx, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + gq = lane // 4 + tq = lane % 4 + storage = cutlass.utils.SmemAllocator().allocate(self.shared_storage) + full = storage.full.data_ptr().toint() + empty = storage.empty.data_ptr().toint() + ring = storage.ring.data_ptr().toint() + red = storage.red.data_ptr().toint() + act_s = storage.act.data_ptr().toint() + out_s = storage.out.data_ptr().toint() + if tidx == 0: + for s in cutlass.range_constexpr(p.ring): + _mbar_init(full + 8 * s, 1) + _mbar_init(empty + 8 * s, WARPS) + _fence_mbar_init() + cute.arch.sync_threads() + w_cta = w_addr + Int64(cta) * Int64(p.cta_bytes) + + if warp == WARPS: + # ---- producer: every stage of this CTA's weight stream through the ring ---- + if lane == 0: + for s in cutlass.range(p.stages, unroll=1): + slot = s % p.ring + if s >= p.ring: + par = ((s // p.ring) - 1) % 2 + done = _mbar_try_wait(empty + 8 * slot, par) + while done == 0: + done = _mbar_try_wait(empty + 8 * slot, par) + # The compute warps' reads of the slot (generic proxy) before the bulk + # copy's writes (async proxy). + _fence_proxy_async() + _mbar_expect_tx(full + 8 * slot, p.stage_bytes) + _bulk_g2s( + ring + slot * p.stage_bytes, w_cta + Int64(s) * p.stage_bytes, + p.stage_bytes, full + 8 * slot, + ) # fmt: skip + else: + # ---- gate/up: K steps 8 j + warp of every m-tile ---- + xb = cute.make_rmem_tensor((2 * GU_STEPS,), Int32) + tok_ok = gq < m + xrow = x_addr + Int64(cutlass.select_(tok_ok, gq, Int32(0)) * HIDDEN) * 2 + for j in cutlass.range_constexpr(GU_STEPS): + col = 16 * (8 * j) + 16 * warp + 2 * tq + v0 = _ld_b32(xrow + Int64(col) * 2) + v1 = _ld_b32(xrow + Int64(col + 8) * 2) + xb[2 * j] = cutlass.select_(tok_ok, v0, Int32(0)) + xb[2 * j + 1] = cutlass.select_(tok_ok, v1, Int32(0)) + acc = [[Float32(0.0)] * 4 for _ in range(p.gt)] + for j in cutlass.range_constexpr(GU_STEPS): + slot = j % p.ring + par = (j // p.ring) % 2 + done = _mbar_try_wait(full + 8 * slot, Int32(par)) + while done == 0: + done = _mbar_try_wait(full + 8 * slot, Int32(par)) + base = ring + slot * p.stage_bytes + for mt in cutlass.range_constexpr(p.gt): + a4 = _lds_v4(base + ((mt * WARPS + warp) * 32 + lane) * 16) + c = acc[mt] + acc[mt] = list(_mma_bf16( + a4[0], a4[1], a4[2], a4[3], xb[2 * j], xb[2 * j + 1], + c[0], c[1], c[2], c[3], + )) # fmt: skip + cute.arch.sync_warp() + if lane == 0: + _mbar_arrive(empty + 8 * slot) + # The warp's partials: red[warp][mt][token][row]. + for mt in cutlass.range_constexpr(p.gt): + rb = red + ((warp * p.gt + mt) * MAX_M) * RED_PITCH * 4 + _sts_f32(rb + ((2 * tq) * RED_PITCH + gq) * 4, acc[mt][0]) + _sts_f32(rb + ((2 * tq + 1) * RED_PITCH + gq) * 4, acc[mt][1]) + _sts_f32(rb + ((2 * tq) * RED_PITCH + gq + 8) * 4, acc[mt][2]) + _sts_f32(rb + ((2 * tq + 1) * RED_PITCH + gq + 8) * 4, acc[mt][3]) + _bar_sync(1, WARPS * 32) + + # ---- activation of this CTA's columns: warp-order sums, BF16, SiTU, BF16 ---- + ic = p.ic + for i0 in cutlass.range_constexpr(0, ic * MAX_M, WARPS * 32): + item = tidx + i0 + if item < ic * MAX_M: + t = item // ic + col = item - t * ic + g = Float32(0.0) + u = Float32(0.0) + for w in cutlass.range_constexpr(WARPS): + gr = col + ur = ic + col + g = g + _lds_f32( + red + (((w * p.gt + gr // 16) * MAX_M + t) * RED_PITCH + gr % 16) * 4 + ) + u = u + _lds_f32( + red + (((w * p.gt + ur // 16) * MAX_M + t) * RED_PITCH + ur % 16) * 4 + ) + g = _bf16_round(g) + u = _bf16_round(u) + a = _mul(_mul(Float32(self.beta), _tanh(_div(g, Float32(self.beta)))), + _sigmoid(g)) # fmt: skip + ub = u + if cutlass.const_expr(self.lbeta is not None): + ub = _mul(Float32(self.lbeta), _tanh(_div(u, Float32(self.lbeta)))) + r = _mul(a, ub) + if t < m: + xa = xch_addr + Int64(t * p.inter + cta * ic + col) * 2 + _st_bf16_global(xa, r) + _bar_sync(1, WARPS * 32) + if tidx == 0: + _fence_acq_rel_gpu() + atom_add_acq_rel_u32(ctr_addr, Int32(1)) + arrived = ld_acquire_u32(ctr_addr) + while arrived < self.sms: + arrived = ld_acquire_u32(ctr_addr) + _bar_sync(1, WARPS * 32) + # The whole activation [m, I] into shared memory (tokens >= m: zeros). + vecs = MAX_M * (p.inter // 8) + for v0 in cutlass.range_constexpr(0, vecs, WARPS * 32): + v = tidx + v0 + if v < vecs: + t = v // (p.inter // 8) + c8 = (v - t * (p.inter // 8)) * 8 + dst = act_s + t * p.act_pitch + c8 * 2 + if t < m: + w4 = ld_volatile_v4(xch_addr + Int64(t * p.inter + c8) * 2) + _sts_v4(dst, w4[0], w4[1], w4[2], w4[3]) + else: + _sts_v4(dst, Int32(0), Int32(0), Int32(0), Int32(0)) + _bar_sync(1, WARPS * 32) + if tidx == 0: + # Every CTA has copied the activation once this count completes: re-arm. + old = atom_add_acq_rel_u32(ctr_addr + Int64(4), Int32(1)) + if old == self.sms - 1: + st_u32(ctr_addr, Int32(0)) + st_u32(ctr_addr + Int64(4), Int32(0)) + + # ---- down: m-tile 8 g + warp of this CTA's list, K = I ---- + f = ld_volatile_v4(out_flags) + cur, nbytes = f[0], f[2] + obase = out_base + Int64(cur) * Int64(nbytes) + arow = act_s + gq * p.act_pitch + stg = out_s + warp * MAX_M * STAGE_PITCH * 4 + for g in cutlass.range_constexpr(p.groups): + acc2 = [[Float32(0.0)] * 4, [Float32(0.0)] * 4] + for hh in cutlass.range_constexpr(p.dpg): + s = GU_STEPS + g * p.dpg + hh + slot = s % p.ring + par = (s // p.ring) % 2 + done = _mbar_try_wait(full + 8 * slot, Int32(par)) + while done == 0: + done = _mbar_try_wait(full + 8 * slot, Int32(par)) + base = ring + slot * p.stage_bytes + for kk in cutlass.range_constexpr(p.gt): + k = hh * p.gt + kk + a4 = _lds_v4(base + ((kk * WARPS + warp) * 32 + lane) * 16) + b0 = _lds_b32(arow + (16 * k + 2 * tq) * 2) + b1 = _lds_b32(arow + (16 * k + 2 * tq + 8) * 2) + c = acc2[k % 2] + acc2[k % 2] = list(_mma_bf16( + a4[0], a4[1], a4[2], a4[3], b0, b1, c[0], c[1], c[2], c[3], + )) # fmt: skip + cute.arch.sync_warp() + if lane == 0: + _mbar_arrive(empty + 8 * slot) + d = [acc2[0][i] + acc2[1][i] for i in range(4)] + _sts_f32(stg + ((2 * tq) * STAGE_PITCH + gq) * 4, d[0]) + _sts_f32(stg + ((2 * tq + 1) * STAGE_PITCH + gq) * 4, d[1]) + _sts_f32(stg + ((2 * tq) * STAGE_PITCH + gq + 8) * 4, d[2]) + _sts_f32(stg + ((2 * tq + 1) * STAGE_PITCH + gq + 8) * 4, d[3]) + cute.arch.sync_warp() + tile = (WARPS * g + warp) * self.sms + cta + tk = lane // 2 + h8 = (lane % 2) * 8 + if (lane < 2 * m) & (tile < DOWN_TILES): + e0, e1, e2, e3 = _lds_v4_f32(stg + (tk * STAGE_PITCH + h8) * 4) + e4, e5, e6, e7 = _lds_v4_f32(stg + (tk * STAGE_PITCH + h8 + 4) * 4) + o = [ + sanitize_bf16x2(pack_bf16x2(e0, e1)), + sanitize_bf16x2(pack_bf16x2(e2, e3)), + sanitize_bf16x2(pack_bf16x2(e4, e5)), + sanitize_bf16x2(pack_bf16x2(e6, e7)), + ] + dst = obase + (Int64(tk * W + rank) * HIDDEN + Int64(tile * 16 + h8)) * 2 + if cutlass.const_expr(self.mc_out): + multimem_st_v4(dst, o[0], o[1], o[2], o[3]) + else: + st_v4(dst, o[0], o[1], o[2], o[3]) + cute.arch.sync_warp() + + +@dsl_user_op +def _st_bf16_global(addr, x, *, loc=None, ip=None): + """Store ``x`` rounded to BF16 (nearest even) into global memory.""" + _asm( + None, + [_lv(addr, loc=loc, ip=ip), _fv(x, loc=loc, ip=ip)], + "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\nst.global.b16 [$0], h;\n}", + "l,f", + loc=loc, + ip=ip, + ) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def can_run(m: int, world: int, sms: int) -> bool: + """Whether the kernel serves ``m`` rows of the ``world``-rank shard on ``sms`` CTAs.""" + if not 1 <= m <= MAX_M or sms < 1: + return False + try: + _Plan(world, sms) + except ValueError: + return False + return True + + +def persist_fragments(gate_up: torch.Tensor, down: torch.Tensor, sms: int) -> torch.Tensor: + """This rank's shared expert weights in the kernel's stream order: ``gate_up`` BF16 + ``[2 I, 7168]`` (gate rows, then up rows; the GatedMLP shard) and ``down`` BF16 ``[7168, I]``; + returns int32 ``[sms, stages * stage_bytes / 4]`` (CTA ``c``: its gate/up stages, then its + down stages; padding m-tiles zero).""" + two_i, k = gate_up.shape + inter = two_i // 2 + if k != HIDDEN or tuple(down.shape) != (HIDDEN, inter): + raise ValueError(f"expected [2 I, {HIDDEN}] and [{HIDDEN}, I], got {gate_up.shape} " + f"and {down.shape}") # fmt: skip + if gate_up.dtype != torch.bfloat16 or down.dtype != torch.bfloat16: + raise ValueError("BF16 weights expected") + p = _Plan(INTER // inter, sms) + ic, gt = p.ic, p.gt + out = [] + for c in range(sms): + lo, hi = c * ic, (c + 1) * ic + rows = torch.cat((gate_up[lo:hi], gate_up[inter + lo : inter + hi])) + # [gt mt, 2 h, 8 gq, 56 j, 8 w, 2 hk, 4 tq, 2 e] -> [j, mt, w, gq, tq, hk, h, e] + v = rows.contiguous().view(gt, 2, 8, GU_STEPS, WARPS, 2, 4, 2) + gu = v.permute(3, 0, 4, 2, 6, 5, 1, 7).reshape(-1) + tiles = torch.zeros(p.groups * WARPS * 16, inter, dtype=torch.bfloat16, device=down.device) + for i in range(p.groups * WARPS): + t = i * sms + c + if t < DOWN_TILES: + tiles[16 * i : 16 * i + 16] = down[16 * t : 16 * t + 16] + # [g, 8 w, 2 h, 8 gq, dpg hh, gt kk, 2 hk, 4 tq, 2 e] -> [g, hh, kk, w, gq, tq, hk, h, e] + v = tiles.view(p.groups, WARPS, 2, 8, p.dpg, gt, 2, 4, 2) + dn = v.permute(0, 4, 5, 1, 3, 7, 6, 2, 8).reshape(-1) + out.append(torch.cat((gu, dn)).view(torch.int32)) + res = torch.stack(out) + assert res.shape[1] * 4 == p.cta_bytes, (res.shape, p.cta_bytes) + return res.contiguous() + + +class PersistWorkspace: + """The kernel's exchange buffer (BF16 ``[8, I]``) and its two counters (zero between + launches).""" + + def __init__(self, world: int, device): + self.xch = torch.zeros(MAX_M, INTER // world, dtype=torch.bfloat16, device=device) + self.ctr = torch.zeros(2, dtype=torch.int32, device=device) + + +def _set_carveout(compiled, device_index: int) -> None: + from . import k3_launch_attrs as attrs + + status, cu_device = cuda.cuDeviceGet(device_index) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuDeviceGet: {status}") + attr = cuda.CUfunction_attribute + for _, kernel in attrs._kernels(compiled, device_index): + (status,) = cuda.cuKernelSetAttribute( + attr.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, attrs.MAX_SHARED_CARVEOUT, + kernel, cu_device, + ) # fmt: skip + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuKernelSetAttribute: {status}") + + +@functools.cache +def _compiled( + world: int, m: int, sms: int, mc_out: bool, beta: float, lbeta: Optional[float], dev: int +): + kernel = K3SharedPersistKernel(world, m, sms, mc_out, beta, lbeta) + z = Int64(0) + with torch.cuda.device(dev): + compiled = cute.compile( + kernel, z, z, z, z, z, z, Int32(0), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) # fmt: skip + _set_carveout(compiled, dev) + return compiled + + +def run_k3_shared_persist( + *, + x: torch.Tensor, + frags: torch.Tensor, + world: int, + rank: int, + sms: int, + ws: PersistWorkspace, + out_base: int, + out_flags: torch.Tensor, + beta: float, + linear_beta: Optional[float], + mc_out: bool = True, +) -> None: + """Launch :class:`K3SharedPersistKernel` on the current stream. + + Args: + x: The shared expert's input ``[m, 7168]`` BF16 (contiguous, complete). + frags: :func:`persist_fragments` of this rank's weights for ``sms`` CTAs. + world / rank: The tail call's ranks and this rank's slot. + ws: The exchange buffer and counters (:class:`PersistWorkspace`). + out_base / out_flags: The tail call's Lamport workspace (multicast base, or with + ``mc_out=False`` a local base) and its rotation state; the kernel writes slot + ``(token, rank)`` of every token. + beta / linear_beta: SiTU's ``beta`` and ``linear_beta`` (``None``: no up clamp). + """ + m = x.shape[0] + if x.dtype != torch.bfloat16 or tuple(x.shape[1:]) != (HIDDEN,) or not x.is_contiguous(): + raise ValueError("x must be contiguous BF16 [m, 7168]") + if not can_run(m, world, sms): + raise ValueError(f"m={m} world={world} sms={sms} not supported") + lbeta = None if linear_beta is None else float(linear_beta) + _compiled(world, m, sms, mc_out, float(beta), lbeta, x.device.index)( + Int64(x.data_ptr()), + Int64(frags.data_ptr()), + Int64(ws.xch.data_ptr()), + Int64(ws.ctr.data_ptr()), + Int64(out_base), + Int64(out_flags.data_ptr()), + Int32(rank), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_small_batch_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_small_batch_moe.py new file mode 100644 index 000000000000..6f51ea35f4fe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_small_batch_moe.py @@ -0,0 +1,2567 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Small-batch NVFP4 routed experts: FC1 + SiTU + NVFP4 requantization + FC2 + combine for at +most 32 tokens in one persistent, dynamically scheduled grid (SM100 family; dispatched on Rubin +by ``moe.py`` through ``k3_small_batch_moe_op``). + +It reads the fused FC12 kernel's inputs: the NVFP4 tokens with linear scales, the Rubin FC1 +weight layout (``[gate16, up16]`` row groups), the FC2 weights and the tile-128 ``moe_sort`` +tables, and writes every element of the routed output. + +* Work list: every FC1 unit (one 128-row W1 block of one tile, tile-major), then every FC2 task + (one 128-channel m-tile of the output over one slice of consecutive tiles, slice-major). A tile + is one local expert and its <= 32 tokens. The slice count is + ``min(max_slices, ceil(nt / slice_tiles))`` for the ``nt`` tiles of the step, computed on the + device. CTAs claim units from a global counter. +* FC2's token operand (the FC1 output rows of the tile and their block scales) is streamed per + 256-wide k-tile next to the weight stage it multiplies: every ring stage carries the weights + (128 x 256 NVFP4 + scales) and a token part (``n_tok`` x 256 NVFP4 + one 2 KB scale block), so + experts with the whole intermediate (expert parallelism) are served. The weight TMA warp streams + the weights; the token warp waits until the tile's FC1 units are stored, then fills the token + part of the same stages (one full barrier for both parts). +* Combine: the warpgroup that owns a task accumulates ``alpha2[e] * scale[token, slot] * acc`` of + its tiles per token in FP32. Red combine (``deterministic_combine=False``): every task adds its + BF16-rounded per-token sums into the output (zeroed by the kernel first) with vector reds, in an + unspecified order. Slice combine (``deterministic_combine=True``): every task stores one FP32 + partial per (m-tile, slice) plus a token mask; the task that completes the m-tile's count sums + the partials in slice order and writes the BF16 rows once, so the output does not depend on the + schedule. +* Token-indexed layout (``token_indexed``, for MoE-TP slices of the intermediate): row + ``tile * 128 + j`` of the FC1 output is token j's intermediate for the tile's expert, for every + token of the step, so FC2's accumulator column j is token j in every tile and the FC2 epilogue + sums ``scale[tile, j] * acc[:, j]`` in registers, the per-tile scales coming from a shared-memory + slot table built once (expanded slot -> tile, ``alpha2 x weight``). Same FP32 operations in the + same order as the permuted layout, so the slice combine gives the same bits. +* Dependencies: after the stores of an FC1 unit the epilogue adds 1 to the tile's counter + (GPU-scope release); the token warp acquires ``n1 / 128`` (plus the zeroed output slices with + the red combine) before loading a tile's rows. Every FC2 task adds 1 more per tile after its + wait; the add that completes the count stores zero, and the CTA that exits last resets the work + and exit counters, so the counter buffer is zero after every launch. FC1 units have no + dependency and every CTA runs its FC1 units before its FC2 tasks, so the grid cannot deadlock + with any number of resident CTAs. +* Die-major weights (``die_major``, for two locality domains): every expert's first half of W1 + row blocks and of W2 output m-tiles lives in the first half of the weight tensors (domain 0), + the second halves in the second half (domain 1), as tensors ``[2 L, rows / 2, cols]`` (slot e = + first half of expert e, slot L + e = second half). FC1 units and FC2 tasks are claimed per + domain from four counters (own domain's FC1, the other's, own FC2, the other's), so both domains + hold the same work for any routing and every weight load is domain-local until the steals. The + units and their arithmetic are the default ones (same bits). +* Warp roles: epilogue warpgroups (warpgroup 0 also converts the token scales and runs every FC1 + epilogue; FC2 tasks alternate between the warpgroups), the MMA warp, the weight TMA warp, the + token warp. +* Published routing (``published_routing``, token-indexed layout): instead of the ``moe_sort`` + tables the kernel takes the step's top-k ids / weights, which the routing kernel before it + publishes token by token (stores, then a GPU-scope release add to the counter word + ``CTR_ROUTE``). The CTAs poll that word, build this rank's tables in shared memory (the active + local experts in ascending order are the tiles, ``moe_sort``'s tile order at <= 32 tokens) and + start the weight stream without waiting for the routing grid; the token warp and the epilogue + wait for it before they read the tokens or write the output. Same tables, same arithmetic. +* Output push (``output_push_world``): the task that combines an m-tile (slice combine), or the + task whose arrival completes the m-tile after its reds (red combine: it reads the m-tile's rows + back), also stores the BF16 rows into slot ``(token, rank)`` of every rank's current buffer of + a one-shot MNNVL all-reduce (``k3_lamport_ar.push_slot``, ``-0`` sent as ``+0``), for + a consumer that finishes that all-reduce. Without a local expert the token warps push zero rows. +* PDL: barrier setup and TMEM allocation run before ``griddepcontrol.wait``; the dependents are + released right after it, with published routing right after the tables (every consumer of the + output waits for the whole grid). +""" + +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 +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm, math +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + fmin, + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +from . import k3_lamport_ar as pushed_ar + +#: L2 eviction hint for the streamed expert weights (CacheHintSm90::EVICT_FIRST). +_EVICT_FIRST = 0x12F0000000000000 +#: Weight rows per MMA (UMMA M). +_MMA_M = 128 +#: K extent of one pipeline stage (four 64-wide block-scaled MMAs). +_KBLK = 256 +_SF_VEC = 16 +#: Bytes of one 128-row x 256-K block of scale factors (four 512-byte 32x4x4 atoms). +_SF_KBLK_BYTES = 2048 +#: Token tiles (UMMA N). +SUPPORTED_N_TOK = (8, 16, 32) +#: Largest token count of the kernel. +MAX_TOKENS = SUPPORTED_N_TOK[-1] +#: FP32 padding per staged accumulator row (bank spread of the transposed reads). +_SACC_PAD = 4 +#: Named barrier of epilogue warpgroup ``wg``: ``_EPI_BARRIER_BASE + wg``. +_EPI_BARRIER_BASE = 1 +#: int32 words of the counter buffer: [0] work (die-major: FC1 units of domain 0), [1] exited CTAs, +#: die-major claiming: [CTR_DIE1] FC1 units of domain 1, [CTR_FC2] / [CTR_FC2_DIE1] FC2 tasks of +#: domain 0 / 1 (own 128-byte lines); [CTR_MTILE_BASE + c] FC2 arrivals of m-tile c, +#: [CTR_TILE_BASE + t] FC1 units stored / FC2 tasks started of tile t. +CTR_DIE1 = 32 +CTR_FC2 = 48 +CTR_FC2_DIE1 = 56 +#: Published routing: tokens whose top-k the routing kernel has published (release adds; the CTA +#: that exits last resets it), on a 128-byte line of its own: the CTAs poll it while others claim +#: work. +CTR_ROUTE = 64 +CTR_MTILE_BASE = 96 +MAX_MTILES = 64 +CTR_TILE_BASE = CTR_MTILE_BASE + MAX_MTILES +#: Bytes of the weight part of one ring stage (128 x 256 NVFP4). +_SLOT_BYTES = _MMA_M * _KBLK // 2 +#: Largest slice count (the partial workspace is sized for it). +MAX_SLICES = 32 +#: Published routing: words of the shared-memory active-expert mask, one bit per local expert. +_ROUTE_MASK_WORDS = 28 +#: Published routing: most local experts (bits of the active-expert mask). +MAX_PUBLISHED_LOCAL_EXPERTS = 32 * _ROUTE_MASK_WORDS +#: Slices whose partial loads the slice combine keeps in flight at once. +_COMBINE_CHUNK = 4 +#: Red combine: slices of the output zeroed by the kernel (the CTAs that run units +#: 0 .. _ZERO_SLICES - 1); every tile counter counts them before FC2 may add into the output. +_ZERO_SLICES = 4 + + +def n_tok_for(num_tokens: int) -> int: + """Smallest token tile holding ``num_tokens`` tokens.""" + for n_tok in SUPPORTED_N_TOK: + if num_tokens <= n_tok: + return n_tok + raise ValueError(f"{num_tokens} tokens exceed the small-batch MoE kernel ({MAX_TOKENS})") + + +def counter_words(max_tiles: int) -> int: + """int32 words of the (zero-initialized) counter buffer for up to ``max_tiles`` tiles.""" + return CTR_TILE_BASE + max_tiles + + +def workspace_words(n2: int, n_tok: int, max_slices: int) -> int: + """32-bit words of the slice combine's workspace: token masks, then FP32 partials of every + (m-tile, slice); nothing in it needs initialization.""" + mtiles = n2 // _MMA_M + return mtiles * max_slices * (1 + n_tok * _MMA_M) + + +def _smem_capacity() -> int: + """Shared memory per CTA of the current device.""" + major, minor = torch.cuda.get_device_capability() + return utils.get_smem_capacity_in_bytes(f"sm_{major}{minor}") + + +def _align(v: int, a: int) -> int: + return (v + a - 1) // a * a + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _iv(value, *, loc=None, ip=None): + return Int32(value).ir_value(loc=loc, ip=ip) + + +@dsl_user_op +def _ldg_u32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _sts_u32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), Float32(value).ir_value(loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.f32 $0, [$1];", + "=f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _lds_v4_f32(addr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_iv(addr, loc=loc, ip=ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + loc=loc, + ip=ip, + ) + return tuple(Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(4)) + + +@dsl_user_op +def _stg_v2_b32(addr, v0, v1, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(v0, loc=loc, ip=ip), _iv(v1, loc=loc, ip=ip)], + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u8(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.global.u8 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_v4_zero(addr, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip)], + "st.global.v4.b32 [$0], {0, 0, 0, 0};", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_f32_pred(addr, value, pred, *, loc=None, ip=None): + _asm( + None, + [ + Int64(addr).ir_value(loc=loc, ip=ip), + Float32(value).ir_value(loc=loc, ip=ip), + _iv(pred, loc=loc, ip=ip), + ], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.global.f32 [$0], $1; }", + "l,f,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ldg_cg_f32_pred(addr, pred, *, loc=None, ip=None): + """``[addr]`` (L2) if ``pred`` != 0, else 0.0; no branch, so a run of them is in flight at once.""" + return Float32( + _asm( + T.f32(), + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(pred, loc=loc, ip=ip)], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; mov.b32 $0, 0; @p ld.global.cg.f32 $0, [$1]; }", + "=f,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ldg_cg_v4(addr, *, loc=None, ip=None): + """16 bytes at ``addr`` read from L2 (other CTAs' reds), as four 32-bit words.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.cg.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)) + + +@dsl_user_op +def _stg_bf16_from_f32(addr, value, *, loc=None, ip=None): + """Round ``value`` to BF16 (nearest even) and store the 2 bytes.""" + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Float32(value).ir_value(loc=loc, ip=ip)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; st.global.b16 [$0], h; }", + "l,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _red_add_v8_bf16_from_f32_smem(dst_addr, src_smem_addr, *, loc=None, ip=None): + """``dst[0:8] += bf16(src[0:8])``: 8 FP32 read from shared memory, rounded to BF16, one + vector red.""" + _asm( + None, + [Int64(dst_addr).ir_value(loc=loc, ip=ip), _iv(src_smem_addr, loc=loc, ip=ip)], + "{\n.reg .f32 f0, f1, f2, f3, f4, f5, f6, f7;\n.reg .b32 a, b, c, d;\n" + "ld.shared.v4.f32 {f0, f1, f2, f3}, [$1];\n" + "ld.shared.v4.f32 {f4, f5, f6, f7}, [$1+16];\n" + "cvt.rn.bf16x2.f32 a, f1, f0;\ncvt.rn.bf16x2.f32 b, f3, f2;\n" + "cvt.rn.bf16x2.f32 c, f5, f4;\ncvt.rn.bf16x2.f32 d, f7, f6;\n" + "red.global.add.noftz.v4.bf16x2 [$0], {a, b, c, d};\n}", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _signal_add_release(addr, *, loc=None, ip=None): + """GPU-scope release add of 1: every write that happens-before it (in particular the CTA's + writes ordered by a preceding ``bar.sync``) is visible to an acquire that observes it.""" + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip)], + "red.release.gpu.global.add.s32 [$0], 1;", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _red_add_release(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "red.release.gpu.global.add.s32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.acquire.gpu.global.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.relaxed.gpu.global.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_relaxed_f32(addr, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.relaxed.gpu.global.f32 $0, [$1];", + "=f,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _red_shared_or(addr, value, *, loc=None, ip=None): + _asm( + None, + [_iv(addr, loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "red.shared.or.b32 [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bf16x2_bits(lo, hi, *, loc=None, ip=None): + """Two FP32 values rounded to BF16 (nearest even), ``lo`` in the low half.""" + return Int32( + _asm( + T.i32(), + [Float32(lo).ir_value(loc=loc, ip=ip), Float32(hi).ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bit(n, *, loc=None, ip=None): + """``1 << n`` for ``0 <= n < 32``.""" + return Int32( + _asm(T.i32(), [_iv(n, loc=loc, ip=ip)], "shl.b32 $0, 1, $1;", "=r,r", loc=loc, ip=ip) + ) + + +@dsl_user_op +def _atom_add_relaxed(addr, value, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "atom.relaxed.gpu.global.add.s32 $0, [$1], $2;", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _atom_add_acq_rel(addr, value, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "atom.acq_rel.gpu.global.add.s32 $0, [$1], $2;", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), _iv(value, loc=loc, ip=ip)], + "st.relaxed.gpu.global.b32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_g2s(dst_smem, src_gmem, nbytes, mbar, *, loc=None, ip=None): + """``nbytes`` contiguous bytes global -> shared, crediting ``mbar``.""" + _asm( + None, + [ + _iv(dst_smem, loc=loc, ip=ip), + Int64(src_gmem).ir_value(loc=loc, ip=ip), + _iv(nbytes, loc=loc, ip=ip), + _iv(mbar, loc=loc, ip=ip), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];", + "r,l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ballot(pred, *, loc=None, ip=None): + """Warp ballot of ``pred != 0`` (every lane of the warp executes it).""" + return Int32( + _asm( + T.i32(), + [_iv(pred, loc=loc, ip=ip)], + "{ .reg .pred p; setp.ne.b32 p, $1, 0; vote.sync.ballot.b32 $0, p, 0xffffffff; }", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return Int32(_asm(T.i32(), [], "mov.u32 $0, %smid;", "=r", loc=loc, ip=ip)) + + +def pack_sm_domain_words(domains: tuple) -> tuple: + """Bit ``s % 32`` of word ``s // 32`` is the locality domain of SM ``s`` (signed 32-bit).""" + words = [0] * (-(-len(domains) // 32)) + for sm, domain in enumerate(domains): + if domain == 1: + words[sm // 32] |= 1 << (sm % 32) + return tuple(w - (1 << 32) if w >= (1 << 31) else w for w in words) + + +def _s2t_copy_and_partition(sSF: cute.Tensor, tSF: cute.Tensor, sf_dtype): + """tcgen05.cp (SMEM -> TMEM) partitions of a scale-factor operand.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom_s2t = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), sf_dtype) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + +class RubinSmallBatchFusedMoeKernel: + """FC1 + SiTU + NVFP4 quantization + FC2 + combine for at most 32 tokens. + + :param n_tok: Token tile (MMA N) of both GEMMs: 8, 16 or 32. + :param n1: Rows of one expert's W1 (2 x intermediate, ``[gate16, up16]`` row groups). + :param k1: Hidden size (W1 K, W2 rows, output columns). + :param topk: Experts per token. + :param situ_beta: SiTU gate-side constant. + :param situ_linear_beta: SiTU up-side constant. + :param token_indexed: Token-indexed FC1 output layout. + :param claim_free: Ring stages a CTA needs free before it claims its next FC2 task (0: one unit + ahead). + :param max_slices: Largest number of FC2 slices. + :param max_ab_stage: Cap on the ring depth (the rest of shared memory sizes it). + :param deterministic_combine: Slice combine instead of the BF16 vector reds. + :param num_acc: Accumulator stages in TMEM per epilogue warpgroup. + :param num_epi_wg: Epilogue warpgroups (1 or 2). + :param slice_tiles: Target tiles per FC2 slice. + :param queue_depth: Entries of the shared-memory unit queue. + :param use_pdl: Launch with programmatic dependent launch. + :param die_major: Die-major expert weights, work claimed per locality domain. + :param sm_domain_words: :func:`pack_sm_domain_words` of the device (with ``die_major``). + :param published_routing: Tables built in the kernel from the top-k the routing kernel + publishes (token-indexed layout only). + :param output_push_world: Also push the output rows into a one-shot all-reduce over that many + ranks (either combine; 0: off). + """ + + @staticmethod + def supports(n_tok: int, n1: int, k1: int) -> bool: + # FC1 runs in 128-row W1 blocks; FC2's K (the intermediate n1 / 2) needs whole 64-wide + # scale atoms. + return ( + n_tok in SUPPORTED_N_TOK + and n1 % _MMA_M == 0 + and (n1 // 2) % 64 == 0 + and k1 % _KBLK == 0 + and k1 // _MMA_M <= MAX_MTILES + ) + + @staticmethod + def h_sf_tile_bytes(n1: int) -> int: + """Bytes of one tile's FC1-output scales (128 rows; the intermediate padded to a multiple + of 256, so a half last k-tile of FC2 reads zeros there: the buffer's padding must be zero).""" + return _MMA_M * ((n1 // 2 + _KBLK - 1) // _KBLK * _KBLK) // _SF_VEC + + def __init__( + self, + n_tok: int, + n1: int, + k1: int, + topk: int, + situ_beta: float, + situ_linear_beta: float, + token_indexed: bool, + claim_free: int, + max_slices: int, + max_ab_stage: int, + deterministic_combine: bool = False, + num_acc: int = 2, + num_epi_wg: int = 2, + slice_tiles: int = 1, + queue_depth: int = 16, + use_pdl: bool = True, + die_major: bool = False, + sm_domain_words: tuple = (), + published_routing: bool = False, + output_push_world: int = 0, + ): + assert self.supports(n_tok, n1, k1), (n_tok, n1, k1) + self.output_push_world = output_push_world + assert not published_routing or token_indexed, "published routing is token-indexed" + self.published_routing = published_routing + assert not die_major or sm_domain_words, "die-major claiming needs the SM table" + self.die_major = die_major + self.sm_domain_words = tuple(sm_domain_words) + assert num_epi_wg in (1, 2) and num_acc >= 2 + assert 1 <= max_slices <= MAX_SLICES and slice_tiles >= 1 + assert situ_beta > 0 and situ_linear_beta > 0 + self.token_indexed = token_indexed + # Claim policy: 0 = one unit ahead (the next unit is claimed when the current one starts + # loading); n > 0 = once the FC1 units are all claimed, the next unit is claimed only when + # n ring stages are free for it (FC1 units are still claimed right after the current one). + self.claim_free_req = int(claim_free) + self.combine_red = not deterministic_combine + self.ready_target = n1 // _MMA_M + (_ZERO_SLICES if self.combine_red else 0) + self.n_tok = n_tok + self.n1, self.k1 = n1, k1 + self.n2, self.k2 = k1, n1 // 2 + self.topk = topk + self.nkb1, self.nblk1 = k1 // _KBLK, n1 // _MMA_M + # FC2's K (the intermediate) may end in a half k-tile: the TMA zero-fills the weights, + # their scales and the rows past K, the FC1-output scales are padded. + self.nkb2, self.nblk2 = (self.k2 + _KBLK - 1) // _KBLK, self.n2 // _MMA_M + self.h_tile_sf_bytes = self.h_sf_tile_bytes(n1) + self.situ_beta = float(situ_beta) + self.situ_linear_beta = float(situ_linear_beta) + self.num_epi_wg = num_epi_wg + self.num_acc = num_acc + self.nacc_total = num_acc * num_epi_wg + self.slice_tiles = slice_tiles + self.max_slices = max_slices + self.queue_depth = queue_depth + self.use_pdl = use_pdl + self.mma_tiler = (_MMA_M, n_tok, _KBLK) + self.acc_cols = n_tok * self.nacc_total + self.sfa_cols = 16 + self.sfb_cols = 16 + cols = self.acc_cols + self.sfa_cols + self.sfb_cols + self.tmem_cols = 32 + while self.tmem_cols < cols: + self.tmem_cols *= 2 + self.a_stage_bytes = _SLOT_BYTES + _SF_KBLK_BYTES + self.b_stage_bytes = n_tok * _KBLK // 2 + self.fc2_stage_bytes = self.a_stage_bytes + self.b_stage_bytes + _SF_KBLK_BYTES + self.x_bytes = self.nkb1 * self.b_stage_bytes + self.sacc_row_floats = _MMA_M + _SACC_PAD + self.acc_buf_bytes = n_tok * self.sacc_row_floats * 4 + # slice combine workspace geometry (32-bit words): masks [mtile][max_slices], then + # partials [mtile][max_slices][n_tok][128] + self.mask_word_base = 0 + self.part_word_base = self.nblk2 * max_slices + + # token-indexed slot table: one entry per expanded slot (token x top-k) of the step: its + # tile (-1: another rank's expert) and its FP32 scale; per epilogue warpgroup two buffers + # of a tile's per-token scales and token masks. + self.n_x = n_tok * topk + self.n_sv = 2 * num_epi_wg * n_tok + # Per-stage bytes: weights + their scales + FC2 token rows + their scale block. + per_stage = _SLOT_BYTES + _SF_KBLK_BYTES + self.b_stage_bytes + _SF_KBLK_BYTES + 16 + fixed = ( + self.x_bytes + + self.nkb1 * _SF_KBLK_BYTES + + (self.acc_buf_bytes if num_epi_wg > 1 else 0) + + self.acc_buf_bytes + + 2 * num_epi_wg * n_tok * 4 + + 2 * queue_depth * 4 + + 8 * self.n_x + + 4 * self.n_sv + + 8 * num_epi_wg + + 64 + + 4 * 1024 # barriers, TMEM holder, flags and alignment slack + ) + # published routing: tile -> expert table (at most one tile per expanded slot), the + # active-expert mask with its word prefix counts, the tile count + self.route_t2e = self.n_x if published_routing else 4 + if published_routing: + fixed += 4 * self.route_t2e + 8 * _ROUTE_MASK_WORDS + 64 + self.num_ab_stage = min(max_ab_stage, (_smem_capacity() - fixed) // per_stage) + assert self.num_ab_stage >= 2, "ring too small" + self.claim_free = min(self.claim_free_req, self.num_ab_stage) + if self.die_major: + assert self.nblk1 % 2 == 0 and self.nblk2 % 2 == 0, "die-major halves need even blocks" + S = self.num_ab_stage + # Region after the weight ring: [B2 stages | x] (token-operand view), then + # [SFB2 stages | x scales] (token-scale view), then warpgroup 1's token accumulator. + self.off_b2 = 0 + self.off_x = S * self.b_stage_bytes + self.off_sfb2 = _align(self.off_x + self.x_bytes, 1024) + self.off_xsf = self.off_sfb2 + S * _SF_KBLK_BYTES + self.off_acc1 = _align(self.off_xsf + self.nkb1 * _SF_KBLK_BYTES, 128) + end = self.off_acc1 + (self.acc_buf_bytes if num_epi_wg > 1 else 0) + self.region_bytes = _align(end, 1024) + self.epi_warps = 4 * num_epi_wg + self.mma_warp = self.epi_warps + self.tma_warp = self.epi_warps + 1 + self.tok_warp = self.epi_warps + 2 + self.threads = 32 * (self.epi_warps + 3) + + # ------------------------------------------------------------------ host + @cute.jit + def __call__( + self, + x: cute.Tensor, + w1: cute.Tensor, + w1_sf: cute.Tensor, + h: cute.Tensor, + w2: cute.Tensor, + w2_sf: cute.Tensor, + x_sf_ptr: cute.Pointer, + h_ptr: cute.Pointer, + h_sf_ptr: cute.Pointer, + out_ptr: cute.Pointer, + alpha1: cute.Tensor, + alpha2: cute.Tensor, + tile_idx_to_expert_idx: cute.Tensor, + tile_idx_to_mn_limit: cute.Tensor, + permuted_idx_to_expanded_idx: cute.Tensor, + expanded_idx_to_permuted_idx: cute.Tensor, + num_non_exiting_tiles: cute.Tensor, + global_sf: cute.Tensor, + token_final_scales: cute.Tensor, + ctr_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + orig_m: Int32, + expert_offset: Int32, + push_mc: Int64, + push_flags: Int64, + push_rank: Int32, + num_ctas: cutlass.Constexpr, + stream: cuda.CUstream, + ): + self.ab_dtype = w1.element_type + self.sf_dtype = w1_sf.element_type + major = utils.LayoutEnum.from_tensor(w1).mma_major_mode() + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.ab_dtype, + major, + major, + self.sf_dtype, + _SF_VEC, + tcgen05.CtaGroup.ONE, + self.mma_tiler[:2], + ) + S = self.num_ab_stage + a_layout = sm100_utils.make_smem_layout_a(tiled_mma, self.mma_tiler, self.ab_dtype, S) + sfa_layout = blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler, _SF_VEC, S) + # Token-operand views: stages [0, S) are the ring's FC2 token parts, [S, S + nkb1) the + # step's tokens (FC1); the same for the scales. + self.b_view_stages = S + self.nkb1 + b_layout = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.ab_dtype, self.b_view_stages + ) + sfb_layout = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, self.mma_tiler, _SF_VEC, self.b_view_stages + ) + w1_sf = cute.make_tensor( + w1_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF(w1.shape, _SF_VEC) + ) + w2_sf = cute.make_tensor( + w2_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF(w2.shape, _SF_VEC) + ) + cl = cute.tiled_divide(cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) + a_op = sm100_utils.cluster_shape_to_tma_atom_A((1, 1), tiled_mma.thr_id) + b_op = sm100_utils.cluster_shape_to_tma_atom_B((1, 1), tiled_mma.thr_id) + a_one = cute.slice_(a_layout, (None, None, None, 0)) + sfa_one = cute.slice_(sfa_layout, (None, None, None, 0)) + b_one = cute.slice_(b_layout, (None, None, None, 0)) + tma_a1, tA1 = cute.nvgpu.make_tiled_tma_atom_A( + a_op, w1, a_one, self.mma_tiler, tiled_mma, cl.shape + ) + tma_sfa1, tSFA1 = cute.nvgpu.make_tiled_tma_atom_A( + a_op, w1_sf, sfa_one, self.mma_tiler, tiled_mma, cl.shape, internal_type=cutlass.Int16 + ) + tma_a2, tA2 = cute.nvgpu.make_tiled_tma_atom_A( + a_op, w2, a_one, self.mma_tiler, tiled_mma, cl.shape + ) + tma_sfa2, tSFA2 = cute.nvgpu.make_tiled_tma_atom_A( + a_op, w2_sf, sfa_one, self.mma_tiler, tiled_mma, cl.shape, internal_type=cutlass.Int16 + ) + tma_b1, tB1 = cute.nvgpu.make_tiled_tma_atom_B( + b_op, x, b_one, self.mma_tiler, tiled_mma, cl.shape + ) + tma_b2, tB2 = cute.nvgpu.make_tiled_tma_atom_B( + b_op, h, b_one, self.mma_tiler, tiled_mma, cl.shape + ) + + NACC, NWG, Q = self.nacc_total, self.num_epi_wg, self.queue_depth + + @cute.struct + class SharedStorage: + a_full: cute.struct.MemRange[cutlass.Int64, S] + a_empty: cute.struct.MemRange[cutlass.Int64, S] + acc_full: cute.struct.MemRange[cutlass.Int64, NACC] + acc_empty: cute.struct.MemRange[cutlass.Int64, NACC] + uq_full: cute.struct.MemRange[cutlass.Int64, Q] + uq_empty: cute.struct.MemRange[cutlass.Int64, Q] + tok_full: cutlass.Int64 + sfb_full: cutlass.Int64 + tmem_ready: cutlass.Int64 + tmem_free: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sUnits: cute.struct.Align[cute.struct.MemRange[Int32, Q], 16] + sStages: cute.struct.Align[cute.struct.MemRange[Int32, Q], 16] + sFlag: cute.struct.Align[cute.struct.MemRange[Int32, 4], 16] + sA: cute.struct.Align[ + cute.struct.MemRange[self.ab_dtype, cute.cosize(a_layout.outer)], 1024 + ] + sRegion: cute.struct.Align[cute.struct.MemRange[cutlass.Uint8, self.region_bytes], 1024] + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(sfa_layout)], 1024 + ] + sAcc: cute.struct.Align[ + cute.struct.MemRange[Float32, self.n_tok * self.sacc_row_floats], 16 + ] + sTok: cute.struct.Align[cute.struct.MemRange[Int32, NWG * self.n_tok], 16] + sScale: cute.struct.Align[cute.struct.MemRange[Float32, NWG * self.n_tok], 16] + sTl: cute.struct.Align[cute.struct.MemRange[Int32, self.n_x], 16] + sSc: cute.struct.Align[cute.struct.MemRange[Float32, self.n_x], 16] + sSv: cute.struct.Align[cute.struct.MemRange[Float32, self.n_sv], 16] + sSvMask: cute.struct.Align[cute.struct.MemRange[Int32, 2 * NWG], 16] + sT2e: cute.struct.Align[cute.struct.MemRange[Int32, self.route_t2e], 16] + sAm: cute.struct.Align[cute.struct.MemRange[Int32, _ROUTE_MASK_WORDS], 16] + sApre: cute.struct.Align[cute.struct.MemRange[Int32, _ROUTE_MASK_WORDS], 16] + sNt: cute.struct.Align[cute.struct.MemRange[Int32, 4], 16] + + self.shared_storage = SharedStorage + self.kernel( + tiled_mma, + tma_a1, + tA1, + tma_sfa1, + tSFA1, + tma_a2, + tA2, + tma_sfa2, + tSFA2, + tma_b1, + tB1, + tma_b2, + tB2, + x_sf_ptr, + h_ptr, + h_sf_ptr, + out_ptr, + alpha1, + alpha2, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + expanded_idx_to_permuted_idx, + num_non_exiting_tiles, + global_sf, + token_final_scales, + ctr_ptr, + ws_ptr, + orig_m, + expert_offset, + push_mc, + push_flags, + push_rank, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + cl, + ).launch( + grid=(num_ctas, 1, 1), + block=[self.threads, 1, 1], + cluster=(1, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + use_pdl=self.use_pdl, + ) + + def _b_view(self, st, layout): + base = st.sRegion.data_ptr() + self.off_b2 + return cute.make_tensor( + cute.recast_ptr(base, layout.inner, dtype=self.ab_dtype), layout.outer + ) + + def _sfb_view(self, st, layout): + base = st.sRegion.data_ptr() + self.off_sfb2 + return cute.make_tensor(cute.recast_ptr(base, dtype=self.sf_dtype), layout) + + def _region_addr(self, st, off: int): + return (st.sRegion.data_ptr() + off).toint() + + @cute.jit + def _wg_barrier(self, wg: Int32): + for i in cutlass.range_constexpr(self.num_epi_wg): + if wg == i: + cute.arch.barrier(barrier_id=_EPI_BARRIER_BASE + i, number_of_threads=128) + + @cute.jit + def _fc2_task(self, v: Int32, nt: Int32, ns: Int32): + """FC2 task ``v`` (index past the FC1 units): (m-tile, combine slice, tiles [t0, t1)).""" + s = v // self.nblk2 + c = v - s * self.nblk2 + t0, t1 = self._slice_bounds(s, nt, ns) + return c, s, t0, t1 + + @cute.jit + def _slice_bounds(self, s: Int32, nt: Int32, ns: Int32): + """Tiles [t0, t1) of slice ``s`` of ``ns`` over ``nt`` tiles.""" + return s * nt // ns, (s + 1) * nt // ns + + # ------------------------------------------------------------------ unit queue + @cute.jit + def _q_write(self, st, j: Int32, u: Int32, g: Int32): + """Publish the CTA's ``j``-th unit (-1: no more units) and its first ring index.""" + Q = self.queue_depth + q = j % Q + if j >= Q: + cute.arch.mbarrier_wait(st.uq_empty.data_ptr() + q, ((j // Q) & 1) ^ 1) + units = st.sUnits.get_tensor(cute.make_layout(Q)) + stages = st.sStages.get_tensor(cute.make_layout(Q)) + with cute.arch.elect_one(): + units[q] = u + stages[q] = g + cute.arch.mbarrier_arrive(st.uq_full.data_ptr() + q) + + @cute.jit + def _q_read(self, st, j: Int32): + """The CTA's ``j``-th unit and its first ring index (whole warp).""" + Q = self.queue_depth + q = j % Q + cute.arch.mbarrier_wait(st.uq_full.data_ptr() + q, (j // Q) & 1) + units = st.sUnits.get_tensor(cute.make_layout(Q)) + stages = st.sStages.get_tensor(cute.make_layout(Q)) + u = units[q] + g = stages[q] + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(st.uq_empty.data_ptr() + q) + return u, g + + # ------------------------------------------------------------------ kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_a1: cute.CopyAtom, + mA1: cute.Tensor, + tma_sfa1: cute.CopyAtom, + mSFA1: cute.Tensor, + tma_a2: cute.CopyAtom, + mA2: cute.Tensor, + tma_sfa2: cute.CopyAtom, + mSFA2: cute.Tensor, + tma_b1: cute.CopyAtom, + mB1: cute.Tensor, + tma_b2: cute.CopyAtom, + mB2: cute.Tensor, + x_sf_ptr: cute.Pointer, + h_ptr: cute.Pointer, + h_sf_ptr: cute.Pointer, + out_ptr: cute.Pointer, + alpha1: cute.Tensor, + alpha2: cute.Tensor, + tile_idx_to_expert_idx: cute.Tensor, + tile_idx_to_mn_limit: cute.Tensor, + permuted_idx_to_expanded_idx: cute.Tensor, + expanded_idx_to_permuted_idx: cute.Tensor, + num_non_exiting_tiles: cute.Tensor, + global_sf: cute.Tensor, + token_final_scales: cute.Tensor, + ctr_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + orig_m: Int32, + expert_offset: Int32, + push_mc: Int64, + push_flags: Int64, + push_rank: Int32, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + cl: cute.Layout, + ): + S = self.num_ab_stage + NACC, Q = self.nacc_total, self.queue_depth + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + nctas_, _, _ = cute.arch.grid_dim() + nctas = Int32(nctas_) + + st = utils.SmemAllocator().allocate(self.shared_storage) + if warp_idx == self.tma_warp: + for atom in (tma_a1, tma_sfa1, tma_a2, tma_sfa2, tma_b1, tma_b2): + cpasync.prefetch_descriptor(atom) + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(S): + cute.arch.mbarrier_init(st.a_full.data_ptr() + i, 1) + cute.arch.mbarrier_init(st.a_empty.data_ptr() + i, 1) + for i in cutlass.range_constexpr(NACC): + cute.arch.mbarrier_init(st.acc_full.data_ptr() + i, 1) + cute.arch.mbarrier_init(st.acc_empty.data_ptr() + i, 4) + for i in cutlass.range_constexpr(Q): + cute.arch.mbarrier_init(st.uq_full.data_ptr() + i, 1) + cute.arch.mbarrier_init(st.uq_empty.data_ptr() + i, self.epi_warps + 2) + cute.arch.mbarrier_init(st.tok_full, 1) + cute.arch.mbarrier_init(st.sfb_full, 4) + cute.arch.mbarrier_init(st.tmem_ready, 1) + cute.arch.mbarrier_init(st.tmem_free, self.epi_warps) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + if warp_idx == self.mma_warp: + cute.arch.alloc_tmem(self.tmem_cols, st.tmem_holding_buf) + cute.arch.relinquish_tmem_alloc_permit() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(st.tmem_ready) + if cutlass.const_expr(self.use_pdl and not self.published_routing): + griddepcontrol_wait() + griddepcontrol_launch_dependents() + # The CTA's first claim after the grid dependency: the previous launch on this counter + # buffer may still run until then (die-major: claimed by the weight producer). With + # published routing the routing kernel's producer waited for every earlier grid, so the + # claim (die-major: on the own domain's FC1 counter, resolved by the weight producer once + # the tile count is known) runs before the tables. + first = Int32(0) + if cutlass.const_expr(not self.die_major or self.published_routing): + if warp_idx == self.tma_warp: + if cute.arch.lane_idx() == 0: + word = Int64(0) + if cutlass.const_expr(self.die_major): + if self._sm_domain(_smid()) == 1: + word = Int64(CTR_DIE1) + first = _atom_add_relaxed(ctr_ptr.toint() + word * 4, 1) + + t2e_pre = Int32(0) + nt = Int32(0) + if cutlass.const_expr(self.published_routing): + self._published_tables( + st, alpha2, expanded_idx_to_permuted_idx, ctr_ptr, orig_m, expert_offset, tidx + ) + if cutlass.const_expr(self.use_pdl): + # every earlier grid completed before the routing kernel published (its producer + # waited for them); every consumer of the output waits for this whole grid + griddepcontrol_launch_dependents() + nt = st.sNt.get_tensor(cute.make_layout(4))[0] + if warp_idx == self.tma_warp: + lane = cute.arch.lane_idx() + t2e_pre = st.sT2e.get_tensor(cute.make_layout(self.route_t2e))[ + cutlass.min(lane, self.route_t2e - 1) + ] + else: + if warp_idx == self.tma_warp: + lane = cute.arch.lane_idx() + t2e_pre = tile_idx_to_expert_idx[ + cutlass.min(lane, cute.size(tile_idx_to_expert_idx) - 1) + ] + nt = Int32(num_non_exiting_tiles[0]) + w1 = nt * self.nblk1 + ns = cutlass.min(Int32(self.max_slices), (nt + (self.slice_tiles - 1)) // self.slice_tiles) + total = w1 + ns * self.nblk2 + thr_mma = tiled_mma.get_slice(0) + + if warp_idx == self.tma_warp: + self._weight_producer( + st, + thr_mma, + cl, + a_layout, + sfa_layout, + tma_a1, + mA1, + tma_sfa1, + mSFA1, + tma_a2, + mA2, + tma_sfa2, + mSFA2, + tile_idx_to_expert_idx, + ctr_ptr, + first, + t2e_pre, + w1, + total, + nctas, + nt, + ns, + ) + elif warp_idx == self.tok_warp: + self._token_producer( + st, + thr_mma, + cl, + b_layout, + tma_b1, + mB1, + tma_b2, + mB2, + h_sf_ptr, + out_ptr, + ctr_ptr, + orig_m, + nt, + ns, + w1, + bid, + nctas, + (push_mc, push_flags, push_rank), + ) + elif warp_idx == self.mma_warp: + self._mma(st, tiled_mma, a_layout, b_layout, sfa_layout, sfb_layout, w1, nt, ns) + else: + wg = tidx // 128 + et = tidx % 128 + cute.arch.mbarrier_wait(st.tmem_ready, 0) + acc_ptr = cute.arch.retrieve_tmem_ptr( + Float32, alignment=16, ptr_to_buffer_holding_addr=st.tmem_holding_buf + ) + if cutlass.const_expr(self.published_routing and self.use_pdl): + # the token scales and the output come with the routing grid's end + griddepcontrol_wait() + if wg == 0: + if w1 > 0: + self._token_scales_to_smem(st, x_sf_ptr, orig_m, et) + if cutlass.const_expr(self.published_routing): + self._published_slots( + st, + alpha2, + expanded_idx_to_permuted_idx, + token_final_scales, + orig_m, + expert_offset, + tidx, + ) + elif cutlass.const_expr(self.token_indexed): + self._build_slot_table( + st, + alpha2, + tile_idx_to_expert_idx, + expanded_idx_to_permuted_idx, + token_final_scales, + orig_m, + tidx, + ) + self._epilogue( + st, + tiled_mma, + acc_ptr, + h_ptr, + h_sf_ptr, + out_ptr, + alpha1, + alpha2, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + global_sf, + token_final_scales, + ctr_ptr, + ws_ptr, + orig_m, + w1, + nt, + ns, + et, + wg, + (push_mc, push_flags, push_rank), + ) + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(st.tmem_free) + + # ------------------------------------------------------------------ producers + @cute.jit + def _reserve_stage(self, a_empty, g): + """Wait until ring stage ``g`` may be overwritten.""" + S = self.num_ab_stage + stage = g % S + if g >= S: + cute.arch.mbarrier_wait(a_empty + stage, ((g // S) & 1) ^ 1) + return stage + + @cute.jit + def _load_stage( + self, + tma_a, + tma_sfa, + tAsA, + tAgA, + tAsSFA, + tAgSFA, + blk, + kb, + expert, + a_full, + a_empty, + g, + tx_bytes: cutlass.Constexpr, + ): + stage = self._reserve_stage(a_empty, g) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(a_full + stage, tx_bytes) + # the weights are read once per step: L2 evict-first + cute.copy( + tma_a, + tAgA[(None, blk, kb, expert)], + tAsA[(None, stage)], + tma_bar_ptr=a_full + stage, + cache_policy=Int64(_EVICT_FIRST), + ) + cute.copy( + tma_sfa, + tAgSFA[(None, blk, kb, expert)], + tAsSFA[(None, stage)], + tma_bar_ptr=a_full + stage, + cache_policy=Int64(_EVICT_FIRST), + ) + return g + 1 + + @cute.jit + def _sm_domain(self, smid: Int32) -> Int32: + """Locality domain of the SM with ``%smid`` ``smid`` (0 beyond the table).""" + word = Int32(0) + for i in cutlass.range_constexpr(len(self.sm_domain_words)): + if (smid >> 5) == i: + word = Int32(self.sm_domain_words[i]) + return (word >> (smid & 31)) & 1 + + @cute.jit + def _dm_claim(self, ctr, own_die: Int32, phase: Int32, w1: Int32, nt: Int32, ns: Int32): + """Die-major claiming (whole warp; lane 0 claims): phase 0 / 1 = FC1 units of the own / the + other domain (domain d: row blocks [d nblk1 / 2, (d + 1) nblk1 / 2) of every tile, + tile-major), phase 2 / 3 = FC2 tasks of the own / the other domain (m-tiles + [d nblk2 / 2, ...) of every slice, slice-major), 4 = done. Returns the unit in the default + numbering (an FC1 unit ``tile * nblk1 + blk`` or ``w1 +`` the FC2 task) and the phase.""" + lane = cute.arch.lane_idx() + h1 = self.nblk1 // 2 + h2 = self.nblk2 // 2 + n1d = nt * h1 + n2d = ns * h2 + idx = Int32(0) + ph = phase + if lane == 0: + searching = Int32(1) + while searching != 0: + if ph >= 4: + idx = w1 + 2 * n2d # past the last unit + searching = Int32(0) + else: + die = own_die + if (ph & 1) == 1: + die = 1 - own_die + if ph < 2: + word = Int64(0) + if die == 1: + word = Int64(CTR_DIE1) + i = _atom_add_relaxed(ctr + word * 4, 1) + if i < n1d: + tile = i // h1 + idx = tile * self.nblk1 + die * h1 + (i - tile * h1) + searching = Int32(0) + else: + ph = ph + 1 + else: + word = Int64(CTR_FC2) + if die == 1: + word = Int64(CTR_FC2_DIE1) + i = _atom_add_relaxed(ctr + word * 4, 1) + if i < n2d: + s2 = i // h2 + idx = w1 + s2 * self.nblk2 + die * h2 + (i - s2 * h2) + searching = Int32(0) + else: + ph = ph + 1 + idx = cute.arch.shuffle_sync(idx, 0) + ph = cute.arch.shuffle_sync(ph, 0) + return idx, ph + + @cute.jit + def _expert_of(self, st, tile_idx_to_expert_idx, t2e_pre: Int32, tile: Int32): + """Expert of ``tile`` (warp-uniform): from the prefetched lanes for the first 32 tiles + (published routing: the shared-memory table).""" + expert = Int32(0) + if tile < 32: + expert = cute.arch.shuffle_sync(t2e_pre, tile) + else: + if cutlass.const_expr(self.published_routing): + expert = st.sT2e.get_tensor(cute.make_layout(self.route_t2e))[tile] + else: + expert = tile_idx_to_expert_idx[tile] + return expert + + @cute.jit + def _weight_producer( + self, + st, + thr_mma, + cl, + a_layout, + sfa_layout, + tma_a1, + mA1, + tma_sfa1, + mSFA1, + tma_a2, + mA2, + tma_sfa2, + mSFA2, + tile_idx_to_expert_idx, + ctr_ptr, + first: Int32, + t2e_pre: Int32, + w1: Int32, + total: Int32, + nctas: Int32, + nt: Int32, + ns: Int32, + ): + """Fetches the CTA's units from the work counter (one ahead, or with ``claim_free`` once + enough ring stages are free), publishes each with its first ring index and streams its + weight stages; the CTA that exits last resets the work and exit counters.""" + a_full, a_empty = st.a_full.data_ptr(), st.a_empty.data_ptr() + sA = st.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner) + sSFA = st.sSFA.get_tensor(sfa_layout) + a_cta = cute.make_layout(cute.slice_(cl, (0, 0, None, 0)).shape) + tiler = cute.slice_(self.mma_tiler, (None, 0, None)) + parts = [] + for mA, mSFA, tma_a, tma_sfa in ( + (mA1, mSFA1, tma_a1, tma_sfa1), + (mA2, mSFA2, tma_a2, tma_sfa2), + ): + tCgA = thr_mma.partition_A(cute.local_tile(mA, tiler, (None, None, None))) + tCgSFA = thr_mma.partition_A(cute.local_tile(mSFA, tiler, (None, None, None))) + tAsA, tAgA = cpasync.tma_partition( + tma_a, 0, a_cta, cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3) + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_sfa, 0, a_cta, cute.group_modes(sSFA, 0, 3), cute.group_modes(tCgSFA, 0, 3) + ) + parts.append((tAsA, tAgA, cute.filter_zeros(tAsSFA), cute.filter_zeros(tAgSFA))) + tAsA1, tAgA1, tAsSFA1, tAgSFA1 = parts[0] + tAsA2, tAgA2, tAsSFA2, tAgSFA2 = parts[1] + lane = cute.arch.lane_idx() + ctr = ctr_ptr.toint() + cur = cute.arch.shuffle_sync(first, 0) + own_die = Int32(0) + queue = Int32(0) + nloc = Int32(0) + if cutlass.const_expr(self.die_major): + nloc = Int32(cute.size(mA1, mode=[2])) // 2 + own_die = self._sm_domain(_smid()) + if cutlass.const_expr(self.published_routing): + # the claim before the tables: an own-domain FC1 unit, else on to the other's + h1 = self.nblk1 // 2 + if cur < nt * h1: + tile0 = cur // h1 + cur = tile0 * self.nblk1 + own_die * h1 + (cur - tile0 * h1) + else: + cur, queue = self._dm_claim(ctr, own_die, Int32(1), w1, nt, ns) + else: + cur, queue = self._dm_claim(ctr, own_die, queue, w1, nt, ns) + g = Int32(0) + j = Int32(0) + while cur < total: + nxt = Int32(0) + if cutlass.const_expr(self.claim_free == 0): + if cutlass.const_expr(self.die_major): + nxt, queue = self._dm_claim(ctr, own_die, queue, w1, nt, ns) + else: + if lane == 0: + nxt = _atom_add_relaxed(ctr, 1) + self._q_write(st, j, cur, g) + if cur < w1: + tile = cur // self.nblk1 + blk = cur - tile * self.nblk1 + expert = self._expert_of(st, tile_idx_to_expert_idx, t2e_pre, tile) + tblk = blk + texp = expert + if cutlass.const_expr(self.die_major): + # die-major slot of this row block: second half -> slot L + e + if blk >= self.nblk1 // 2: + tblk = blk - self.nblk1 // 2 + texp = expert + nloc + for kb in cutlass.range(self.nkb1, unroll=1): + g = self._load_stage( + tma_a1, + tma_sfa1, + tAsA1, + tAgA1, + tAsSFA1, + tAgSFA1, + tblk, + kb, + texp, + a_full, + a_empty, + g, + self.a_stage_bytes, + ) + else: + c, _s, t0, t1 = self._fc2_task(cur - w1, nt, ns) + tc = c + eoff = Int32(0) + if cutlass.const_expr(self.die_major): + if c >= self.nblk2 // 2: + tc = c - self.nblk2 // 2 + eoff = nloc + for t in cutlass.range(t0, t1, unroll=1): + expert = self._expert_of(st, tile_idx_to_expert_idx, t2e_pre, t) + eoff + for kb in cutlass.range(self.nkb2, unroll=1): + g = self._load_stage( + tma_a2, + tma_sfa2, + tAsA2, + tAgA2, + tAsSFA2, + tAgSFA2, + tc, + kb, + expert, + a_full, + a_empty, + g, + self.fc2_stage_bytes, + ) + if cutlass.const_expr(self.claim_free > 0): + # while FC1 units are left (they never wait on other CTAs) the next unit is claimed + # right away; FC2 tasks only once claim_free ring stages are free for them + peek = Int32(0) + if lane == 0: + if cutlass.const_expr(self.die_major): + # FC1 work is left for this CTA until it is past both domains' FC1 queues + if queue >= 2: + peek = w1 + else: + peek = _ld_relaxed_u32(ctr) + peek = cute.arch.shuffle_sync(peek, 0) + gw = g + (self.claim_free - 1) + if peek >= w1: + if gw >= self.num_ab_stage: + cute.arch.mbarrier_wait( + a_empty + gw % self.num_ab_stage, ((gw // self.num_ab_stage) & 1) ^ 1 + ) + if cutlass.const_expr(self.die_major): + nxt, queue = self._dm_claim(ctr, own_die, queue, w1, nt, ns) + else: + if lane == 0: + nxt = _atom_add_relaxed(ctr, 1) + cur = cute.arch.shuffle_sync(nxt, 0) + j += 1 + self._q_write(st, j, Int32(-1), g) + if lane == 0: + # every fetch of this CTA happens before its release; the last CTA's acquire sees all + if _atom_add_acq_rel(ctr + 4, 1) == nctas - 1: + _st_relaxed_u32(ctr, 0) + _st_relaxed_u32(ctr + 4, 0) + # a routing kernel may have published into it for a launch with tables too + _st_relaxed_u32(ctr + CTR_ROUTE * 4, 0) + if cutlass.const_expr(self.die_major): + _st_relaxed_u32(ctr + CTR_DIE1 * 4, 0) + _st_relaxed_u32(ctr + CTR_FC2 * 4, 0) + _st_relaxed_u32(ctr + CTR_FC2_DIE1 * 4, 0) + + @cute.jit + def _token_producer( + self, + st, + thr_mma, + cl, + b_layout, + tma_b1, + mB1, + tma_b2, + mB2, + h_sf_ptr, + out_ptr, + ctr_ptr, + orig_m: Int32, + nt: Int32, + ns: Int32, + w1: Int32, + bid, + nctas: Int32, + push, + ): + """The step's tokens (FC1's token operand) once; then, for every tile of every FC2 task, + once FC1 has stored all of the tile's rows, its rows and scales of each k-tile into the + token part of the ring stage that holds the matching weights.""" + S = self.num_ab_stage + b_cta = cute.make_layout(cute.slice_(cl, (0, None, 0, 0)).shape) + tiler = cute.slice_(self.mma_tiler, (0, None, None)) + sB = self._b_view(st, b_layout) + tCgB1 = thr_mma.partition_B(cute.local_tile(mB1, tiler, (None, None, None))) + tBsB1, tBgB1 = cpasync.tma_partition( + tma_b1, 0, b_cta, cute.group_modes(sB, 0, 3), cute.group_modes(tCgB1, 0, 3) + ) + tCgB2 = thr_mma.partition_B(cute.local_tile(mB2, tiler, (None, None, None))) + tBsB2, tBgB2 = cpasync.tma_partition( + tma_b2, 0, b_cta, cute.group_modes(sB, 0, 3), cute.group_modes(tCgB2, 0, 3) + ) + if cutlass.const_expr(self.published_routing and self.use_pdl): + # the tokens (and the output this warp zeroes) come with the routing grid's end + griddepcontrol_wait() + if w1 > 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(st.tok_full, self.x_bytes) + for kb in cutlass.range(self.nkb1, unroll=1): + cute.copy( + tma_b1, + tBgB1[(None, 0, kb, 0)], + tBsB1[(None, S + kb)], + tma_bar_ptr=st.tok_full, + ) + a_full, a_empty = st.a_full.data_ptr(), st.a_empty.data_ptr() + sfb2_base = self._region_addr(st, self.off_sfb2) + h_sf_base = h_sf_ptr.toint() + tile_sf_bytes = self.h_tile_sf_bytes + ctr = ctr_ptr.toint() + lane = cute.arch.lane_idx() + out_bytes = orig_m * (self.n2 * 2) + if nt == 0: + # no local expert: no FC2 task writes the output, CTA 0 zeroes it + if bid == 0: + out_base = out_ptr.toint() + for off in cutlass.range(lane * 16, out_bytes, 32 * 16, unroll=1): + _stg_v4_zero(out_base + Int64(off)) + if cutlass.const_expr(self.output_push_world > 0): + # No FC2 task pushes the output either, so every CTA pushes zero rows: the + # consumer polls every rank's slot of every token. + self._push_zero_rows(push, orig_m, Int32(bid) * 32 + lane, nctas * 32) + + slice_bytes = (out_bytes + _ZERO_SLICES * 16 - 1) // (_ZERO_SLICES * 16) * 16 + e = Int32(0) + u, g = self._q_read(st, e) + while u >= 0: + if cutlass.const_expr(self.combine_red): + if u < _ZERO_SLICES: + # zero slice u of the output, then count it on every tile: FC2 tasks of a tile + # add into the output only after all slices (the warp barrier orders every + # lane's stores before the releases) + lo = u * slice_bytes + hi = cutlass.min(lo + slice_bytes, out_bytes) + for off in cutlass.range(lo + lane * 16, hi, 32 * 16, unroll=1): + _stg_v4_zero(out_ptr.toint() + Int64(off)) + cute.arch.sync_warp() + for t in cutlass.range(lane, nt, 32, unroll=1): + _red_add_release(ctr + Int64(CTR_TILE_BASE + t) * 4, 1) + if u >= w1: + _c, _s, t0, t1 = self._fc2_task(u - w1, nt, ns) + for t in cutlass.range(t0, t1, unroll=1): + ready_addr = ctr + Int64(CTR_TILE_BASE + t) * 4 + ready = _ld_acquire_u32(ready_addr) + while ready < self.ready_target: + ready = _ld_acquire_u32(ready_addr) + cute.arch.fence_proxy("async.global") + h_sf_src = h_sf_base + Int64(t) * tile_sf_bytes + for kb in cutlass.range(self.nkb2, unroll=1): + # the weight producer expects this part's bytes on the same barrier; the + # stage is free once the MMAs of its previous use are done + stage = self._reserve_stage(a_empty, g) + cute.copy( + tma_b2, + tBgB2[(None, t * (_MMA_M // self.n_tok), kb, 0)], + tBsB2[(None, stage)], + tma_bar_ptr=a_full + stage, + ) + with cute.arch.elect_one(): + _bulk_g2s( + sfb2_base + stage * _SF_KBLK_BYTES, + h_sf_src + kb * _SF_KBLK_BYTES, + _SF_KBLK_BYTES, + (a_full + stage).toint(), + ) + g += 1 + # counter reset: every FC2 task of the tile adds 1 after its wait; the add + # that completes the count zeroes it + if lane == 0: + old = _atom_add_relaxed(ready_addr, 1) + if old + 1 == self.ready_target + self.nblk2: + _st_relaxed_u32(ready_addr, 0) + e += 1 + u, g = self._q_read(st, e) + + # ------------------------------------------------------------------ MMA + @cute.jit + def _mma( + self, + st, + tiled_mma, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + w1: Int32, + nt: Int32, + ns: Int32, + ): + S = self.num_ab_stage + NA, NWG = self.num_acc, self.num_epi_wg + a_full, a_empty = st.a_full.data_ptr(), st.a_empty.data_ptr() + acc_full, acc_empty = st.acc_full.data_ptr(), st.acc_empty.data_ptr() + acc_ptr = cute.arch.retrieve_tmem_ptr( + Float32, alignment=16, ptr_to_buffer_holding_addr=st.tmem_holding_buf + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tAcc_base = cute.make_tensor( + acc_ptr, tiled_mma.make_fragment_C(cute.append(acc_shape, self.nacc_total)).layout + ) + tCtSFA = cute.make_tensor( + cute.recast_ptr(acc_ptr + self.acc_cols, dtype=self.sf_dtype), + blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, self.mma_tiler, _SF_VEC, cute.slice_(sfa_layout, (None, None, None, 0)) + ), + ) + tCtSFB = cute.make_tensor( + cute.recast_ptr(acc_ptr + self.acc_cols + self.sfa_cols, dtype=self.sf_dtype), + blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, self.mma_tiler, _SF_VEC, cute.slice_(sfb_layout, (None, None, None, 0)) + ), + ) + sSFA = st.sSFA.get_tensor(sfa_layout) + s2t_sfa, tCsSFA, tCtSFA_s2t = _s2t_copy_and_partition(sSFA, tCtSFA, self.sf_dtype) + s2t_sfb, tCsSFB, tCtSFB_s2t = _s2t_copy_and_partition( + self._sfb_view(st, sfb_layout), tCtSFB, self.sf_dtype + ) + tCrA = tiled_mma.make_fragment_A(st.sA.get_tensor(a_layout.outer, swizzle=a_layout.inner)) + tCrB = tiled_mma.make_fragment_B(self._b_view(st, b_layout)) + nkp = cute.size(tCrA, mode=[2]) + + g = Int32(0) + jw0 = Int32(0) + jw1 = Int32(0) + f2 = Int32(0) + first1 = Int32(1) + e = Int32(0) + u, _g0 = self._q_read(st, e) + while u >= 0: + # FC1 unit: one tile block over nkb1 k-tiles, epilogue on warpgroup 0; FC2 task: the + # tiles of its slice, each over nkb2 k-tiles, epilogue on the task's warpgroup + owner = Int32(0) + ntl = Int32(1) + nst = Int32(self.nkb1) + fc2 = u >= w1 + if fc2: + owner = f2 % NWG + _c, _s, t0, t1 = self._fc2_task(u - w1, nt, ns) + ntl = t1 - t0 + nst = Int32(self.nkb2) + f2 += 1 + elif first1 > 0: + first1 = Int32(0) + cute.arch.mbarrier_wait(st.tok_full, 0) + cute.arch.mbarrier_wait(st.sfb_full, 0) + jj = jw0 + if owner != 0: + jj = jw1 + for _t in cutlass.range(ntl, unroll=1): + acc_stage = owner * NA + jj % NA + if jj >= NA: + cute.arch.mbarrier_wait(acc_empty + acc_stage, ((jj // NA) & 1) ^ 1) + tAcc = tAcc_base[(None, None, None, acc_stage)] + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for kb in cutlass.range(nst, unroll=1): + stage = g % S + cute.arch.mbarrier_wait(a_full + stage, (g // S) & 1) + bidx = stage + if u < w1: + bidx = S + kb + cute.copy(s2t_sfa, tCsSFA[(None, None, None, None, stage)], tCtSFA_s2t) + cute.copy(s2t_sfb, tCsSFB[(None, None, None, None, bidx)], tCtSFB_s2t) + for kp in cutlass.range_constexpr(nkp): + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[(None, None, kp)].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB[(None, None, kp)].iterator) + cute.gemm( + tiled_mma, + tAcc[(None, 0, 0)], + tCrA[(None, 0, kp, stage)], + tCrB[(None, 0, kp, bidx)], + tAcc[(None, 0, 0)], + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + with cute.arch.elect_one(): + tcgen05.commit(a_empty + stage) + g += 1 + with cute.arch.elect_one(): + tcgen05.commit(acc_full + acc_stage) + jj += 1 + if owner == 0: + jw0 = jj + else: + jw1 = jj + e += 1 + u, _g0 = self._q_read(st, e) + cute.arch.mbarrier_wait(st.tmem_free, 0) + cute.arch.dealloc_tmem(acc_ptr, self.tmem_cols) + + # ------------------------------------------------------------------ epilogues + @cute.jit + def _token_scales_to_smem(self, st, x_sf_ptr, orig_m: Int32, tidx: Int32): + """Token scales, linear [orig_m, k1 / 16], into the MMA's 128x4-blocked SMEM layout (rows + past orig_m zero), every load in flight before the first store.""" + n_tok = self.n_tok + nwords = self.nkb1 * 4 * n_tok + per_thread = (nwords + 127) // 128 + sf_row_bytes = self.k1 // _SF_VEC + x_sf_base = x_sf_ptr.toint() + sfb_base = self._region_addr(st, self.off_xsf) + vals = cute.make_rmem_tensor((per_thread,), Int32) + for t in cutlass.range_constexpr(per_thread): + i = tidx + t * 128 + r = i % n_tok + rest = i // n_tok + atom = rest % 4 + kb = cutlass.min(rest // 4, self.nkb1 - 1) + rr = cutlass.min(r, orig_m - 1) + vals[t] = _ldg_u32(x_sf_base + Int64(rr * sf_row_bytes + kb * 16 + atom * 4)) + for t in cutlass.range_constexpr(per_thread): + i = tidx + t * 128 + r = i % n_tok + rest = i // n_tok + atom = rest % 4 + kb = rest // 4 + if i < nwords: + v = vals[t] + if r >= orig_m: + v = Int32(0) + _sts_u32( + sfb_base + kb * _SF_KBLK_BYTES + atom * 512 + (r % 32) * 16 + (r // 32) * 4, + v, + ) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(st.sfb_full) + + @cute.jit + def _build_slot_table( + self, + st, + alpha2, + tile_idx_to_expert_idx, + expanded_idx_to_permuted_idx, + token_final_scales, + orig_m: Int32, + tidx: Int32, + ): + """Token-indexed layout: for every expanded slot ``x = token * topk + k`` of the step, its + tile (``expanded_idx_to_permuted_idx[x] // 128``, -1 for another rank's expert or a + padding token) and its scale ``alpha2[expert] * token_final_scales[token, k]`` (the + factor the permuted epilogue applies), by all epilogue threads; independent loads, one + round.""" + nthr = 128 * self.num_epi_wg + tl_t = st.sTl.get_tensor(cute.make_layout(self.n_x)) + sc_t = st.sSc.get_tensor(cute.make_layout(self.n_x)) + for it in cutlass.range_constexpr((self.n_x + nthr - 1) // nthr): + x = tidx + it * nthr + if x < self.n_x: + t = Int32(-1) + sc = Float32(0.0) + if x < orig_m * self.topk: + p = Int32(expanded_idx_to_permuted_idx[x]) + if p >= 0: + t = p // _MMA_M + token = x // self.topk + slot_k = x - token * self.topk + sc = alpha2[tile_idx_to_expert_idx[t]] * token_final_scales[(token, slot_k)] + tl_t[x] = t + sc_t[x] = sc + cute.arch.barrier(barrier_id=_EPI_BARRIER_BASE + self.num_epi_wg, number_of_threads=nthr) + + @cute.jit + def _published_tables( + self, st, alpha2, ids_t, ctr_ptr, orig_m: Int32, expert_offset: Int32, tidx: Int32 + ): + """Published routing: once the routing kernel has published every token's top-k + (``CTR_ROUTE``), this rank's active-expert mask with its word prefix counts, the tile -> + expert table (the active local experts in ascending order) and the tile count into shared + memory, by the epilogue warps; every warp of the CTA waits at the closing barrier.""" + nthr = 128 * self.num_epi_wg + bar = _EPI_BARRIER_BASE + self.num_epi_wg + if tidx < nthr: + ids_base = ids_t.iterator.toint() + num_local = Int32(cute.size(alpha2)) + am_base = st.sAm.data_ptr().toint() + am = st.sAm.get_tensor(cute.make_layout(_ROUTE_MASK_WORDS)) + apre = st.sApre.get_tensor(cute.make_layout(_ROUTE_MASK_WORDS)) + t2e = st.sT2e.get_tensor(cute.make_layout(self.route_t2e)) + ntt = st.sNt.get_tensor(cute.make_layout(4)) + if tidx < _ROUTE_MASK_WORDS: + am[tidx] = Int32(0) + if tidx == 0: + # relaxed polls, one acquire once complete + addr = ctr_ptr.toint() + Int64(CTR_ROUTE * 4) + done = _ld_relaxed_u32(addr) + while done < orig_m: + done = _ld_relaxed_u32(addr) + done = _ld_acquire_u32(addr) + cute.arch.barrier(barrier_id=bar, number_of_threads=nthr) + for x in cutlass.range(tidx, orig_m * self.topk, nthr, unroll=1): + local = _ld_relaxed_u32(ids_base + Int64(x) * 4) - expert_offset + if local >= 0 and local < num_local: + _red_shared_or(am_base + (local >> 5) * 4, _bit(local & 31)) + cute.arch.barrier(barrier_id=bar, number_of_threads=nthr) + nwords = (num_local + 31) >> 5 + if tidx < 32: + pre = Int32(0) + tot = Int32(0) + for w in cutlass.range(nwords, unroll=1): + c = cute.arch.popc(am[w]) + if w < tidx: + pre = pre + c + tot = tot + c + if tidx < nwords: + apre[tidx] = pre + if tidx == 0: + ntt[0] = tot + cute.arch.barrier(barrier_id=bar, number_of_threads=nthr) + for e in cutlass.range(tidx, num_local, nthr, unroll=1): + word = am[e >> 5] + bit = _bit(e & 31) + if (word & bit) != 0: + t2e[apre[e >> 5] + cute.arch.popc(word & (bit - 1))] = e + cute.arch.barrier() + + @cute.jit + def _published_slots( + self, + st, + alpha2, + ids_t, + vals_t, + orig_m: Int32, + expert_offset: Int32, + tidx: Int32, + ): + """Published routing: the token-indexed slot table (expanded slot -> tile, ``alpha2 x + weight``, as :meth:`_build_slot_table`) from the published top-k and the mask of + :meth:`_published_tables`, by the epilogue warps.""" + nthr = 128 * self.num_epi_wg + ids_base = ids_t.iterator.toint() + vals_base = vals_t.iterator.toint() + num_local = Int32(cute.size(alpha2)) + am = st.sAm.get_tensor(cute.make_layout(_ROUTE_MASK_WORDS)) + apre = st.sApre.get_tensor(cute.make_layout(_ROUTE_MASK_WORDS)) + tl_t = st.sTl.get_tensor(cute.make_layout(self.n_x)) + sc_t = st.sSc.get_tensor(cute.make_layout(self.n_x)) + for x in cutlass.range(tidx, self.n_x, nthr, unroll=1): + t = Int32(-1) + sc = Float32(0.0) + if x < orig_m * self.topk: + local = _ld_relaxed_u32(ids_base + Int64(x) * 4) - expert_offset + if local >= 0 and local < num_local: + word = am[local >> 5] + t = apre[local >> 5] + cute.arch.popc(word & (_bit(local & 31) - 1)) + sc = alpha2[local] * _ld_relaxed_f32(vals_base + Int64(x) * 4) + tl_t[x] = t + sc_t[x] = sc + cute.arch.barrier(barrier_id=_EPI_BARRIER_BASE + self.num_epi_wg, number_of_threads=nthr) + + @cute.jit + def _tile_scales(self, st, t: Int32, sbuf: Int32, et: Int32, wg: Int32): + """Token-indexed layout: token j's scale for tile ``t`` (0 if the tile's expert does not + serve it) into buffer ``sbuf`` of the warpgroup (thread j scans token j's top-k slots) and + the tile's token mask (warp ballot); the caller's warpgroup barrier publishes them.""" + tl_t = st.sTl.get_tensor(cute.make_layout(self.n_x)) + sc_t = st.sSc.get_tensor(cute.make_layout(self.n_x)) + sv_t = st.sSv.get_tensor( + cute.make_layout( + (self.n_tok, 2, self.num_epi_wg), stride=(1, self.n_tok, 2 * self.n_tok) + ) + ) + svm_t = st.sSvMask.get_tensor(cute.make_layout((2, self.num_epi_wg), stride=(1, 2))) + hit = Int32(0) + if et < self.n_tok: + sv = Float32(0.0) + for k in cutlass.range_constexpr(self.topk): + x = et * self.topk + k + if tl_t[x] == t: + sv = sc_t[x] + hit = Int32(1) + sv_t[et, sbuf, wg] = sv + if et < 32: + bits = _ballot(hit) + if et == 0: + svm_t[sbuf, wg] = bits + + @cute.jit + def _situ(self, up, gate, alpha_val, out): + """SiTU of 16 (up, gate) pairs, ``beta tanh(g / beta) sigmoid(g) lbeta tanh(u / lbeta)`` + with ``tanh(v) = 2 sigmoid(2 v) - 1``, in packed f32x2.""" + beta = self.situ_beta + linear_beta = self.situ_linear_beta + LOG2_E = Float32(1.4426950408889634) + neg_log2e_pair = (-LOG2_E, -LOG2_E) + one_pair = (Float32(1.0), Float32(1.0)) + inv_2beta = Float32(2.0 / beta) + two_beta = Float32(2.0 * beta) + neg_beta = Float32(-beta) + inv_2lbeta = Float32(2.0 / linear_beta) + two_lbeta = Float32(2.0 * linear_beta) + neg_lbeta = Float32(-linear_beta) + + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = (cute.math.exp2(neg[0], fastmath=True), cute.math.exp2(neg[1], fastmath=True)) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + alpha_pair = (Float32(alpha_val), Float32(alpha_val)) + for i in cutlass.range_constexpr(0, 16, 2): + gg = cute.arch.mul_packed_f32x2((gate[i], gate[i + 1]), alpha_pair) + uu = cute.arch.mul_packed_f32x2((up[i], up[i + 1]), alpha_pair) + sigmoid_g = _sigmoid(gg[0], gg[1]) + gs = _sigmoid(*cute.arch.mul_packed_f32x2(gg, (inv_2beta, inv_2beta))) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), (neg_beta, neg_beta) + ) + us = _sigmoid(*cute.arch.mul_packed_f32x2(uu, (inv_2lbeta, inv_2lbeta))) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), (neg_lbeta, neg_lbeta) + ) + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_g) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) + out[i] = out_pair[0] + out[i + 1] = out_pair[1] + + @cute.jit + def _epilogue( + self, + st, + tiled_mma, + acc_ptr, + h_ptr, + h_sf_ptr, + out_ptr, + alpha1, + alpha2, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + global_sf, + token_final_scales, + ctr_ptr, + ws_ptr, + orig_m: Int32, + w1: Int32, + nt: Int32, + ns: Int32, + et: Int32, + wg: Int32, + push, + ): + """Every epilogue warp walks the CTA's units. FC1 units (warpgroup 0): the accumulator + into a token-major staging buffer, SiTU and NVFP4 quantization of 16 outputs of one row + per thread, stores, tile counter release. FC2 tasks (alternating warpgroups): per tile, + FP32 ``scale * acc`` added into the warpgroup's token-major accumulator; then the task's + BF16-rounded sums added into the output (red combine), or the touched tokens' rows as the + task's partial, the m-tile's arrival, and on the last arrival the slice-ordered sum into + the output (slice combine).""" + NA, NWG = self.num_acc, self.num_epi_wg + n_tok = self.n_tok + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tAcc_base = cute.make_tensor( + acc_ptr, tiled_mma.make_fragment_C(cute.append(acc_shape, self.nacc_total)).layout + ) + copy_atom = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(n_tok)), Float32) + tAcc_epi = cute.flat_divide(tAcc_base[((None, None), 0, 0, None)], (_MMA_M, n_tok)) + tiled_t2r = tcgen05.make_tmem_copy(copy_atom, tAcc_epi[(None, None, 0, 0, 0)]) + thr_t2r = tiled_t2r.get_slice(et) + tTR_tAcc = thr_t2r.partition_S(tAcc_epi) + tTR_cAcc = thr_t2r.partition_D(cute.make_identity_tensor((_MMA_M, n_tok))) + rAcc = cute.make_rmem_tensor(tTR_cAcc.shape, Float32) + acc_full, acc_empty = st.acc_full.data_ptr(), st.acc_empty.data_ptr() + + # FC1 epilogue state + sacc_base = st.sAcc.data_ptr().toint() + col_bytes = self.sacc_row_floats * 4 + h_base = h_ptr.toint() + h_sf_base = h_sf_ptr.toint() + h_row_bytes = self.k2 // 2 + tile_sf_bytes = self.h_tile_sf_bytes + r = et >> 2 + q = et & 3 + norm_const = global_sf[0] + rcp_limit = Float32(1.0 / 6.0) + fp32_max = Float32(3.40282346638528859812e38) + up = cute.make_rmem_tensor((16,), Float32) + gate = cute.make_rmem_tensor((16,), Float32) + tCompute = cute.make_rmem_tensor((16,), Float32) + rSF = cute.make_rmem_tensor((2,), self.sf_dtype) + rQ = cute.make_rmem_tensor((16,), cutlass.Float4E2M1FN) + rPv = cute.make_rmem_tensor((2,), Float32) + ctr = ctr_ptr.toint() + + # FC2 epilogue state: the warpgroup's token-major FP32 accumulator ([token][channel], + # padded rows; warpgroup 0 reuses the FC1 staging buffer, whose last use precedes it) + tacc_base = sacc_base + if cutlass.const_expr(NWG > 1): + if wg == 1: + tacc_base = self._region_addr(st, self.off_acc1) + stok = st.sTok.get_tensor(cute.make_layout((n_tok, NWG), stride=(1, n_tok))) + sscale = st.sScale.get_tensor(cute.make_layout((n_tok, NWG), stride=(1, n_tok))) + sflag = st.sFlag.get_tensor(cute.make_layout(4)) + ws = ws_ptr.toint() + out_base = out_ptr.toint() + out_row_bytes = self.n2 * 2 + if cutlass.const_expr(self.output_push_world > 0): + # read once, past the grid dependency (as every read of the state must be) + push = self._push_target(push) + racc = cute.make_rmem_tensor((n_tok,), Float32) + # token-indexed layout: the FC2 task's per-token sums of this thread's channel (rAcc's + # columns), the per-tile scale buffers of the warpgroup and the buffer of the next tile + rtok = cute.make_rmem_tensor(tTR_cAcc.shape, Float32) + sv_t = st.sSv.get_tensor(cute.make_layout((n_tok, 2, NWG), stride=(1, n_tok, 2 * n_tok))) + svm_t = st.sSvMask.get_tensor(cute.make_layout((2, NWG), stride=(1, 2))) + sbuf = Int32(0) + + jj = Int32(0) + f2 = Int32(0) + e = Int32(0) + u, _g = self._q_read(st, e) + while u >= 0: + if u < w1: + if wg == 0: + tile = u // self.nblk1 + blk = u - tile * self.nblk1 + alpha_val = Float32(0.0) + if cutlass.const_expr(self.published_routing): + alpha_val = alpha1[ + st.sT2e.get_tensor(cute.make_layout(self.route_t2e))[tile] + ] + else: + alpha_val = alpha1[tile_idx_to_expert_idx[tile]] + tok = Int32(0) + n_valid = Int32(n_tok) + if cutlass.const_expr(self.token_indexed): + # row j of the tile = token j, for every token of the step + tok = r + else: + n_valid = cutlass.min( + cutlass.max(tile_idx_to_mn_limit[tile] - tile * _MMA_M, Int32(0)), + Int32(n_tok), + ) + if r < n_valid: + tok = permuted_idx_to_expanded_idx[tile * _MMA_M + r] // self.topk + acc_stage = jj % NA + cute.arch.mbarrier_wait(acc_full + acc_stage, (jj // NA) & 1) + cute.copy(tiled_t2r, tTR_tAcc[(None, None, None, 0, 0, acc_stage)], rAcc) + cute.arch.fence_view_async_tmem_load() + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(acc_empty + acc_stage) + for i in cutlass.range_constexpr(cute.size(rAcc)): + row, col = tTR_cAcc[i] + # W1 rows 32 q + i (gate) and 32 q + 16 + i (up) of the block go to + # staging rows 64 + 16 (q ^ 1) + i and 16 q + i: a warp's stores hit 32 + # distinct banks and the reads below are two-way instead of four-way. + g = ((row >> 4) & 1) ^ 1 + srow = (((row >> 5) ^ g) << 4) + (row & 15) + (g << 6) + _sts_f32(sacc_base + col * col_bytes + srow * 4, rAcc[i]) + self._wg_barrier(Int32(0)) + if r < n_valid: + # outputs 16 q .. 16 q + 15 of the block + src = sacc_base + tok * col_bytes + src_up = src + q * 64 + src_gate = src + 256 + (q ^ 1) * 64 + for vv in cutlass.range_constexpr(4): + g0, g1, g2, g3 = _lds_v4_f32(src_gate + vv * 16) + u0, u1, u2, u3 = _lds_v4_f32(src_up + vv * 16) + up[4 * vv] = u0 + up[4 * vv + 1] = u1 + up[4 * vv + 2] = u2 + up[4 * vv + 3] = u3 + gate[4 * vv] = g0 + gate[4 * vv + 1] = g1 + gate[4 * vv + 2] = g2 + gate[4 * vv + 3] = g3 + self._situ(up, gate, alpha_val, tCompute) + frg = cute.logical_divide(tCompute, cute.make_layout(_SF_VEC)) + acc_frg = frg.load() + abs_ir = math.absf(acc_frg.ir_value()) + abs_frg = type(acc_frg)(abs_ir, acc_frg.shape, acc_frg.dtype) + amax = abs_frg[None, 0].reduce(cute.ReductionOp.MAX, Float32(0.0), 0) + pv = cute.arch.mul_packed_f32x2( + (amax, Float32(0.0)), (rcp_limit, rcp_limit) + ) + pv = cute.arch.mul_packed_f32x2((pv[0], pv[1]), (norm_const, norm_const)) + rPv[0] = pv[0] + rPv[1] = pv[1] + rSF.store(rPv.load().to(self.sf_dtype)) + sf_f32 = rSF.load().to(Float32) + acc_scale = cute.arch.mul_packed_f32x2( + (cute.arch.rcp_approx(sf_f32[0]), cute.arch.rcp_approx(sf_f32[1])), + (norm_const, norm_const), + ) + scale0 = fmin(acc_scale[0], fp32_max, nan=True) + for ei in cutlass.range_constexpr(0, 16, 2): + tCompute[ei], tCompute[ei + 1] = cute.arch.mul_packed_f32x2( + (tCompute[ei], tCompute[ei + 1]), (scale0, scale0) + ) + rQ.store(tCompute.load().to(cutlass.Float4E2M1FN)) + rQw = cute.recast_tensor(rQ, Int32) + rSFb = cute.recast_tensor(rSF, cutlass.Uint8) + grow = tile * _MMA_M + r + _stg_v2_b32( + h_base + Int64(grow) * h_row_bytes + blk * 32 + q * 8, rQw[0], rQw[1] + ) + _stg_u8( + h_sf_base + + Int64(tile) * tile_sf_bytes + + blk * 512 + + (r % 32) * 16 + + (r // 32) * 4 + + q, + Int32(rSFb[0]), + ) + self._wg_barrier(Int32(0)) + # the barrier orders every thread's stores of this unit before the release + if et == 0: + _signal_add_release(ctr + Int64(CTR_TILE_BASE + tile) * 4) + jj += 1 + else: + owner = f2 % NWG + f2 += 1 + if owner == wg: + c, s, t0, t1 = self._fc2_task(u - w1, nt, ns) + mask = Int32(0) + if cutlass.const_expr(self.token_indexed): + # accumulator column j = token j in every tile: per-token FP32 sums in + # registers (tile order), then this thread's channel of every token row + # into the warpgroup's token-major staging buffer for the combine + for i in cutlass.range_constexpr(cute.size(rAcc)): + rtok[i] = Float32(0.0) + for t in cutlass.range(t0, t1, unroll=1): + self._tile_scales(st, t, sbuf, et, wg) + self._wg_barrier(wg) + acc_stage = wg * NA + jj % NA + cute.arch.mbarrier_wait(acc_full + acc_stage, (jj // NA) & 1) + cute.copy( + tiled_t2r, tTR_tAcc[(None, None, None, 0, 0, acc_stage)], rAcc + ) + cute.arch.fence_view_async_tmem_load() + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(acc_empty + acc_stage) + for i in cutlass.range_constexpr(cute.size(rAcc)): + row, col = tTR_cAcc[i] + rtok[i] = rtok[i] + rAcc[i] * sv_t[col, sbuf, wg] + mask = mask | svm_t[sbuf, wg] + sbuf = 1 - sbuf + jj += 1 + for i in cutlass.range_constexpr(cute.size(rAcc)): + row, col = tTR_cAcc[i] + _sts_f32(tacc_base + col * col_bytes + row * 4, rtok[i]) + else: + # this thread's channel column of the token accumulator + for tk in cutlass.range_constexpr(n_tok): + _sts_f32(tacc_base + tk * col_bytes + et * 4, Float32(0.0)) + for t in cutlass.range(t0, t1, unroll=1): + n_valid = cutlass.min( + cutlass.max(tile_idx_to_mn_limit[t] - t * _MMA_M, Int32(0)), + Int32(n_tok), + ) + scale = Float32(0.0) + token = Int32(0) + if et < n_valid: + expanded = permuted_idx_to_expanded_idx[t * _MMA_M + et] + token = expanded // self.topk + slot_k = expanded - token * self.topk + scale = ( + alpha2[tile_idx_to_expert_idx[t]] + * token_final_scales[(token, slot_k)] + ) + # the previous tile's accumulation still reads the previous metadata + self._wg_barrier(wg) + if et < n_valid: + stok[et, wg] = token + sscale[et, wg] = scale + self._wg_barrier(wg) + acc_stage = wg * NA + jj % NA + cute.arch.mbarrier_wait(acc_full + acc_stage, (jj // NA) & 1) + cute.copy( + tiled_t2r, tTR_tAcc[(None, None, None, 0, 0, acc_stage)], rAcc + ) + cute.arch.fence_view_async_tmem_load() + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(acc_empty + acc_stage) + for i in cutlass.range_constexpr(cute.size(rAcc)): + row, col = tTR_cAcc[i] + if col < n_valid: + addr = tacc_base + stok[col, wg] * col_bytes + row * 4 + _sts_f32(addr, _lds_f32(addr) + rAcc[i] * sscale[col, wg]) + for i in cutlass.range(n_valid, unroll=1): + mask = mask | (Int32(1) << stok[i, wg]) + jj += 1 + if cutlass.const_expr(self.combine_red): + # BF16-rounded per-token sums of the task added into the output: 16 + # threads per touched token row, 8 channels (16 bytes) each + self._wg_barrier(wg) + ch8 = (et % 16) * 8 + for rb in cutlass.range_constexpr(0, n_tok, 8): + tk = rb + et // 16 + if ((mask >> tk) & 1) != 0: + _red_add_v8_bf16_from_f32_smem( + out_base + Int64(tk) * out_row_bytes + (c * _MMA_M + ch8) * 2, + tacc_base + tk * col_bytes + ch8 * 4, + ) + # the next task's accumulator reset must not overtake these reads + self._wg_barrier(wg) + if cutlass.const_expr(self.output_push_world > 0): + self._red_push( + ctr, out_base, out_row_bytes, sflag, c, ns, orig_m, et, wg, push + ) + else: + self._fc2_combine( + ws, + ctr, + out_base, + out_row_bytes, + tacc_base, + col_bytes, + sflag, + racc, + mask, + c, + s, + ns, + orig_m, + et, + wg, + push, + ) + e += 1 + u, _g = self._q_read(st, e) + + @cute.jit + def _red_push( + self, + ctr, + out_base, + out_row_bytes, + sflag, + c, + ns, + orig_m, + et, + wg, + push, + ): + """Red combine with the output push: the task's arrival on its m-tile (acq_rel, after the + warpgroup barrier that follows its reds); the task that completes the m-tile's ``ns`` + arrivals reads the m-tile's BF16 output rows of every token (every task's reds precede + its arrival) and stores them into the token's slot of the all-reduce, as + :meth:`_push_rows` does for the slice combine (``push``: :meth:`_push_target`).""" + if et == 0: + old = _atom_add_acq_rel(ctr + Int64(CTR_MTILE_BASE + c) * 4, 1) + flag = Int32(0) + if old == ns - 1: + flag = Int32(1) + sflag[wg] = flag + self._wg_barrier(wg) + if sflag[wg] != 0: + push_mc, push_rank, cur, nbytes = push + # 16 threads per token row, 8 channels (16 bytes) each, as the reds + ch8 = (et % 16) * 8 + for rb in cutlass.range_constexpr(0, self.n_tok, 8): + tk = rb + et // 16 + if tk < orig_m: + w = _ldg_cg_v4(out_base + Int64(tk) * out_row_bytes + (c * _MMA_M + ch8) * 2) + dst = pushed_ar.push_slot( + push_mc, + push_mc, + cur, + nbytes, + tk, + push_rank, + c * _MMA_M + ch8, + self.output_push_world, + self.n2, + True, + ) + pushed_ar.push_v4(dst, w[0], w[1], w[2], w[3], True) + if et == 0: + _st_relaxed_u32(ctr + Int64(CTR_MTILE_BASE + c) * 4, 0) + # sflag is reused by the next task's arrival + self._wg_barrier(wg) + + @cute.jit + def _fc2_combine( + self, + ws, + ctr, + out_base, + out_row_bytes, + tacc_base, + col_bytes, + sflag, + racc, + mask, + c, + s, + ns, + orig_m, + et, + wg, + push, + ): + """Slice combine: the task's partial (rows of the tokens it touched, each thread its + channel) and the m-tile's arrival; the last of the m-tile's ``ns`` tasks sums the slices in + slice order and writes the BF16 output rows of every token.""" + n_tok = self.n_tok + pslot = c * self.max_slices + s + part = ws + (Int64(self.part_word_base) + Int64(pslot) * (n_tok * _MMA_M)) * 4 + for tk in cutlass.range_constexpr(n_tok): + _stg_f32_pred( + part + (tk * _MMA_M + et) * 4, + _lds_f32(tacc_base + tk * col_bytes + et * 4), + (mask >> tk) & 1, + ) + if et == 0: + _st_relaxed_u32(ws + Int64(self.mask_word_base + pslot) * 4, mask) + self._wg_barrier(wg) + # arrival (release of every thread's stores, ordered by the barrier); the last of the + # m-tile's ns tasks combines + if et == 0: + old = _atom_add_acq_rel(ctr + Int64(CTR_MTILE_BASE + c) * 4, 1) + flag = Int32(0) + if old == ns - 1: + flag = Int32(1) + sflag[wg] = flag + self._wg_barrier(wg) + if sflag[wg] != 0: + for tk in cutlass.range_constexpr(n_tok): + racc[tk] = Float32(0.0) + # slices in order, _COMBINE_CHUNK at a time with every load of the chunk in flight + # (predicated on the slice's token mask, 0.0 where absent) + rv = cute.make_rmem_tensor((_COMBINE_CHUNK * n_tok,), Float32) + # lane l holds slice l's token mask (ns <= max_slices <= 32) + lane = cute.arch.lane_idx() + my_mask = Int32(0) + if lane < ns: + my_mask = _ld_relaxed_u32( + ws + Int64(self.mask_word_base + c * self.max_slices + lane) * 4 + ) + for s0 in cutlass.range(0, ns, _COMBINE_CHUNK, unroll=1): + for q in cutlass.range_constexpr(_COMBINE_CHUNK): + sq = cutlass.min(s0 + q, ns - 1) + ps2 = c * self.max_slices + sq + m2 = cute.arch.shuffle_sync(my_mask, sq) + if s0 + q >= ns: + m2 = Int32(0) + p2 = ws + (Int64(self.part_word_base) + Int64(ps2) * (n_tok * _MMA_M)) * 4 + for tk in cutlass.range_constexpr(n_tok): + rv[q * n_tok + tk] = _ldg_cg_f32_pred( + p2 + (tk * _MMA_M + et) * 4, (m2 >> tk) & 1 + ) + for q in cutlass.range_constexpr(_COMBINE_CHUNK): + for tk in cutlass.range_constexpr(n_tok): + racc[tk] = racc[tk] + rv[q * n_tok + tk] + for tk in cutlass.range_constexpr(n_tok): + if tk < orig_m: + _stg_bf16_from_f32( + out_base + Int64(tk) * out_row_bytes + (c * _MMA_M + et) * 2, + racc[tk], + ) + if cutlass.const_expr(self.output_push_world > 0): + self._push_rows(racc, push, c, orig_m, et) + if et == 0: + _st_relaxed_u32(ctr + Int64(CTR_MTILE_BASE + c) * 4, 0) + + @cute.jit + def _push_target(self, push): + """``(push_mc, push_rank, cur, nbytes)`` of the output push: the multicast base, this + rank, and the all-reduce call's current buffer and bytes per buffer + (:func:`k3_lamport_ar.lamport_state`, valid past the grid dependency).""" + push_mc, push_flags, push_rank = push + cur, _, nbytes, _, _, _, _, _ = pushed_ar.lamport_state(push_flags) + return push_mc, push_rank, cur, nbytes + + @cute.jit + def _push_rows(self, racc, push, c, orig_m: Int32, et: Int32): + """Output push: lane groups of eight gather their eight channels of every token (BF16) + and the group's first lane stores them into the token's slot of the all-reduce + (``push``: :meth:`_push_target`).""" + push_mc, push_rank, cur, nbytes = push + lane = cute.arch.lane_idx() + group = lane - (lane & 7) + for tk in cutlass.range_constexpr(self.n_tok): + vals = [cute.arch.shuffle_sync(racc[tk], group + i) for i in range(8)] + if (lane & 7) == 0 and tk < orig_m: + words = [_bf16x2_bits(vals[2 * i], vals[2 * i + 1]) for i in range(4)] + dst = pushed_ar.push_slot( + push_mc, + push_mc, + cur, + nbytes, + Int32(tk), + push_rank, + c * _MMA_M + et, + self.output_push_world, + self.n2, + True, + ) + pushed_ar.push_v4(dst, words[0], words[1], words[2], words[3], True) + + @cute.jit + def _push_zero_rows(self, push, orig_m: Int32, gtid: Int32, nthreads: Int32): + """Output push without a local expert: the vectors :meth:`_push_rows` stores for an + all-zero output (``+0`` words), grid-strided over the token warps.""" + push_mc, push_rank, cur, nbytes = self._push_target(push) + vecs = self.n2 // 8 + zero = Int32(0) + for i in cutlass.range(gtid, orig_m * vecs, nthreads, unroll=1): + tk = i // vecs + dst = pushed_ar.push_slot( + push_mc, + push_mc, + cur, + nbytes, + tk, + push_rank, + (i - tk * vecs) * 8, + self.output_push_world, + self.n2, + True, + ) + pushed_ar.push_v4(dst, zero, zero, zero, zero, True) + + # ------------------------------------------------------------------ wrapper + @cute.jit + def wrapper( + self, + x_ptr: cute.Pointer, + w1_ptr: cute.Pointer, + x_sf_ptr: cute.Pointer, + w1_sf_ptr: cute.Pointer, + h_ptr: cute.Pointer, + h_sf_ptr: cute.Pointer, + alpha1_ptr: cute.Pointer, + w2_ptr: cute.Pointer, + w2_sf_ptr: cute.Pointer, + alpha2_ptr: cute.Pointer, + out_ptr: cute.Pointer, + tile_idx_to_group_idx_ptr: cute.Pointer, + tile_idx_to_mn_limit_ptr: cute.Pointer, + permuted_idx_to_expanded_idx_ptr: cute.Pointer, + num_non_exiting_tiles_ptr: cute.Pointer, + global_sf_ptr: cute.Pointer, + token_final_scales_ptr: cute.Pointer, + expanded_idx_to_permuted_idx_ptr: cute.Pointer, + ctr_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + orig_m: Int32, + m: Int32, + num_local: Int32, + expert_offset: Int32, + push_mc: Int64, + push_flags: Int64, + push_rank: Int32, + num_ctas: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """``x`` [orig_m, k1] tokens and ``x_sf`` their linear scales; ``h`` / ``h_sf`` the FC1 + output buffers ([m, k2] rows, 128x4-blocked scales); ``out`` the BF16 [orig_m, k1] output + (fully written); ``ctr`` int32 [counter_words(m // 128)] zeros; ``ws`` + [workspace_words(...)] scratch of the slice combine; ``expanded_idx_to_permuted_idx`` + [orig_m * topk] (read with the token-indexed layout only; published routing: the top-k ids, + ``token_final_scales`` the weights, ``expert_offset`` this rank's first expert). Die-major + weights are ``[2 L, rows / 2, cols]`` (``num_local`` is L).""" + n1, k1, n2, k2 = self.n1, self.k1, self.n2, self.k2 + L = num_local + num_tiles = m // _MMA_M + + def kmajor(ptr, rows, cols, batch): + return cute.make_tensor( + ptr, layout=cute.make_ordered_layout((rows, cols, batch), order=(1, 0, 2)) + ) + + def blocked_sf(ptr, rows, cols, batch): + return cute.make_tensor( + ptr, + layout=cute.make_ordered_layout( + (32, 4, rows // 128, 4, cols // (_SF_VEC * 4), batch), order=(2, 1, 4, 0, 3, 5) + ), + ) + + wn1, wn2, wl = n1, n2, L + if cutlass.const_expr(self.die_major): + wn1, wn2, wl = n1 // 2, n2 // 2, 2 * L + self( + kmajor(x_ptr, orig_m, k1, 1), + kmajor(w1_ptr, wn1, k1, wl), + blocked_sf(w1_sf_ptr, wn1, k1, wl), + kmajor(h_ptr, m, k2, 1), + kmajor(w2_ptr, wn2, k2, wl), + blocked_sf(w2_sf_ptr, wn2, k2, wl), + x_sf_ptr, + h_ptr, + h_sf_ptr, + out_ptr, + cute.make_tensor(alpha1_ptr, layout=cute.make_layout((L,))), + cute.make_tensor(alpha2_ptr, layout=cute.make_layout((L,))), + cute.make_tensor(tile_idx_to_group_idx_ptr, layout=cute.make_layout((num_tiles,))), + cute.make_tensor(tile_idx_to_mn_limit_ptr, layout=cute.make_layout((num_tiles,))), + cute.make_tensor(permuted_idx_to_expanded_idx_ptr, layout=cute.make_layout((m,))), + cute.make_tensor( + expanded_idx_to_permuted_idx_ptr, layout=cute.make_layout((orig_m * self.topk,)) + ), + cute.make_tensor(num_non_exiting_tiles_ptr, layout=cute.make_layout((1,))), + cute.make_tensor(global_sf_ptr, layout=cute.make_layout((1,))), + cute.make_tensor( + token_final_scales_ptr, + layout=cute.make_ordered_layout((orig_m, self.topk), order=(1, 0)), + ), + ctr_ptr, + ws_ptr, + orig_m, + expert_offset, + push_mc, + push_flags, + push_rank, + num_ctas, + stream, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_small_batch_moe_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_small_batch_moe_op.py new file mode 100644 index 000000000000..e195b3240bc5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_small_batch_moe_op.py @@ -0,0 +1,391 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""":func:`k3r_nvfp4_small_batch_moe`: the small-batch NVFP4 routed-experts kernel. + +Wraps :class:`RubinSmallBatchFusedMoeKernel` (FC1 + SiTU + NVFP4 requantization + FC2 + combine +of at most 32 tokens in one persistent grid). The caller passes the tile-128 ``moe_sort`` tables +and a counter buffer from :func:`new_small_batch_moe_counters` that it keeps for the layer (every +launch leaves it zero again). +""" + +import functools + +import torch + +from tensorrt_llm._torch.cute_dsl_utils import sm_reserve as global_sm_reserve + +__all__ = [ + "SMALL_BATCH_MOE_MAX_TOKENS", + "die_major_layout", + "k3r_nvfp4_small_batch_moe", + "new_small_batch_moe_counters", + "small_batch_moe_route_word", + "small_batch_moe_supports", + "small_batch_moe_takes_published", +] + +#: Largest token count of the kernel (MMA N 32). +SMALL_BATCH_MOE_MAX_TOKENS = 32 +#: Tiles the counter buffer holds (moe_sort's padded rows / 128 at <= 32 tokens). +_MAX_TILES = 1024 +#: (claim_free, max_slices) of the expert-parallel (permuted) layout and of the token-indexed +#: layout of MoE-TP slices. Small intermediates make FC2 tasks short, so the token-indexed schedule +#: claims an FC2 task only once 4 ring stages are free for it and splits the tiles into more +#: slices. +_PERMUTED_SCHEDULE = (0, 16) +_TOKEN_INDEXED_SCHEDULE = (4, 24) +#: Ring stage cap: die-major weights stream from the SM's own domain and keep more stages busy. +_RING_STAGES = 8 +_DIE_MAJOR_RING_STAGES = 12 + + +def small_batch_moe_supports(intermediate_size: int, hidden_size: int) -> bool: + """Whether the kernel serves experts of this per-rank shape (NVFP4, gated FC1).""" + from .k3_small_batch_moe import RubinSmallBatchFusedMoeKernel + + # whole 256-wide FC2 k-tiles: the FC1-output scale buffer then needs no zero padding + return intermediate_size % 256 == 0 and RubinSmallBatchFusedMoeKernel.supports( + SMALL_BATCH_MOE_MAX_TOKENS, 2 * intermediate_size, hidden_size + ) + + +def small_batch_moe_takes_published(num_local_experts: int) -> bool: + """Whether the kernel builds its tables from a published top-k (``published_routing``) of + ``num_local_experts`` local experts.""" + from .k3_small_batch_moe import MAX_PUBLISHED_LOCAL_EXPERTS + + return num_local_experts <= MAX_PUBLISHED_LOCAL_EXPERTS + + +def die_major_layout(t: torch.Tensor) -> torch.Tensor: + """``[L, R, C]`` expert weights or block scales -> ``[2 L, R / 2, C]``: slot e holds rows + ``[0, R / 2)`` of expert e, slot L + e rows ``[R / 2, R)``. The halves keep whole 128-row scale + blocks, so the 128x4-blocked scales split the same way.""" + rows = t.size(1) + assert rows % 256 == 0, "die-major halves need whole 128-row blocks" + return torch.cat([t[:, : rows // 2], t[:, rows // 2 :]], dim=0).contiguous() + + +def new_small_batch_moe_counters(device: torch.device) -> torch.Tensor: + """A zeroed counter buffer for one caller.""" + from .k3_small_batch_moe import counter_words + + return torch.zeros(counter_words(_MAX_TILES), dtype=torch.int32, device=device) + + +def small_batch_moe_route_word(counters: torch.Tensor) -> torch.Tensor: + """The int32 word of ``counters`` into which a routing kernel publishes the step's tokens + for a launch with ``published_routing``.""" + from .k3_small_batch_moe import CTR_ROUTE + + return counters[CTR_ROUTE : CTR_ROUTE + 1] + + +@functools.cache +def _persistent_ctas(device_index: int, sm_reserve: int) -> int: + """One CTA per SM less the ``sm_reserve`` SMs left free, never below one.""" + num_sms = torch.cuda.get_device_properties(device_index).multi_processor_count + return max(1, num_sms - sm_reserve) + + +class _SmallBatchMoeRunner: + """Compiles (once per configuration) and launches :class:`RubinSmallBatchFusedMoeKernel`.""" + + kernel_cache: dict = {} + + @classmethod + def run( + cls, + x: torch.Tensor, + x_sf: torch.Tensor, + w1: torch.Tensor, + w1_sf: torch.Tensor, + alpha1: torch.Tensor, + w2: torch.Tensor, + w2_sf: torch.Tensor, + alpha2: torch.Tensor, + out: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor | None, + tile_idx_to_mn_limit: torch.Tensor | None, + permuted_idx_to_expanded_idx: torch.Tensor | None, + expanded_idx_to_permuted_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor | None, + global_sf: torch.Tensor, + token_final_scales: torch.Tensor, + counters: torch.Tensor, + top_k: int, + situ_beta: float, + situ_linear_beta: float, + token_indexed: bool, + num_ctas: int, + deterministic_combine: bool, + die_major: bool, + published_routing: bool, + expert_offset: int, + push_mc: int, + push_flags: torch.Tensor | None, + push_rank: int, + push_world: int, + sm_reserve: int, + ) -> None: + import cutlass + import cutlass.cute as cute + from cuda.bindings import driver as cuda + from cutlass import Int32, Int64 + + from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, make_ptr + + from .k3_small_batch_moe import ( + MAX_SLICES, + RubinSmallBatchFusedMoeKernel, + counter_words, + n_tok_for, + workspace_words, + ) + + orig_m = x.size(0) + num_local, n1, k1 = w1.size(0), w1.size(1), w2.size(1) + if die_major: + num_local, n1, k1 = num_local // 2, 2 * n1, 2 * k1 + sm_domain_words = () + if die_major: + raise ValueError("die-major expert weights are not supported") + k2 = n1 // 2 + assert k2 % 256 == 0, "the FC1-output scale buffer is not zero-padded" + tables = ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + ) + if published_routing: + from tensorrt_llm._torch.moe.kernels.blackwell.kimi_k3_routing import sort_buffer_sizes + + if any(t is not None for t in tables): + raise ValueError("published routing takes no moe_sort tables") + if not small_batch_moe_takes_published(num_local): + raise ValueError(f"published routing serves no {num_local} local experts") + m = sort_buffer_sizes(orig_m, num_local)[1] + # the kernel reads none of the tables; any device address fills their slots + tables = (counters,) * len(tables) + elif any(t is None for t in tables): + raise ValueError("the moe_sort tables are required without published routing") + else: + m = permuted_idx_to_expanded_idx.size(0) + assert m // 128 <= _MAX_TILES and counters.numel() >= counter_words(m // 128) + assert expanded_idx_to_permuted_idx.numel() == orig_m * top_k + n_tok = n_tok_for(orig_m) + if num_ctas <= 0: + num_ctas = _persistent_ctas(x.device.index, sm_reserve + global_sm_reserve()) + h = torch.empty(m, k2 // 2, dtype=torch.uint8, device=x.device) + h_sf = torch.empty( + (m // 128) * RubinSmallBatchFusedMoeKernel.h_sf_tile_bytes(n1), + dtype=torch.uint8, + device=x.device, + ) + ws = h_sf # the red combine reads no workspace + if deterministic_combine: + ws = torch.empty( + workspace_words(k1, SMALL_BATCH_MOE_MAX_TOKENS, MAX_SLICES), + dtype=torch.int32, + device=x.device, + ) + + def ptr(dtype, t, align=None): + if align is None: + return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem) + return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=align) + + ptrs = [ + ptr(cutlass.Float4E2M1FN, x, 32), + ptr(cutlass.Float4E2M1FN, w1, 32), + ptr(cutlass.Float8E4M3FN, x_sf, 16), + ptr(cutlass.Float8E4M3FN, w1_sf, 16), + ptr(cutlass.Float4E2M1FN, h, 32), + ptr(cutlass.Float8E4M3FN, h_sf, 16), + ptr(cutlass.Float32, alpha1), + ptr(cutlass.Float4E2M1FN, w2, 32), + ptr(cutlass.Float8E4M3FN, w2_sf, 16), + ptr(cutlass.Float32, alpha2), + ptr(cutlass.BFloat16, out, 16), + ptr(cutlass.Int32, tables[0]), + ptr(cutlass.Int32, tables[1]), + ptr(cutlass.Int32, tables[2]), + ptr(cutlass.Int32, tables[3]), + ptr(cutlass.Float32, global_sf), + ptr(cutlass.Float32, token_final_scales), + ptr(cutlass.Int32, expanded_idx_to_permuted_idx), + ptr(cutlass.Int32, counters), + ptr(cutlass.Int32, ws, 16), + ] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + key = ( + n_tok, + n1, + k1, + top_k, + situ_beta, + situ_linear_beta, + token_indexed, + num_ctas, + deterministic_combine, + die_major, + sm_domain_words, + published_routing, + push_world, + TRTLLM_ENABLE_PDL, + ) + compiled = cls.kernel_cache.get(key) + if compiled is None: + claim_free, max_slices = ( + _TOKEN_INDEXED_SCHEDULE if token_indexed else _PERMUTED_SCHEDULE + ) + kernel = RubinSmallBatchFusedMoeKernel( + n_tok, + n1, + k1, + top_k, + situ_beta, + situ_linear_beta, + token_indexed=token_indexed, + claim_free=claim_free, + max_slices=max_slices, + max_ab_stage=_DIE_MAJOR_RING_STAGES if die_major else _RING_STAGES, + deterministic_combine=deterministic_combine, + use_pdl=TRTLLM_ENABLE_PDL, + die_major=die_major, + sm_domain_words=sm_domain_words, + published_routing=published_routing, + output_push_world=push_world, + ) + compiled = cute.compile( + kernel.wrapper, + *ptrs, + orig_m, + m, + num_local, + expert_offset, + Int64(push_mc), + Int64(push_flags.data_ptr() if push_flags is not None else 0), + Int32(push_rank), + num_ctas, + stream, + ) + cls.kernel_cache[key] = compiled + compiled( + *ptrs, + orig_m, + m, + num_local, + expert_offset, + Int64(push_mc), + Int64(push_flags.data_ptr() if push_flags is not None else 0), + Int32(push_rank), + stream, + ) + + +def k3r_nvfp4_small_batch_moe( + input: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor | None, + tile_idx_to_mn_limit: torch.Tensor | None, + permuted_idx_to_expanded_idx: torch.Tensor | None, + expanded_idx_to_permuted_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor | None, + global_sf: torch.Tensor, + token_final_scales: torch.Tensor, + counters: torch.Tensor, + top_k: int, + situ_beta: float, + situ_linear_beta: float, + token_indexed: bool, + num_ctas: int = 0, + deterministic_combine: bool = False, + die_major: bool = False, + published_routing: bool = False, + expert_offset: int = 0, + push_mc: int = 0, + push_flags: torch.Tensor | None = None, + push_rank: int = 0, + push_world: int = 0, + sm_reserve: int = 0, +) -> None: + """Routed NVFP4 SiTU MoE of at most 32 tokens: writes the sum of this rank's experts' + weighted down-projections into every element of ``output`` [tokens, hidden]. + + ``input`` / ``input_scale``: NVFP4 tokens and their linear scales; ``fc1_*`` / ``fc2_*``: the + local experts' NVFP4 weights, 128x4-blocked scales and per-expert alphas (FC1 rows in the Rubin + ``[gate16, up16]`` order); the six ``moe_sort`` tables of tile size 128; ``global_sf``: the FC2 + input scale; ``counters``: from :func:`new_small_batch_moe_counters`, zero again afterwards. + ``token_indexed``: the token-indexed FC1 output layout (MoE-TP slices of the intermediate). + ``num_ctas`` > 0 overrides the persistent CTA count, which is otherwise one per SM less + ``sm_reserve`` (SMs left to concurrent kernels) and ``cute_dsl_utils.sm_reserve()``. + ``deterministic_combine``: the schedule-independent slice combine instead of BF16 vector reds + (a test reference; the reds' order, so the output's bits, may change from run to run). + ``die_major``: the four weight tensors hold :func:`die_major_layout` bytes under their + ``[L, rows, cols]`` shapes (see ``RubinSmallBatchFusedMoeKernel``). ``published_routing`` + (``token_indexed``): instead of the tables, ``expanded_idx_to_permuted_idx`` holds the top-k + ids, which a routing kernel publishes into :func:`small_batch_moe_route_word` of ``counters`` + (the other four tables are None); ``expert_offset``: this rank's first expert. + ``push_world`` > 0: also push the output rows into slot ``(token, push_rank)`` of the + one-shot MNNVL all-reduce over ``push_world`` ranks whose multicast base is ``push_mc`` and + state ``push_flags`` (its consumer finishes it), with either combine. + """ + if die_major: + fc1_weight, fc1_weight_scale, fc2_weight, fc2_weight_scale = ( + t.view(2 * t.size(0), t.size(1) // 2, *t.shape[2:]) + for t in (fc1_weight, fc1_weight_scale, fc2_weight, fc2_weight_scale) + ) + _SmallBatchMoeRunner.run( + input, + input_scale, + fc1_weight, + fc1_weight_scale, + fc1_alpha, + fc2_weight, + fc2_weight_scale, + fc2_alpha, + output, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + expanded_idx_to_permuted_idx, + num_non_exiting_tiles, + global_sf, + token_final_scales, + counters, + top_k, + float(situ_beta), + float(situ_linear_beta), + token_indexed, + num_ctas, + deterministic_combine, + die_major, + published_routing, + expert_offset, + push_mc, + push_flags, + push_rank, + push_world, + sm_reserve, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_ull_front.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_ull_front.py new file mode 100644 index 000000000000..be2596eeff77 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/k3_ull_front.py @@ -0,0 +1,1482 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi-K3 MoE front for small decode steps (``M <= 32`` rows) over the TP group. + +Replaces, per MoE layer and rank, the router GEMM (``k3_router_exact`` / ``dsv3_router_gemm_op``), +the routing tail (``k3_route_tail``) and the replicated latent down-projection with its NVFP4 +quantization (``k3_latent_down_fp4``, the side stream it runs on and its fork / join) by two +kernels on one stream: + +* **Head** (:class:`K3UllFrontHeadKernel`, one grid, cluster size = the latent GEMM's split-K): + + - *Latent CTAs*: the ``k3_fp8w_lowm_gemm`` split-K MXFP8 GEMM on this rank's 128-row tiles of + latent_down (28 tiles over ``W`` ranks, rank ``r`` owning tiles ``[t0(r), t1(r))``), same + tactic as the replicated latent_down of ``k3_latent_down_fp4``. Its all-reduce push epilogue + (one-shot) multicasts the BF16 rows into slot ``(token, r)`` of the front workspace of every + rank (``-0`` sent as ``+0``). + - *Router CTAs*: ``k3_router_exact``'s arithmetic on all 896 experts (every rank computes all + logits, 12.8 MB of BF16 router weight streamed before the wait, and writes them into its own + workspace only, so the routing does not wait for an exchange). A cluster owns + ``8 * ROUTER_TILES`` experts; its CTAs split the 14 K slices of 512 and each warp runs one + ``mma.sync`` chain per (slice, 16-row tile, 8-expert tile) exactly as + ``router_gemm_mma_kernel`` does; the cluster's last CTA collects the 14 slice partials over + distributed shared memory, sums them in slice order from 0 and stores the FP32 logits into + ``[token][expert]``. Each 16-row tile therefore gets the bits ``dsv3_router_gemm_op`` gives + for those 16 rows. + + Both roles read their weights before ``griddepcontrol.wait`` and release the dependent right + after it. + +* **Route** (:class:`K3UllFrontRouteKernel`, PDL dependent of the head, launched while the head + runs and never waiting for it before it polls): a ``k3_route_tail`` cluster whose logits loads + poll the workspace (sentinel ``0x80000000``) instead of reading a tensor, so the top-16 and the + ``moe_sort`` tables of this rank are ``k3_route_tail``'s (bit-identical to ``noaux_tc_op`` + + ``moe_sort`` on the same logits); and gather CTAs that poll the latent slots of every rank and + write the NVFP4 quantization of the gathered BF16 latent (``fp4_quantize``'s arithmetic, linear + scale layout), i.e. what ``k3_latent_down_fp4`` returns. Every reader re-arms what it read. + +Workspace (:class:`UllFrontWorkspace`): an MNNVL multicast buffer of two halves; consecutive calls +alternate between them (the half is a launch argument through a constant flags tensor, so a CUDA +graph replays the same alternation). Rank ``r`` writes half ``h`` of a peer only after every rank +finished the call that last used ``h``: between two calls on a half there is the other half's call, +whose route kernel on ``r`` waited for data every peer pushed after its own previous call. +""" + +from __future__ import annotations + +import dataclasses +import functools +import os +import struct +import threading +from typing import Dict, Optional, Tuple + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cute.runtime import make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +from . import k3_fp8w_lowm_gemm as lowm +from . import k3_lamport_ar as lamport +from .k3_fused_router import _EVICT_FIRST, _mma_bf16, sort_buffer_sizes +from .k3_latent_fp8_gemm import ( + _abs_max_f32x16, + _e2m1x8, + _ldg_f32, + _nvfp4_scales_x2, + _rcp_approx_ftz, + _stg_u16, + _stg_v4_b32, +) +from .k3_route_tail import ( + MAX_CTAS, + TILE, + K3RouteTailKernel, + _cluster_arrive_release, + _cluster_wait_acquire, + _griddepcontrol_launch_dependents, + _griddepcontrol_wait, + _ld_nc_v4_f32, + _lds_b32, + _st_dsmem_b32, + _sts_b32, + _sts_v4_f32, + _syncwarp, +) +from .k3_router_exact import _ldg_v4_b32_if, _lds_v4_b32, _lds_v4_f32, _st_async_v4_f32 + +#: Kimi-K3 shapes. +HIDDEN = 7168 +LATENT = 3584 +NUM_EXPERTS = 896 +TOP_K = 16 +#: 128-row tiles of latent_down (its output columns), split over the ranks. +LATENT_TILES = LATENT // TILE +#: Largest rows per call. +M_MAX = 32 +#: Router K slices (``router_gemm_mma_kernel``'s warps) and their width. +NUM_SLICES = 14 +SLICE_K = HIDDEN // NUM_SLICES +CHUNKS = SLICE_K // 32 +_BOX_K = 64 +_BOXES_PER_SLICE = SLICE_K // _BOX_K +_BOX_BYTES = 8 * _BOX_K * 2 +#: One 8-expert x 16-row fragment set: 32 lanes x 4 FP32. +_FRAG_BYTES = 32 * 4 * 4 +#: 8-expert tiles per router cluster (28 clusters, 84 CTAs at split 3). +ROUTER_TILES = 4 +#: Gather clusters of the route kernel (of ``num_ctas`` CTAs each). +GATHER_CLUSTERS = 1 +SENTINEL = lamport.SENTINEL +#: NVFP4 group (fp4_quantize's scaling vector) and the gather item (32 latent columns). +_FP4_VEC = 16 +_ITEM_COLS = 32 + +assert HIDDEN % (NUM_SLICES * 32) == 0 + + +# --------------------------------------------------------------------------- +# Rank geometry +# --------------------------------------------------------------------------- + + +def latent_tiles_of(rank: int, world: int) -> Tuple[int, int]: + """``(first tile, tile count)`` of latent_down owned by ``rank`` (the first ``28 % W`` ranks + own one tile more).""" + base, rem = divmod(LATENT_TILES, world) + count = base + (1 if rank < rem else 0) + return rank * base + min(rank, rem), count + + +def slot_cols(world: int) -> int: + """Features of one (token, rank) latent slot: the largest per-rank tile count x 128.""" + base, rem = divmod(LATENT_TILES, world) + return (base + (1 if rem else 0)) * TILE + + +def experts_per_rank(world: int) -> int: + return NUM_EXPERTS // world + + +def supports_world(world: int) -> bool: + return ( + world >= 2 + and NUM_EXPERTS % world == 0 + and (NUM_EXPERTS // world) % 8 == 0 + and LATENT_TILES >= world + ) + + +def _align(a: int, b: int) -> int: + return (a + b - 1) // b * b + + +def logits_offset(world: int) -> int: + """Byte offset of the logits ``fp32 [M_MAX, 896]`` in a workspace half (after the latent + slots ``bf16 [M_MAX, W, slot_cols]``).""" + return _align(M_MAX * world * slot_cols(world) * 2, 1024) + + +def half_bytes(world: int) -> int: + return _align(logits_offset(world) + M_MAX * NUM_EXPERTS * 4, 1 << 16) + + +# --------------------------------------------------------------------------- +# PTX helpers +# --------------------------------------------------------------------------- + + +def _asm(res, operands, text, constraints, *, loc=None, ip=None): + return llvm.inline_asm( + res, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_volatile_b32(addr, *, loc=None, ip=None): + return Int32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.b32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_poll_v4(addr, *, loc=None, ip=None): + """16 bytes of the workspace with the poll's load.""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4)) + + +@dsl_user_op +def _red_add_release_gpu(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "red.release.gpu.global.add.s32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_sentinel_b32(addr, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n.reg .b32 t;\nmov.b32 t, 0x80000000;\nst.global.b32 [$0], t;\n}", + "l", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sanitize_f32_bits(word, *, loc=None, ip=None): + """``-0.0`` (the sentinel) sent as ``+0.0``.""" + return Int32( + _asm( + T.i32(), + [Int32(word).ir_value(loc=loc, ip=ip)], + "{\n.reg .pred p;\nsetp.eq.b32 p, $1, 0x80000000;\nselp.b32 $0, 0, $1, p;\n}", + "=r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bf16x2_to_f32x2(word, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Int32(word).ir_value(loc=loc, ip=ip)], + "{\n.reg .b32 l, h;\nshl.b32 l, $2, 16;\nand.b32 h, $2, 0xffff0000;\n" + "mov.b32 $0, l;\nmov.b32 $1, h;\n}", + "=f,=f,r", + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), res, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)), + ) + + +@cute.jit +def _poll_words(base: Int64, step: int, n: int): + """``n`` FP32 words at ``base + i * step`` bytes, polled (every load of a sweep in flight at + once) until none holds the sentinel.""" + words = cute.make_rmem_tensor((n,), Int32) + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for i in cutlass.range_constexpr(n): + words[i] = _ld_volatile_b32(base + Int64(i * step)) + for i in cutlass.range_constexpr(n): + if words[i] == Int32(SENTINEL): + ready = Int32(0) + return [_i2f(words[i]) for i in range(n)] + + +@cute.jit +def _poll_raw4(base: Int64, step: int, n: int): + """``n`` 16-byte vectors at ``base + i * step`` bytes polled until none holds the sentinel; + returns the ``4 n`` 32-bit words.""" + words = cute.make_rmem_tensor((4 * n,), Int32) + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for i in cutlass.range_constexpr(n): + w = _ld_poll_v4(base + Int64(i * step)) + for j in cutlass.range_constexpr(4): + words[4 * i + j] = w[j] + ready = ready & lamport.clean4(*w) + return [words[i] for i in range(4 * n)] + + +@cute.jit +def _poll_words4(base: Int64, step: int, n: int): + """``n`` 16-byte vectors at ``base + i * step`` bytes, polled like :func:`_poll_words`; returns + the ``4 n`` words as FP32.""" + words = cute.make_rmem_tensor((4 * n,), Int32) + ready = Int32(0) + while ready == 0: + ready = Int32(1) + for i in cutlass.range_constexpr(n): + w = _ld_poll_v4(base + Int64(i * step)) + for j in cutlass.range_constexpr(4): + words[4 * i + j] = w[j] + ready = ready & lamport.clean4(*w) + return [_i2f(words[i]) for i in range(4 * n)] + + +# --------------------------------------------------------------------------- +# Head: latent GEMM CTAs + router CTAs +# --------------------------------------------------------------------------- + + +class K3UllFrontHeadKernel: + """The head grid: the latent GEMM's ``n_tiles * split_k`` CTAs, then the router clusters. + + Args: + gemm: The latent GEMM (``K3Fp8wLowMGemmKernel`` with the one-shot push, ``ar_slot_n`` = + :func:`slot_cols`) of this rank's tiles. + router_tiles: 8-expert tiles per router cluster. + m_tiles: 16-row tiles of the router (``ceil(M_max / 16)``). + logits_off: Byte offset of the logits in a workspace half. + """ + + def __init__( + self, + gemm: "lowm.K3Fp8wLowMGemmKernel", + router_tiles: int, + m_tiles: int, + logits_off: int, + use_pdl: bool, + ): + self.gemm = gemm + self.S = gemm.split_k + self.threads = gemm.threads + self.num_warps = self.threads // 32 + self.rt = router_tiles + self.m_tiles = m_tiles + self.logits_off = logits_off + self.use_pdl = use_pdl + self.n_tiles_r = NUM_EXPERTS // 8 + self.router_clusters = (self.n_tiles_r + router_tiles - 1) // router_tiles + self.n_gemm_ctas = gemm.n_tiles * self.S + self.grid = self.n_gemm_ctas + self.router_clusters * self.S + # Slices of cluster rank r: [NUM_SLICES * r // S, NUM_SLICES * (r + 1) // S). + self.max_slices = max( + NUM_SLICES * (r + 1) // self.S - NUM_SLICES * r // self.S for r in range(self.S) + ) + last = self.S - 1 + self.last_slices = NUM_SLICES - NUM_SLICES * last // self.S + if self.max_slices > gemm.a_stages: + raise ValueError("the router uses one GEMM weight-stage barrier per local slice") + self.sw_bytes = self.max_slices * router_tiles * _BOXES_PER_SLICE * _BOX_BYTES + self.part_bytes = NUM_SLICES * m_tiles * router_tiles * _FRAG_BYTES + self.stage_bytes = M_MAX * 8 * router_tiles * 4 + # The router CTAs keep their buffers in the GEMM's shared storage from its weight ring on + # (one allocation for both roles): weights, slice partials, logits (checked at launch). + self.router_bytes = self.sw_bytes + self.part_bytes + self.stage_bytes + if m_tiles * router_tiles * 32 > self.threads: + raise ValueError("the router's combine needs one thread per fragment") + + @cute.jit + def __call__( + self, + w_ptr: cute.Pointer, + x_ptr: cute.Pointer, + ws_ptr: cute.Pointer, + m: Int32, + peer_ptrs: Int64, + mc_base: Int64, + flags: Int64, + rank: Int32, + rw_ptr: cute.Pointer, + x_addr: Int64, + stream: cuda.CUstream, + ): + gemm_args = self.gemm.kernel_args( + w_ptr, + x_ptr, + None, + None, + ws_ptr, + x_ptr, + x_ptr, + x_ptr, + x_ptr, + m, + peer_ptrs, + mc_base, + flags, + rank, + ) + tma_atom_r, tma_r, r_layout = self.router_args(rw_ptr) + # The weight ring starts at the first 1 KB boundary after the barriers. + avail = self.gemm.shared_storage.size_in_bytes() - 1024 + if cutlass.const_expr(self.router_bytes > avail): + raise ValueError( + f"router buffers ({self.router_bytes} B) exceed the GEMM storage ({avail} B)" + ) + self.kernel(*gemm_args, tma_atom_r, tma_r, r_layout, x_addr).launch( + grid=(self.grid, 1, 1), + block=(self.threads, 1, 1), + cluster=(self.S, 1, 1), + smem=self.gemm.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.jit + def router_args(self, rw_ptr: cute.Pointer): + """``(tma_atom, tma_tensor, smem_layout)`` of the router weight slice (8-row x 64-column + 128-byte-swizzled boxes, as ``k3_router_exact``).""" + mR = cute.make_tensor(rw_ptr, cute.make_layout((NUM_EXPERTS, HIDDEN), stride=(HIDDEN, 1))) + atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + boxes = self.max_slices * self.rt * _BOXES_PER_SLICE + r_layout = cute.tile_to_shape(atom, (8, _BOX_K, boxes), order=(0, 1, 2)) + tma_atom_r, tma_r = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mR, + cute.tile_to_shape(atom, (8, _BOX_K), order=(0, 1)), + (8, _BOX_K), + ) + return tma_atom_r, tma_r, r_layout + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_x: cute.CopyAtom, + mXt: cute.Tensor, + mWS: cute.Tensor, + mOut: cute.Tensor, + out1_ptr: cute.Pointer, + aux0_ptr: cute.Pointer, + aux1_ptr: cute.Pointer, + m: Int32, + a_layout: cute.ComposedLayout, + b_layout: cute.ComposedLayout, + sfa_layout: cute.Layout, + sfb_layout: cute.Layout, + x_layout: cute.Layout, + tma_atom_g: Optional[cute.CopyAtom], + mGt: Optional[cute.Tensor], + ar_peers: Optional[Int64], + ar_mc: Optional[Int64], + ar_flags: Optional[Int64], + ar_rank: Optional[Int32], + tma_atom_r: cute.CopyAtom, + mRt: cute.Tensor, + r_layout: cute.ComposedLayout, + x_addr: Int64, + ): + bidx, _, _ = cute.arch.block_idx() + storage = utils.SmemAllocator().allocate(self.gemm.shared_storage) + if bidx < self.n_gemm_ctas: + self.gemm.cta( + tiled_mma, + tma_atom_a, + mA, + tma_atom_x, + mXt, + mWS, + mOut, + out1_ptr, + aux0_ptr, + aux1_ptr, + m, + a_layout, + b_layout, + sfa_layout, + sfb_layout, + x_layout, + tma_atom_g, + mGt, + ar_peers, + ar_mc, + ar_flags, + ar_rank, + storage, + ) + else: + self.router_cta( + (bidx - self.n_gemm_ctas) // self.S, + storage, + tma_atom_r, + mRt, + r_layout, + x_addr, + m, + ar_flags, + ar_peers, + ar_rank, + ) + + @cute.jit + def router_cta( + self, + cl: Int32, + storage, + tma_atom_r: cute.CopyAtom, + mRt: cute.Tensor, + r_layout: cute.ComposedLayout, + x_addr: Int64, + m: Int32, + flags: Int64, + peer_ptrs: Int64, + ar_rank: Int32, + ): + """Router cluster ``cl``: experts ``[8 * rt * cl, 8 * rt * (cl + 1))`` of the router, + written to this rank's workspace only.""" + S = self.S + tidx, _, _ = cute.arch.thread_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + crank = Int32(0) + if cutlass.const_expr(S > 1): + crank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + # Barriers of the GEMM's storage (a_full: one per local slice) and its weight ring. + w_full = storage.a_full.data_ptr() + mbox_full = storage.mbox_full + ring = storage.sA.data_ptr() + w_base = ring.toint() + part_base = w_base + self.sw_bytes + stage_base = part_base + self.part_bytes + lo = (crank * NUM_SLICES) // S + hi = ((crank + 1) * NUM_SLICES) // S + n_sl = hi - lo + is_last = crank == S - 1 + tile0 = cl * self.rt + + if warp == 0: + with cute.arch.elect_one(): + for s in cutlass.range_constexpr(self.max_slices): + cute.arch.mbarrier_init(w_full + s, 1) + cute.arch.mbarrier_init(mbox_full, 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + if warp == 0: + with cute.arch.elect_one(): + if is_last: + # One arrival plus the other cluster ranks' slice partials. + cute.arch.mbarrier_arrive_and_expect_tx( + mbox_full, + (NUM_SLICES - self.last_slices) * self.m_tiles * self.rt * _FRAG_BYTES, + ) + if cutlass.const_expr(S > 1): + # Peers write into the last CTA's partials only after its barriers are set up. + cute.arch.cluster_arrive_relaxed() + + # The weight slices of this CTA, before the wait. + sW = cute.make_tensor( + cute.recast_ptr(ring, swizzle_=r_layout.inner, dtype=cutlass.BFloat16), r_layout.outer + ) + gW = cute.local_tile(mRt, (8, _BOX_K), (None, None)) + tWs, tWg = cpasync.tma_partition( + tma_atom_r, + 0, + cute.make_layout(1), + cute.group_modes(sW, 0, 2), + cute.group_modes(gW, 0, 2), + ) + if warp == 0: + for ls in cutlass.range(n_sl, unroll=1): + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + w_full + ls, self.rt * _BOXES_PER_SLICE * _BOX_BYTES + ) + for t in cutlass.range_constexpr(self.rt): + # Tiles past this rank's experts read zero rows (TMA bounds) and are not pushed. + for j in cutlass.range_constexpr(_BOXES_PER_SLICE): + cute.copy( + tma_atom_r, + tWg[(None, tile0 + t, (lo + ls) * _BOXES_PER_SLICE + j)], + tWs[(None, (ls * self.rt + t) * _BOXES_PER_SLICE + j)], + tma_bar_ptr=w_full + ls, + cache_policy=Int64(_EVICT_FIRST), + ) + + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_wait() + # Every earlier grid has completed: the dependent (the route kernel) may launch. + _griddepcontrol_launch_dependents() + + if cutlass.const_expr(S > 1): + cute.arch.cluster_wait() + + # router_gemm_mma_kernel fragments: lane (g, q) holds activation rows g and g + 8 of the + # 16-row tile and weight row g at k = chunk + 8 q .. + 8. + g = lane >> 2 + q = lane & 3 + mt_n = self.m_tiles + for item in cutlass.range(warp, n_sl * mt_n, self.num_warps, unroll=1): + ls = item // mt_n + mt = item - ls * mt_n + s = lo + ls + r0 = mt * 16 + g + v0 = Int32(0) + if r0 < m: + v0 = Int32(1) + v1 = Int32(0) + if r0 + 8 < m: + v1 = Int32(1) + a0_base = x_addr + Int64(r0 * (HIDDEN * 2) + (s * SLICE_K + q * 8) * 2) + a1_base = a0_base + Int64(8 * HIDDEN * 2) + acts = [] + for cc in cutlass.range_constexpr(CHUNKS): + a_off = Int64(cc * 64) + acts.append( + (_ldg_v4_b32_if(v0, a0_base + a_off), _ldg_v4_b32_if(v1, a1_base + a_off)) + ) + cute.arch.mbarrier_wait(w_full + ls, 0) + accs = [[Float32(0.0)] * 4 for _ in range(self.rt)] + for cc in cutlass.range_constexpr(CHUNKS): + a0, a1 = acts[cc] + for t in cutlass.range_constexpr(self.rt): + bw = self._wfrag(w_base, ls, t, g, q, cc) + acc = _mma_bf16((a0[0], a1[0], a0[1], a1[1]), bw[0], bw[1], accs[t]) + accs[t] = list(_mma_bf16((a0[2], a1[2], a0[3], a1[3]), bw[2], bw[3], acc)) + for t in cutlass.range_constexpr(self.rt): + acc = accs[t] + off = ((s * mt_n + mt) * self.rt + t) * _FRAG_BYTES + lane * 16 + if is_last: + _sts_v4_f32(part_base + off, acc[0], acc[1], acc[2], acc[3]) + else: + _st_async_v4_f32( + part_base + off, + acc[0], + acc[1], + acc[2], + acc[3], + mbox_full.toint(), + Int32(S - 1), + ) + + if is_last: + cute.arch.mbarrier_wait(mbox_full, 0) + cute.arch.barrier() + # Thread (m tile, expert tile, lane): the 14 slice partials in slice order from 0. + frags = mt_n * self.rt * 32 + if tidx < frags: + mt = tidx // (self.rt * 32) + t = (tidx // 32) - mt * self.rt + ln = tidx & 31 + tot = [Float32(0.0)] * 4 + for s in cutlass.range_constexpr(NUM_SLICES): + off = ((s * mt_n + mt) * self.rt + t) * _FRAG_BYTES + ln * 16 + v = _lds_v4_f32(part_base + off) + for i in cutlass.range_constexpr(4): + tot[i] = tot[i] + v[i] + for i in cutlass.range_constexpr(4): + row = mt * 16 + (ln >> 2) + (i >> 1) * 8 + col = t * 8 + (ln & 3) * 2 + (i & 1) + _sts_f32(stage_base + (row * (8 * self.rt) + col) * 4, tot[i]) + cute.arch.barrier() + cur, _, nbytes, _, _, _, _, _ = lamport.lamport_state(flags) + vecs = 2 * self.rt + for p in cutlass.range(tidx, m * vecs, self.threads, unroll=1): + row = p // vecs + v = p - row * vecs + local = tile0 * 8 + v * 4 + if local < NUM_EXPERTS: + w = _lds_v4_b32(stage_base + (row * (8 * self.rt) + v * 4) * 4) + s0 = _sanitize_f32_bits(w[0]) + s1 = _sanitize_f32_bits(w[1]) + s2 = _sanitize_f32_bits(w[2]) + s3 = _sanitize_f32_bits(w[3]) + off = Int64((row * NUM_EXPERTS + local) * 4) + rel = Int64(cur) * Int64(nbytes) + Int64(self.logits_off) + off + lamport.st_v4( + lamport.ld_nc_u64(peer_ptrs + Int64(ar_rank) * 8) + rel, + s0, + s1, + s2, + s3, + ) + if cutlass.const_expr(S > 1): + # No CTA leaves while a peer may still write into its shared memory. + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + @cute.jit + def _wfrag(self, w_base: Int32, ls: Int32, t: int, g: Int32, q: Int32, cc: int): + """Weight row ``g`` of expert tile ``t`` at chunk ``cc`` of local slice ``ls``.""" + box = (ls * self.rt + t) * _BOXES_PER_SLICE + cc // 2 + unit = (4 * (cc % 2) + q) ^ g + return _lds_v4_b32(w_base + box * _BOX_BYTES + g * 128 + (unit << 4)) + + +@dsl_user_op +def _sts_f32(addr, value, *, loc=None, ip=None): + _asm( + None, + [Int32(addr).ir_value(loc=loc, ip=ip), Float32(value).ir_value(loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + loc=loc, + ip=ip, + ) + + +# --------------------------------------------------------------------------- +# Route: routing tail on the polled logits + latent gather / NVFP4 +# --------------------------------------------------------------------------- + + +class K3UllFrontRouteKernel(K3RouteTailKernel): + """``k3_route_tail``'s cluster (top-16 + ``moe_sort`` of one rank) on the logits the head + pushes, then ``gather_clusters`` clusters that gather and quantize the latent. + + The grid is ``(1 + gather_clusters) * num_ctas`` CTAs; cluster 0 routes. + + ``publish`` (the experts kernel builds its tables itself, ``PublishedSortTables``): no + ``moe_sort`` tables and no cluster mask exchange; after each token's ids / weights are stored + the team adds 1 to the 32-bit word at ``num_tiles_addr`` (GPU-scope release), and every CTA + releases the dependent grid at its start, so the experts' CTAs can be resident and polling + before the top-k is published (they poll that word; their reads of the latent / output still + wait for this grid's end). + """ + + def __init__( + self, + num_ctas: int, + num_warps: int, + warps_per_token: int, + local_cap: int, + world: int, + use_pdl: bool, + publish: bool = False, + ): + super().__init__(num_ctas, num_warps, warps_per_token, 1, not publish, use_pdl, local_cap) + self.publish = publish + # k3_route_tail's approximate candidate pass (``tanh.approx`` scores pick the candidates, + # ``noaux_tc``'s exact score ranks them; the same selection, see + # K3RouteTailKernel.select_token). + self.approx = True + self.world = world + self.slot_n = slot_cols(world) + self.logits_off = logits_offset(world) + self.tiles_base, self.tiles_rem = divmod(LATENT_TILES, world) + self.gather_clusters = GATHER_CLUSTERS + self.grid = (1 + GATHER_CLUSTERS) * num_ctas + + def load_logits(self, row: Int64, g: Int32, stride: int, n: int): + """The logits polled from the workspace (all ``n`` loads in flight per sweep).""" + return _poll_words(row + Int64(g * 4), stride * 4, n) + + def load_logits4(self, row: Int64, lane: Int32, n: int): + return _poll_words4(row + Int64(lane * 16), 128 * 4, n) + + @cute.jit + def __call__( + self, + m: Int32, + uc_local: Int64, + flags: Int64, + bias_addr: Int64, + ids_addr: Int64, + vals_addr: Int64, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + scale_bits: Int64, + codes_addr: Int64, + sf_addr: Int64, + gsf_addr: Int64, + stream: cuda.CUstream, + ): + teams = self.teams + + @cute.struct + class SharedStorage: + bias: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, NUM_EXPERTS], 16] + mask: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.local_cap], 16] + masks: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, self.num_ctas * self.local_cap], 16 + ] + base: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.local_cap], 16] + ids: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, 32 * TOP_K], 16] + cand: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, self.num_warps * self.region_slots * 4], 16 + ] + top: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, teams * TOP_K * 2], 16] + xchg: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, teams * 16], 16] + warp_tot: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, self.num_warps], 16] + + self.shared_storage = SharedStorage + self.kernel( + m, + uc_local, + flags, + bias_addr, + ids_addr, + vals_addr, + e2p_addr, + p2e_addr, + tile_expert_addr, + tile_limit_addr, + total_addr, + num_tiles_addr, + local_offset, + local_num, + scale_bits, + codes_addr, + sf_addr, + gsf_addr, + ).launch( + grid=(self.grid, 1, 1), + block=(self.threads, 1, 1), + cluster=(self.num_ctas, 1, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + m: Int32, + uc_local: Int64, + flags: Int64, + bias_addr: Int64, + ids_addr: Int64, + vals_addr: Int64, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + scale_bits: Int64, + codes_addr: Int64, + sf_addr: Int64, + gsf_addr: Int64, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + if cutlass.const_expr(self.publish and self.use_pdl): + # the experts poll the published top-k and wait for this grid before the latent + _griddepcontrol_launch_dependents() + # The halves' flags never change: no ordering needed to read them. + cur, _, nbytes, _, _, _, _, _ = lamport.lamport_state(flags) + half = uc_local + Int64(cur) * Int64(nbytes) + if bidx < self.num_ctas: + self.route_cta( + bidx, + tidx, + m, + half + Int64(self.logits_off), + bias_addr, + ids_addr, + vals_addr, + e2p_addr, + p2e_addr, + tile_expert_addr, + tile_limit_addr, + total_addr, + num_tiles_addr, + local_offset, + local_num, + scale_bits, + ) + else: + self.gather_cta( + bidx - self.num_ctas, + tidx, + m, + half, + codes_addr, + sf_addr, + gsf_addr, + ) + if cutlass.const_expr(self.use_pdl): + # Complete only after the head (whose CTAs are done pushing by now). + _griddepcontrol_wait() + + @cute.jit + def route_cta( + self, + rank: Int32, + tidx: Int32, + m: Int32, + logits_addr: Int64, + bias_addr: Int64, + ids_addr: Int64, + vals_addr: Int64, + e2p_addr: Int64, + p2e_addr: Int64, + tile_expert_addr: Int64, + tile_limit_addr: Int64, + total_addr: Int64, + num_tiles_addr: Int64, + local_offset: Int32, + local_num: Int32, + scale_bits: Int64, + ): + """``K3RouteTailKernel.kernel`` without the grid-dependency wait before the logits (they + are polled), re-arming each token's logits once its team is done with them.""" + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + storage = utils.SmemAllocator().allocate(self.shared_storage) + bias_smem = storage.bias.data_ptr().toint() + mask_smem = storage.mask.data_ptr().toint() + masks_smem = storage.masks.data_ptr().toint() + base_smem = storage.base.data_ptr().toint() + ids_smem = storage.ids.data_ptr().toint() + cand0 = storage.cand.data_ptr().toint() + top0 = storage.top.data_ptr().toint() + xchg0 = storage.xchg.data_ptr().toint() + warp_tot_smem = storage.warp_tot.data_ptr().toint() + team = warp // self.team_warps + team_warp = warp - team * self.team_warps + cand_smem = cand0 + team * (self.team_warps * self.region_slots * 16) + top_smem = top0 + team * (TOP_K * 8) + xchg_smem = xchg0 + team * 64 + + per_cta = (m + self.num_ctas - 1) // self.num_ctas + t_begin = rank * per_cta + t_end = t_begin + per_cta + if t_end > m: + t_end = m + + for j in cutlass.range(tidx, local_num, self.threads, unroll=1): + _sts_b32(mask_smem + j * 4, Int32(0)) + for e4 in cutlass.range(tidx, NUM_EXPERTS // 4, self.threads, unroll=1): + v = _ld_nc_v4_f32(bias_addr + Int64(e4 * 16)) + _sts_v4_f32(bias_smem + e4 * 16, v[0], v[1], v[2], v[3]) + cute.arch.barrier() + + stride = self.team_lanes + gl = team_warp * 32 + lane + for token in cutlass.range(t_begin + team, t_end, self.teams, unroll=1): + self.topk_token( + token, + token - t_begin, + team, + team_warp, + lane, + Int32(M_MAX), + logits_addr, + bias_smem, + ids_addr, + vals_addr, + scale_bits, + Int64(0), + cand_smem, + top_smem, + xchg_smem, + ids_smem, + mask_smem, + local_offset, + local_num, + ) + # Every lane of the team is done with the token's logits: re-arm them. + self.team_sync(team) + row = logits_addr + Int64(token) * Int64(NUM_EXPERTS * 4) + for s in cutlass.range_constexpr(self.slots): + _st_sentinel_b32(row + Int64((s * stride + gl) * 4)) + if cutlass.const_expr(self.publish): + # the team's first warp stored the ids / weights (lanes < 16): publish the token + if team_warp == 0: + _syncwarp() + if lane == 0: + _red_add_release_gpu(num_tiles_addr, 1) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + if cutlass.const_expr(not self.publish): + cute.arch.barrier() + if cutlass.const_expr(self.num_ctas > 1): + for j in cutlass.range(tidx, local_num, self.threads, unroll=1): + word = _lds_b32(mask_smem + j * 4) + dst = masks_smem + (rank * self.local_cap + j) * 4 + for c in cutlass.range_constexpr(self.num_ctas): + _st_dsmem_b32(dst, Int32(c), word) + _cluster_arrive_release() + _cluster_wait_acquire() + self.sort( + tidx, + warp, + lane, + rank, + t_begin, + t_end, + mask_smem, + masks_smem, + base_smem, + ids_smem, + warp_tot_smem, + e2p_addr, + p2e_addr, + tile_expert_addr, + tile_limit_addr, + total_addr, + num_tiles_addr, + local_offset, + local_num, + ) + + @cute.jit + def gather_cta( + self, + gidx: Int32, + tidx: Int32, + m: Int32, + half: Int64, + codes_addr: Int64, + sf_addr: Int64, + gsf_addr: Int64, + ): + """Items of 32 latent columns of one token: poll the owner rank's slot, quantize two + 16-element NVFP4 groups (``fp4_quantize``), store, re-arm.""" + nthreads = self.gather_clusters * self.num_ctas * self.threads + gtid = gidx * self.threads + tidx + gsf = _ldg_f32(gsf_addr) + rgsf = _rcp_approx_ftz(gsf) + rcp6 = _rcp_approx_ftz(Float32(6.0)) + items_per_row = LATENT // _ITEM_COLS + big = self.tiles_base + 1 + split = self.tiles_rem * big + for p in cutlass.range(gtid, m * items_per_row, nthreads, unroll=1): + t = p // items_per_row + j = p - t * items_per_row + tile = j >> 2 + owner = Int32(0) + t0 = Int32(0) + if tile < split: + owner = tile // big + t0 = owner * big + else: + owner = self.tiles_rem + (tile - split) // self.tiles_base + t0 = split + (owner - self.tiles_rem) * self.tiles_base + lcol = j * _ITEM_COLS - t0 * TILE + src = half + Int64(((t * self.world + owner) * self.slot_n + lcol) * 2) + words = _poll_raw4(src, 16, 4) + vals = [] + for i in cutlass.range_constexpr(16): + lo_f, hi_f = _bf16x2_to_f32x2(words[i]) + vals.extend((lo_f, hi_f)) + amax0 = _abs_max_f32x16(vals[:16]) + amax1 = _abs_max_f32x16(vals[16:]) + sf2, os0, os1 = _nvfp4_scales_x2(amax0, amax1, gsf, rgsf, rcp6) + _stg_v4_b32( + codes_addr + Int64(t) * (LATENT // 2) + Int64(j * (_ITEM_COLS // 2)), + _e2m1x8(vals[0:8], os0), + _e2m1x8(vals[8:16], os0), + _e2m1x8(vals[16:24], os1), + _e2m1x8(vals[24:32], os1), + ) + _stg_u16( + sf_addr + Int64(t) * (LATENT // _FP4_VEC) + Int64(j * (_ITEM_COLS // _FP4_VEC)), + sf2, + ) + for v in cutlass.range_constexpr(4): + lamport.st_sentinel_v4(src + Int64(v * 16)) + if cutlass.const_expr(self.use_pdl): + _griddepcontrol_launch_dependents() + + +@dsl_user_op +def _i2f(word, *, loc=None, ip=None): + return Float32( + _asm( + T.f32(), + [Int32(word).ir_value(loc=loc, ip=ip)], + "mov.b32 $0, $1;", + "=f,r", + loc=loc, + ip=ip, + ) + ) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class UllFrontWorkspace: + """The front's MNNVL multicast workspace of one TP group: two halves of + :func:`half_bytes`; ``flags[h]`` is the constant rotation state (``k3_lamport_ar`` layout) that + points the kernels at half ``h``.""" + + world: int + rank: int + uc_local: int + mc_base: int + peer_ptrs: torch.Tensor + flags: Tuple[torch.Tensor, torch.Tensor] + half_bytes: int + handle: object = None + calls: int = 0 + + def next_flags(self) -> torch.Tensor: + """The flags of the next call (the halves alternate call by call).""" + flags = self.flags[self.calls & 1] + self.calls += 1 + return flags + + +_WORKSPACES: Dict[object, UllFrontWorkspace] = {} +_lock = threading.Lock() + + +def ull_front_workspace(mapping, comm=None) -> UllFrontWorkspace: + """The front workspace of ``mapping``'s TP group, created on first use (collective: every rank + of the group calls it at the same point). ``comm``: the group's MPI communicator (default: the + one of the MNNVL all-reduce workspace).""" + ws = _WORKSPACES.get(mapping) + if ws is not None: + return ws + from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace + from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer + + from ..layers import mnnvl_workspace_mpi_comm + + if comm is None: + comm = mnnvl_workspace_mpi_comm( + get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16) + ) + world, rank = mapping.tp_size, mapping.tp_rank + if not supports_world(world): + raise ValueError(f"k3_ull_front does not serve TP world {world}") + nbytes = half_bytes(world) + use_fabric = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node() + handle = McastGPUBuffer(2 * nbytes, world, rank, mapping.local_rank, use_fabric, comm.py2f()) + buf = handle.get_uc_buffer(rank, (2 * nbytes // 4,), torch.int32, 0) + buf.fill_(SENTINEL) + torch.cuda.synchronize() + comm.Barrier() + peers = [handle.get_uc_buffer(r, (1,), torch.uint8, 0).data_ptr() for r in range(world)] + flags = tuple( + torch.tensor([h, 0, nbytes, 0, 0, 0, 0, 0, 0], dtype=torch.int32, device=buf.device) + for h in (0, 1) + ) + ws = UllFrontWorkspace( + world=world, + rank=rank, + uc_local=buf.data_ptr(), + mc_base=handle.get_mc_buffer((1,), torch.uint8, 0).data_ptr(), + peer_ptrs=torch.tensor(peers, dtype=torch.int64, device=buf.device), + flags=flags, + half_bytes=nbytes, + handle=handle, + ) + _WORKSPACES[mapping] = ws + return ws + + +@dataclasses.dataclass +class UllFrontPlan: + """One MoE layer's front operands on one rank (views of the model's weights plus the packed + scale of the rank's latent_down tiles).""" + + world: int + rank: int + latent_w: torch.Tensor # FP8 [128 * tiles of the rank, 7168] + latent_ws: torch.Tensor # DeepGEMM-packed UE8M0 scale of those rows (MN-major int32) + router_w: torch.Tensor # BF16 [896, 7168] + bias: torch.Tensor # FP32 [896] + routed_scaling_factor: float + global_sf: torch.Tensor # FP32 [1]: the NVFP4 global scale of the routed experts' input + local_offset: int # this rank's routed experts (moe_sort tables) + local_num: int + + +def make_plan( + latent_w: torch.Tensor, + latent_scale: torch.Tensor, + router_w: torch.Tensor, + bias: torch.Tensor, + routed_scaling_factor: float, + global_sf: torch.Tensor, + world: int, + rank: int, + local_offset: int, + local_num: int, +) -> UllFrontPlan: + """The front plan of one layer: ``latent_w`` FP8 ``[3584, 7168]`` with its 128x128 block scale + in either layout of ``_Fp8BlockScaleWeightReadLinear`` (Rubin MXFP8 R128c4 or DeepGEMM-packed), + ``router_w`` BF16 ``[896, 7168]``.""" + if not supports_world(world): + raise ValueError(f"k3_ull_front does not serve TP world {world}") + if tuple(latent_w.shape) != (LATENT, HIDDEN) or latent_w.dtype != torch.float8_e4m3fn: + raise ValueError("latent_w must be FP8 [3584, 7168]") + if tuple(router_w.shape) != (NUM_EXPERTS, HIDDEN) or router_w.dtype != torch.bfloat16: + raise ValueError("router_w must be BF16 [896, 7168]") + t0, nt = latent_tiles_of(rank, world) + return UllFrontPlan( + world=world, + rank=rank, + latent_w=latent_w[t0 * TILE : (t0 + nt) * TILE], + latent_ws=lowm.pack_ue8m0_exponents( + lowm.block_exponents_of(latent_scale, LATENT, HIDDEN)[t0 : t0 + nt].contiguous(), + nt * TILE, + HIDDEN, + ), + router_w=router_w, + bias=bias.contiguous(), + routed_scaling_factor=float(routed_scaling_factor), + global_sf=global_sf.reshape(-1)[:1], + local_offset=local_offset, + local_num=local_num, + ) + + +@functools.cache +def latent_tactic(m: int) -> "lowm.K3LowMTactic": + """The latent GEMM's tactic: the low-M kernel's split-K tactic ``k3_latent_down_fp4`` uses at + ``M <= 16`` (``k3_latent_fp8_gemm.route_tactic``, 3 cluster ranks); above, the low-M kernel + with the mid-M route's 4 ranks (the mid-M kernel has no all-gather layout).""" + return lowm.split_tactic(m, LATENT, HIDDEN, 3 if m <= 16 else 4, mid=False) + + +def _m_tiles(mma_n: int) -> int: + return (mma_n + 15) // 16 + + +@functools.cache +def _compiled_head(n: int, tactic, world: int, use_pdl: bool): + gemm = lowm.K3Fp8wLowMGemmKernel( + n, + HIDDEN, + tactic, + use_pdl, + ar_push=lowm.AR_PUSH_ONE_SHOT, + ar_world=world, + ar_slot_n=slot_cols(world), + ) + head = K3UllFrontHeadKernel( + gemm, + ROUTER_TILES, + _m_tiles(tactic.mma_n), + logits_offset(world), + use_pdl, + ) + dummy = 1 << 20 + addr = Int64(dummy) + return cute.compile( + head, + make_ptr(cutlass.Float8E4M3FN, dummy, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, dummy, cute.AddressSpace.gmem, assumed_align=16), + Int32(1), + addr, + addr, + addr, + Int32(0), + make_ptr(cutlass.BFloat16, dummy, cute.AddressSpace.gmem, assumed_align=16), + addr, + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +#: Local-expert capacity of the routing tail's sort build (as ``k3_route_tail``): 224 (EP >= 4), +#: else every expert (e.g. TP8 x EP2: 448 local experts). +_SMALL_LOCAL_CAP = 224 + + +def _local_cap(local_num: int) -> int: + return _SMALL_LOCAL_CAP if local_num <= _SMALL_LOCAL_CAP else NUM_EXPERTS + + +#: Warps routing one token (measured best at M 8 / 16 / 32 on Rubin, hecate 834837). +_TEAM = 2 + + +def route_config(m: int) -> Tuple[int, int, int]: + """``(num_ctas, warps_per_token, num_warps)`` of the route cluster for ``m`` rows: eight CTAs, + one team per token of the CTA (``ceil(m / 8)`` tokens).""" + per_cta = (m + MAX_CTAS - 1) // MAX_CTAS + return MAX_CTAS, _TEAM, max(_TEAM * per_cta, 8) + + +@functools.cache +def _compiled_route( + num_ctas: int, + num_warps: int, + warps_per_token: int, + local_cap: int, + world: int, + use_pdl: bool, + publish: bool = False, +): + kernel = K3UllFrontRouteKernel( + num_ctas, + num_warps, + warps_per_token, + local_cap, + world, + use_pdl, + publish, + ) + addr = Int64(1 << 20) + return cute.compile( + kernel, + Int32(1), + *([addr] * 11), + Int32(0), + Int32(NUM_EXPERTS), + Int64(0), + addr, + addr, + addr, + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +#: Launch shapes compiled so far (see :func:`is_compiled`). +_READY: set = set() + + +def _shape_key(m: int, plan: UllFrontPlan, publish: bool = False) -> tuple: + return ( + publish, + plan.latent_w.shape[0], + latent_tactic(m), + plan.world, + _local_cap(plan.local_num), + route_config(m), + ) + + +def is_compiled(m: int, plan: UllFrontPlan, publish: bool = False) -> bool: + """Whether a call of ``m`` rows with this plan's shapes runs without compiling (required + inside CUDA-graph capture).""" + return _shape_key(m, plan, publish) in _READY + + +def _gptr(dtype, t: torch.Tensor): + return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def can_run(x: torch.Tensor, plan: UllFrontPlan) -> bool: + return ( + x.is_cuda + and torch.cuda.get_device_capability(x.device) in ((10, 0), (10, 3), (10, 7)) + and x.dtype == torch.bfloat16 + and x.dim() == 2 + and x.is_contiguous() + and x.shape[1] == HIDDEN + and 1 <= x.shape[0] <= M_MAX + and x.data_ptr() % 16 == 0 + and 1 <= plan.local_num + and plan.local_offset + plan.local_num <= NUM_EXPERTS + ) + + +def run_k3_ull_front( + x: torch.Tensor, + plan: UllFrontPlan, + ws: UllFrontWorkspace, + *, + after_head=None, + use_pdl: bool = True, + publish: int = 0, +) -> Tuple[Tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: + """Launch the head and the route kernel on the current stream. + + Returns ``(routing, codes, scales)``: ``routing`` = ``k3_route_tail``'s 8 outputs (top-k ids + and weights, then the ``moe_sort`` tables of ``plan.local_offset`` / ``local_num`` at tile 128), + ``codes`` uint8 ``[M, 1792]`` / ``scales`` uint8 ``[M, 224]`` the NVFP4 latent (linear scale + layout), i.e. ``k3_latent_down_fp4``'s outputs. + ``after_head``: called between the two launches (e.g. to fork a side stream after the head). + ``publish``: address of an int32 word (zero) the route kernel adds 1 to per routed token + (release) instead of building the ``moe_sort`` tables (``routing[2:]`` are then not written): + the routed experts' ``published_routing_word``, from which they build the tables themselves. + The first call of each shape compiles (outside CUDA-graph capture).""" + m = x.shape[0] + if not can_run(x, plan): + raise ValueError(f"k3_ull_front: unsupported call x {tuple(x.shape)} {x.dtype}") + world = plan.world + tactic = latent_tactic(m) + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + flags = ws.next_flags() + head = _compiled_head(plan.latent_w.shape[0], tactic, world, use_pdl) + head( + _gptr(cutlass.Float8E4M3FN, plan.latent_w), + _gptr(cutlass.BFloat16, x), + _gptr(cutlass.Uint8, plan.latent_ws), + Int32(m), + Int64(ws.peer_ptrs.data_ptr()), + Int64(ws.mc_base), + Int64(flags.data_ptr()), + Int32(ws.rank), + _gptr(cutlass.BFloat16, plan.router_w), + Int64(x.data_ptr()), + stream, + ) + if after_head is not None: + after_head() + device = x.device + max_tiles, max_permuted = sort_buffer_sizes(m, plan.local_num, TILE) + i32 = dict(dtype=torch.int32, device=device) + routing = ( + torch.empty((m, TOP_K), **i32), + torch.empty((m, TOP_K), dtype=torch.float32, device=device), + torch.empty((max_tiles,), **i32), + torch.empty((max_tiles,), **i32), + torch.empty((m, TOP_K), **i32), + torch.empty((max_permuted,), **i32), + torch.empty((1,), **i32), + torch.empty((1,), **i32), + ) + codes = torch.empty((m, LATENT // 2), dtype=torch.uint8, device=device) + scales = torch.empty((m, LATENT // _FP4_VEC), dtype=torch.uint8, device=device) + tile_expert, tile_limit, e2p, p2e, total, num_tiles = (t.data_ptr() for t in routing[2:]) + if publish: + num_tiles = publish + scale_bits = struct.unpack(" None: + """Validate that a tensor is CUDA bf16.""" + if not tensor.is_cuda or tensor.dtype is not torch.bfloat16: + raise TypeError(f"{name} must be a CUDA bfloat16 tensor") + + +def _require_cuda_fp32(name: str, tensor: torch.Tensor) -> None: + """Validate that a tensor is CUDA fp32.""" + if not tensor.is_cuda or tensor.dtype is not torch.float32: + raise TypeError(f"{name} must be a CUDA float32 tensor") + + +def _as_token_rows(tensor: torch.Tensor) -> torch.Tensor: + """Keep a decode input view when its head and channel axes are packed.""" + if tensor.ndim == 4 and tensor.stride(3) == 1 and tensor.stride(2) == tensor.shape[3]: + return tensor + return tensor.contiguous() + + +def _dummy_tensor( + tag: str, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + fill: float = 0.0, +) -> torch.Tensor: + """Return a cached dummy tensor for optional CUDA arguments.""" + device = torch.device(device) + key = (tag, shape, dtype, device.type, device.index, fill) + tensor = _DUMMY_CACHE.get(key) + if tensor is None: + if fill == 0.0: + tensor = torch.zeros(shape, dtype=dtype, device=device) + elif fill == 1.0: + tensor = torch.ones(shape, dtype=dtype, device=device) + else: + tensor = torch.full(shape, fill, dtype=dtype, device=device) + _DUMMY_CACHE[key] = tensor + return tensor + + +def run_kda_decode_fusion_cuda( + *, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q_t: torch.Tensor, + w_k_t: torch.Tensor, + w_v_t: torch.Tensor, + bias_q: torch.Tensor | None, + bias_k: torch.Tensor | None, + bias_v: torch.Tensor | None, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + A_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor | None, + beta: torch.Tensor, + state: torch.Tensor, + onorm_g: torch.Tensor | None = None, + onorm_weight: torch.Tensor | None = None, + out: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + scale: float = 128**-0.5, + onorm_eps: float = 1e-5, + lower_bound: float | None = None, + use_beta_sigmoid_in_kernel: bool = True, + verbose: bool = False, + update_conv_cache: bool = False, +) -> torch.Tensor: + """Run CUDA KDA decode fusion for the tuned decode shapes. + + ``ssm_state_indices=None`` selects the tuned batch-local static layout, + while a tensor selects the indexed state-pool layout. + """ + for name, tensor in ( + ("x_q", x_q), + ("x_k", x_k), + ("x_v", x_v), + ("w_q_t", w_q_t), + ("w_k_t", w_k_t), + ("w_v_t", w_v_t), + ("cs_q", cs_q), + ("cs_k", cs_k), + ("cs_v", cs_v), + ("g", g), + ("beta", beta), + ): + _require_cuda_bf16(name, tensor) + for name, tensor in (("A_log", A_log), ("state", state)): + _require_cuda_fp32(name, tensor) + + if x_q.ndim != 4 or x_k.ndim != 4 or x_v.ndim != 4: + raise ValueError("x_q, x_k, and x_v must be rank-4 decode tensors") + if x_q.shape[0] != 1 or x_k.shape[0] != 1 or x_v.shape[0] != 1: + raise ValueError("only T=1 decode inputs are supported") + if x_q.shape[-1] != 128 or x_k.shape[-1] != 128 or x_v.shape[-1] != 128: + raise ValueError("only K=128 and V=128 are supported") + + B = x_q.shape[1] + H = x_q.shape[2] + HV = x_v.shape[2] + if x_k.shape[1:3] != (B, H) or x_v.shape[1] != B: + raise ValueError("x_q, x_k, and x_v batch/head dimensions are inconsistent") + if H != HV or H not in (1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96): + raise ValueError( + "CUDA KDA decode fusion supports H == HV in {1,2,3,4,6,8,12,16,24,32,48,64,96}" + ) + if ssm_state_indices is None and not state.is_contiguous(): + raise ValueError("state must be contiguous because it is updated in place") + if out is None: + # The op is inplace-only and never allocates, so supply a buffer here. + # Hot decode paths pass a persistent one instead. + out = x_q.new_empty((B, 1, HV, 128)) + else: + _require_cuda_bf16("out", out) + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + if tuple(out.shape) != (B, 1, HV, 128): + raise ValueError("out must have shape [B, 1, HV, 128]") + + device = x_q.device + apply_onorm = onorm_g is not None + if bias_q is None: + bias_q = _dummy_tensor("bias_q", (H * 128,), torch.bfloat16, device) + if bias_k is None: + bias_k = _dummy_tensor("bias_k", (H * 128,), torch.bfloat16, device) + if bias_v is None: + bias_v = _dummy_tensor("bias_v", (HV * 128,), torch.bfloat16, device) + if dt_bias is None: + dt_bias = _dummy_tensor("dt_bias", (H * 128,), torch.float32, device) + if onorm_g is None: + onorm_g = _dummy_tensor("onorm_g", (1, B, HV, 128), torch.bfloat16, device) + if onorm_weight is None: + onorm_weight = _dummy_tensor("onorm_weight", (128,), torch.float32, device, fill=1.0) + + for name, tensor in ( + ("bias_q", bias_q), + ("bias_k", bias_k), + ("bias_v", bias_v), + ("onorm_g", onorm_g), + ): + _require_cuda_bf16(name, tensor) + for name, tensor in (("dt_bias", dt_bias), ("onorm_weight", onorm_weight)): + _require_cuda_fp32(name, tensor) + + if update_conv_cache: + if any(tensor.ndim != 3 for tensor in (cs_q, cs_k, cs_v)): + raise ValueError("update_conv_cache expects rank-3 conv-state pools") + projection_size = H * 128 + conv_state_width = cs_q.shape[2] + min_slot_stride = 3 * projection_size * conv_state_width + if not ( + cs_q.stride(0) == cs_k.stride(0) == cs_v.stride(0) + and cs_q.stride(0) >= min_slot_stride + and cs_q.stride(1) == cs_k.stride(1) == cs_v.stride(1) == conv_state_width + and cs_q.stride(2) == cs_k.stride(2) == cs_v.stride(2) == 1 + ): + raise ValueError( + "update_conv_cache expects section views of packed " + "[slots, 3 * dim, width] conv states" + ) + + args = ( + _as_token_rows(x_q), + _as_token_rows(x_k), + _as_token_rows(x_v), + w_q_t.contiguous(), + w_k_t.contiguous(), + w_v_t.contiguous(), + bias_q.contiguous(), + bias_k.contiguous(), + bias_v.contiguous(), + cs_q if update_conv_cache else cs_q.contiguous(), + cs_k if update_conv_cache else cs_k.contiguous(), + cs_v if update_conv_cache else cs_v.contiguous(), + A_log.contiguous(), + _as_token_rows(g), + dt_bias.contiguous(), + beta if beta.ndim == 3 and beta.stride(2) == 1 else beta.contiguous(), + _as_token_rows(onorm_g), + onorm_weight.contiguous(), + ssm_state_indices, + state, + ) + + use_lower_bound = lower_bound is not None + lower_bound_value = 0.0 if lower_bound is None else float(lower_bound) + launch_args = ( + bool(apply_onorm), + bool(update_conv_cache), + bool(use_lower_bound), + bool(use_beta_sigmoid_in_kernel), + lower_bound_value, + float(scale), + float(onorm_eps), + ) + kda_decode(*args, *launch_args, output=out) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/_kda_kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/_kda_kernels.py new file mode 100644 index 000000000000..8cbe2199a965 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/_kda_kernels.py @@ -0,0 +1,1805 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA kernels of the Kimi K3 (k3_rubin) package and their dispatch. + +* prefill: the CuTe DSL pipeline (``KdaPrefillRunner``, ``kda_prefill`` in ``ops/ssm/_kda_prefill.py``, + called through the ``ops/ssm/kda_prefill`` entry) on L2-normalized Q/K, with its mid-chunk + snapshot outputs; +* multi-token verify: the v1 / v2 / two-CTA kernels (``kda_mtp_decode_impl`` in ``mtp_ops.py``, + called through the ``ops/ssm/kda_mtp_decode`` entry) on the replay caches main's cache manager + allocates; +* single-token decode: the fused CUDA decode kernel (``_kda_decode.py``); +* the Triton prefill convolution (``fused_kda_conv_prefill``, features.kda_fused_conv_prefill) + and post-convolution layout (``fused_kda_post_conv``, with the chunk padding of + features.kda_prepad_prefill). + +The replay-cache helpers below (``copy`` / ``store`` / ``replay`` / ``commit`` of the verify +kernel's replay caches and the batched seed plan) are kernels of the cache manager's buffers; the +engine's cache manager (main's ``KimiK3HybridCacheManager``) calls main's copies of them. +""" + +from __future__ import annotations + +import dataclasses +import functools +from types import ModuleType +from typing import Optional, Sequence, Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from tensorrt_llm._utils import prefer_pinned + +from . import _kda_decode + +try: + from tensorrt_llm._utils import get_sm_version as _tllm_get_sm_version +except ImportError: # pragma: no cover — source-loader stub path + _tllm_get_sm_version = None + + +# Slot index of a padding sequence: the convolution reads no conv state for it and leaves the pool +# untouched. +PAD_SLOT_ID = -1 + + +def _default_get_sm_version() -> int: + if not torch.cuda.is_available() or torch.cuda.device_count() == 0: + return -1 + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor + + +def get_kda_sm_version() -> int: + """Return the runtime SM version used for KDA kernel selection. + + Prefers ``tensorrt_llm._utils.get_sm_version`` when the real package is + importable so environment-side overrides propagate. Falls back to a + plain CUDA-property probe when we are executing under the source-loader + stub subtree. + """ + if _tllm_get_sm_version is not None: + try: + return int(_tllm_get_sm_version()) + except RuntimeError: + # torch raises RuntimeError when no CUDA device is usable; + # the property probe below handles that case itself. + return _default_get_sm_version() + return _default_get_sm_version() + + +def is_kda_optimized_supported() -> bool: + """The optimized prefill/decode kernels need Blackwell sm_100 or Rubin sm_107.""" + return get_kda_sm_version() in (100, 103, 107) + + +@triton.jit(do_not_specialize=["num_tokens", "num_out_tokens", "beta_row_stride"]) +def _fused_kda_post_conv_kernel( + packed_ptr, + q_out_ptr, + k_out_ptr, + v_out_ptr, + beta_ptr, + beta_out_ptr, + num_tokens, + num_out_tokens, + beta_row_stride, + l2_norm_eps, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, + BLOCK_DIM: tl.constexpr, + HAS_BETA: tl.constexpr, +): + """Normalize and transpose channel-major packed Q/K/V in one launch. + + Rows in ``[num_tokens, num_out_tokens)`` are written as exact zeros, so + the outputs are already zero-padded for the KDA prefill op. When + ``HAS_BETA``, beta is cast to fp32 and zero-padded the same way. + """ + token_offsets = tl.program_id(0) * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS) + head_idx = tl.program_id(1) + dim_offsets = tl.arange(0, BLOCK_DIM) + token_mask = token_offsets < num_tokens + out_token_mask = token_offsets < num_out_tokens + dim_mask = dim_offsets < HEAD_DIM + mask = token_mask[:, None] & dim_mask[None, :] + out_mask = out_token_mask[:, None] & dim_mask[None, :] + + feature_offsets = head_idx * HEAD_DIM + dim_offsets + projection_size = NUM_HEADS * HEAD_DIM + num_tokens_i64 = num_tokens.to(tl.int64) + source_offsets = feature_offsets[None, :].to(tl.int64) * num_tokens_i64 + token_offsets[ + :, None + ].to(tl.int64) + section_stride = projection_size * num_tokens_i64 + output_offsets = ( + token_offsets[:, None].to(tl.int64) * projection_size + feature_offsets[None, :] + ) + + # Masked-off rows load 0 and normalize to exactly +0 (0 / sqrt(eps)). + q = tl.load(packed_ptr + source_offsets, mask=mask, other=0.0).to(tl.float32) + q /= tl.sqrt(tl.sum(q * q, axis=1) + l2_norm_eps)[:, None] + tl.store(q_out_ptr + output_offsets, q, mask=out_mask) + + k = tl.load( + packed_ptr + section_stride + source_offsets, + mask=mask, + other=0.0, + ).to(tl.float32) + k /= tl.sqrt(tl.sum(k * k, axis=1) + l2_norm_eps)[:, None] + tl.store(k_out_ptr + output_offsets, k, mask=out_mask) + + v = tl.load( + packed_ptr + 2 * section_stride + source_offsets, + mask=mask, + other=0.0, + ) + tl.store(v_out_ptr + output_offsets, v, mask=out_mask) + + if HAS_BETA: + token_i64 = token_offsets.to(tl.int64) + beta = tl.load( + beta_ptr + token_i64 * beta_row_stride + head_idx, mask=token_mask, other=0.0 + ).to(tl.float32) + tl.store(beta_out_ptr + token_i64 * NUM_HEADS + head_idx, beta, mask=out_token_mask) + + +def fused_kda_post_conv( + packed: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float = 1e-6, + padded_tokens: Optional[int] = None, + beta: Optional[torch.Tensor] = None, + *, + normalize_qk: bool = True, + out: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, +): + """Convert packed channel-major ``[3 * H * D, T]`` convolution output to KDA Q/K/V. + + A contiguous ``packed`` with ``normalize_qk`` runs the branch's Triton kernel (the CuTe DSL + prefill path): contiguous ``[1, out_tokens, H, D]`` Q/K/V with Q and K L2-normalized, where + ``out_tokens`` is ``padded_tokens`` when given (rows past ``T`` are exact zeros, the chunk + padding of features.kda_prepad_prefill), else ``T``; with ``beta`` (``[..., T, H]``, unit head + stride) a fourth output holds it cast to fp32 and zero-padded to ``[1, out_tokens, H]``. + + Otherwise main's portable layout conversion runs (Q/K normalized only with + ``normalize_qk``, into ``out`` when given); it takes no padding. + """ + if normalize_qk and out is None and packed.is_contiguous(): + return _fused_kda_post_conv_triton( + packed, num_heads, head_dim, l2_norm_eps, padded_tokens, beta + ) + if padded_tokens is not None or beta is not None: + raise ValueError( + "Padded KDA post-conv outputs need a contiguous packed input, Q/K normalization " + "and no caller-provided outputs" + ) + if packed.ndim != 2 or packed.shape[0] != 3 * num_heads * head_dim: + raise ValueError("Expected packed [3 * H * D, T] convolution output") + if packed.shape[1] > 1 and 1 not in packed.stride(): + raise ValueError("Packed KDA post-conv requires a unit stride along tokens or features") + qkv = tuple( + t.t().reshape(1, packed.shape[1], num_heads, head_dim).contiguous() + for t in packed.chunk(3, dim=0) + ) + if normalize_qk: + q, k, v = qkv + + def normalize(t): + f = t.float() + return (f * torch.rsqrt((f * f).sum(-1, keepdim=True) + l2_norm_eps)).to(t.dtype) + + qkv = (normalize(q), normalize(k), v) + if out is not None: + if len(out) != 3 or any( + dst.shape != src.shape + or dst.dtype != src.dtype + or dst.device != src.device + or not dst.is_contiguous() + for dst, src in zip(out, qkv) + ): + raise ValueError("Invalid KDA layout output buffers") + for dst, src in zip(out, qkv): + dst.copy_(src) + return out + return qkv + + +def _beta_rows_and_output( + beta: Optional[torch.Tensor], + num_tokens: int, + out_tokens: int, + num_heads: int, + device: torch.device, +) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Validate an optional ``beta`` and allocate its fp32 ``[1, out_tokens, heads]`` output. + + Returns ``(beta_rows, beta_out)``: the ``[tokens, heads]`` view of ``beta`` + the kernels read and the output they write, or ``(None, None)``. + """ + if beta is None: + return None, None + beta_2d = None + if beta.ndim == 2 and beta.shape == (num_tokens, num_heads): + beta_2d = beta + elif beta.ndim == 3 and beta.shape == (1, num_tokens, num_heads): + beta_2d = beta[0] + if beta_2d is None or beta_2d.stride(-1) != 1 or beta_2d.device != device: + raise ValueError( + f"beta must be [tokens={num_tokens}, heads={num_heads}] or " + f"[1, tokens={num_tokens}, heads={num_heads}] on {device} " + f"with unit head stride, got shape {tuple(beta.shape)}, " + f"stride {beta.stride()}, device {beta.device}" + ) + beta_out = torch.empty((1, out_tokens, num_heads), dtype=torch.float32, device=device) + return beta_2d, beta_out + + +def _fused_kda_post_conv_triton( + packed: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float, + padded_tokens: Optional[int], + beta: Optional[torch.Tensor], +): + """The branch's ``fused_kda_post_conv`` (see :func:`fused_kda_post_conv`).""" + projection_size = num_heads * head_dim + if packed.ndim != 2 or packed.shape[0] != 3 * projection_size: + raise ValueError( + "Packed KDA post-conv expected shape " + f"[{3 * projection_size}, tokens], got {tuple(packed.shape)}" + ) + if not packed.is_contiguous(): + raise ValueError("Packed KDA post-conv requires a contiguous tensor") + + num_tokens = packed.shape[1] + out_tokens = num_tokens if padded_tokens is None else padded_tokens + if out_tokens < num_tokens: + raise ValueError(f"padded_tokens={out_tokens} is smaller than tokens={num_tokens}") + output_shape = (1, out_tokens, num_heads, head_dim) + q_out = torch.empty(output_shape, dtype=packed.dtype, device=packed.device) + k_out = torch.empty_like(q_out) + v_out = torch.empty_like(q_out) + beta_2d, beta_out = _beta_rows_and_output( + beta, num_tokens, out_tokens, num_heads, packed.device + ) + if out_tokens == 0: + return (q_out, k_out, v_out) if beta is None else (q_out, k_out, v_out, beta_out) + + block_tokens = 16 + block_dim = triton.next_power_of_2(head_dim) + grid = (triton.cdiv(out_tokens, block_tokens), num_heads) + _fused_kda_post_conv_kernel[grid]( + packed, + q_out, + k_out, + v_out, + beta_2d if beta_2d is not None else q_out, + beta_out if beta_out is not None else q_out, + num_tokens, + out_tokens, + beta_2d.stride(0) if beta_2d is not None else 0, + l2_norm_eps, + num_heads, + head_dim, + block_tokens, + block_dim, + beta is not None, + num_warps=8, + num_stages=3, + ) + return (q_out, k_out, v_out) if beta is None else (q_out, k_out, v_out, beta_out) + + +# Tile of the fused prefill convolution: tokens per program and warps per program. +_FUSED_CONV_BLOCK_TOKENS = 32 +_FUSED_CONV_NUM_WARPS = 4 + + +@triton.jit +def _kda_conv_silu(acc, out_dtype): + """``causal_conv1d_fwd``'s SiLU, ``y * __frcp_rn(1 + __expf(-y))``, rounded to ``out_dtype``. + + ``__expf`` is spelled as the ``ex2.approx.ftz`` it compiles to. The + reciprocal is ``rcp.rn``'s in-range sequence (MUFU.RCP and one FMA + correction) without its out-of-range branch: ``1 + exp(-y)`` is at least + 1, and from 2**126 on (``y`` below about -87.3) the reciprocal is taken as + 0 instead of a subnormal. Results are bitwise ``causal_conv1d_fwd``'s + except for those ``y``, where SiLU gives -0 instead of a value of + magnitude below 1e-36. + """ + exp_neg = tl.inline_asm_elementwise( + "ex2.approx.ftz.f32 $0, $1;", + "=r,r", + [acc * -1.4426950408889634], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + reciprocal = tl.inline_asm_elementwise( + """{ + .reg .f32 r0, t; + .reg .pred large; + rcp.approx.ftz.f32 r0, $1; + fma.rn.f32 t, $1, r0, 0fBF800000; + neg.f32 t, t; + fma.rn.f32 $0, r0, t, r0; + setp.ge.f32 large, $1, 0f7E800000; + selp.f32 $0, 0f00000000, $0, large; + }""", + "=r,r", + [1.0 + exp_neg], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + return _kda_round(acc * reciprocal, out_dtype) + + +@triton.jit +def _kda_round(x, out_dtype: tl.constexpr): + """Round fp32 ``x`` to nearest-even ``out_dtype``; bf16 two elements per instruction.""" + if out_dtype == tl.bfloat16: + return tl.inline_asm_elementwise( + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,r,r", + [x], + dtype=tl.bfloat16, + is_pure=True, + pack=2, + ) + return x.to(out_dtype) + + +@triton.jit +def _kda_widen(x, MIXED_BF16_FMA: tl.constexpr): + """Exact fp32 value of ``x``; bf16 with one ``add.rn.f32.bf16`` of -0 (SM100+).""" + if MIXED_BF16_FMA and x.dtype == tl.bfloat16: + return tl.inline_asm_elementwise( + "add.rn.f32.bf16 $0, $1, 0f80000000;", + "=r,h", + [x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + return x.to(tl.float32) + + +@triton.jit +def _kda_conv_tap(acc, weight, x, MIXED_BF16_FMA: tl.constexpr): + """``acc + weight * x`` with an fp32 accumulator, one rounding per tap. + + With ``MIXED_BF16_FMA`` (SM100+), bf16 operands use the mixed-precision + ``fma.rn.f32.bf16`` (no conversion instructions); the bf16 product is + exact in fp32 either way. + """ + if MIXED_BF16_FMA and weight.dtype == tl.bfloat16 and x.dtype == tl.bfloat16: + return tl.inline_asm_elementwise( + "fma.rn.f32.bf16 $0, $1, $2, $3;", + "=r,h,h,r", + [weight, x, acc], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + return acc + weight.to(tl.float32) * x.to(tl.float32) + + +@triton.jit +def _kda_store_conv_rows( + out, + offsets, + mask, + section, + q_out_ptr, + k_out_ptr, + v_out_ptr, + l2_norm_eps, + MIXED_BF16_FMA: tl.constexpr, + NORMALIZE_QK: tl.constexpr, +): + """Store ``[rows, head_dim]`` conv outputs: V as is, Q and K L2-normalized per row + with ``NORMALIZE_QK``, else as is. + + The norm is fp32 over the rounded values (the branch's Triton post-conv). Its + division compiles to a multiplication by ``rcp.approx`` of the norm for + norms in ``[2**-126, 2**126]``; the reciprocal is taken once per row here. + """ + if section == 2: + tl.store(v_out_ptr + offsets, out, mask=mask) + elif not NORMALIZE_QK: + if section == 0: + tl.store(q_out_ptr + offsets, out, mask=mask) + else: + tl.store(k_out_ptr + offsets, out, mask=mask) + else: + normalized = _kda_widen(out, MIXED_BF16_FMA) + norm = tl.sqrt(tl.sum(normalized * normalized, axis=1) + l2_norm_eps) + inv_norm = tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", + "=r,r", + [norm], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + normalized = _kda_round(normalized * inv_norm[:, None], out.dtype) + if section == 0: + tl.store(q_out_ptr + offsets, normalized, mask=mask) + else: + tl.store(k_out_ptr + offsets, normalized, mask=mask) + + +@triton.jit(do_not_specialize=["num_tokens", "num_out_tokens", "beta_row_stride"]) +def _fused_kda_conv_prefill_kernel( + src_ptr, + weight_ptr, + q_out_ptr, + k_out_ptr, + v_out_ptr, + beta_ptr, + beta_out_ptr, + num_tokens, + num_out_tokens, + src_stride_token, + src_stride_channel, + weight_stride_channel, + weight_stride_tap, + beta_row_stride, + l2_norm_eps, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, + BLOCK_DIM: tl.constexpr, + CHANNEL_CONTIGUOUS: tl.constexpr, + WEIGHT_CHANNEL_CONTIGUOUS: tl.constexpr, + HAS_BETA: tl.constexpr, + MIXED_BF16_FMA: tl.constexpr, + NORMALIZE_QK: tl.constexpr, +): + """Causal conv + SiLU (+ L2 norm for Q/K) of one head of one Q/K/V section. + + Program ``(tile, section * NUM_HEADS + head)`` produces tokens + ``[tile * BLOCK_TOKENS, (tile + 1) * BLOCK_TOKENS)`` of that head as if + all tokens formed one sequence without history. The first ``WIDTH - 1`` + tokens of every sequence are recomputed by + ``_kda_conv_prefill_edges_kernel``. Tokens in + ``[num_tokens, num_out_tokens)`` are zeros. + + The arithmetic follows ``causal_conv1d_fwd`` (fp32 taps accumulated + oldest first, SiLU, rounding to the activation dtype) and + ``fused_kda_post_conv``. The taps run as a loop: the rows of a tap load + are mostly cache hits on the previous tap's, and only one tap's tile is + live in registers. + """ + channel_block = tl.program_id(1) + section = channel_block // NUM_HEADS + head = channel_block % NUM_HEADS + first_token = tl.program_id(0) * BLOCK_TOKENS + rows = tl.arange(0, BLOCK_TOKENS) + tokens = first_token + rows + dims = tl.arange(0, BLOCK_DIM) + dim_mask = dims < HEAD_DIM + token_mask = tokens < num_tokens + channels = channel_block * HEAD_DIM + dims + + # 64-bit tile base, 32-bit offsets within the tile. + src_tile = ( + src_ptr + + first_token.to(tl.int64) * src_stride_token + + (channel_block * HEAD_DIM).to(tl.int64) * src_stride_channel + ) + if CHANNEL_CONTIGUOUS: + dim_offsets = dims + else: + dim_offsets = dims * src_stride_channel + if WEIGHT_CHANNEL_CONTIGUOUS: + weight_offsets = channels + else: + weight_offsets = channels * weight_stride_channel + # Weights are loaded per tile row so that they share the input tile's layout. + weight_offsets = weight_offsets[None, :] + rows[:, None] * 0 + acc = tl.zeros([BLOCK_TOKENS, BLOCK_DIM], dtype=tl.float32) + for tap in tl.range(WIDTH, loop_unroll_factor=1): + shift = WIDTH - 1 - tap + weight = tl.load( + weight_ptr + tap * weight_stride_tap + weight_offsets, + mask=dim_mask[None, :], + other=0.0, + ) + x = tl.load( + src_tile + ((rows - shift) * src_stride_token)[:, None] + dim_offsets[None, :], + mask=(token_mask & (tokens >= shift))[:, None] & dim_mask[None, :], + other=0.0, + ) + acc = _kda_conv_tap(acc, weight, x, MIXED_BF16_FMA) + + out_tile = first_token.to(tl.int64) * (NUM_HEADS * HEAD_DIM) + head * HEAD_DIM + _kda_store_conv_rows( + _kda_conv_silu(acc, v_out_ptr.dtype.element_ty), + out_tile + (rows * (NUM_HEADS * HEAD_DIM))[:, None] + dims[None, :], + (tokens < num_out_tokens)[:, None] & dim_mask[None, :], + section, + q_out_ptr, + k_out_ptr, + v_out_ptr, + l2_norm_eps, + MIXED_BF16_FMA, + NORMALIZE_QK, + ) + + if HAS_BETA: + if section == 0: + beta_tokens = tokens.to(tl.int64) + beta = tl.load( + beta_ptr + beta_tokens * beta_row_stride + head, mask=token_mask, other=0.0 + ).to(tl.float32) + tl.store( + beta_out_ptr + beta_tokens * NUM_HEADS + head, + beta, + mask=tokens < num_out_tokens, + ) + + +@triton.jit +def _kda_conv_prefill_edges_kernel( + src_ptr, + weight_ptr, + conv_state_ptr, + query_start_loc_ptr, + cache_indices_ptr, + has_initial_state_ptr, + q_out_ptr, + k_out_ptr, + v_out_ptr, + src_stride_token, + src_stride_channel, + weight_stride_channel, + weight_stride_tap, + state_stride_slot, + state_stride_channel, + state_stride_tap, + l2_norm_eps, + pad_slot_id, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK_EDGE: tl.constexpr, + BLOCK_DIM: tl.constexpr, + MIXED_BF16_FMA: tl.constexpr, + NORMALIZE_QK: tl.constexpr, +): + """Sequence edges of the fused conv for one sequence and one head of one section. + + Recomputes the sequence's first ``WIDTH - 1`` tokens with tap inputs + before its start taken from its conv-state slot (zeros without an + initial state or with a ``pad_slot_id`` slot), then stores its last + ``WIDTH - 1`` inputs in the slot. Sequences shorter than ``WIDTH - 1`` + keep the newest entries of their initial state (zeros without one) in + front of their inputs. A ``pad_slot_id`` slot is left untouched. + """ + seq = tl.program_id(0) + channel_block = tl.program_id(1) + section = channel_block // NUM_HEADS + head = channel_block % NUM_HEADS + dims = tl.arange(0, BLOCK_DIM) + dim_mask = dims < HEAD_DIM + channels = channel_block * HEAD_DIM + dims + rows = tl.arange(0, BLOCK_EDGE) + + seq_start = tl.load(query_start_loc_ptr + seq).to(tl.int64) + seq_len = tl.load(query_start_loc_ptr + seq + 1).to(tl.int64) - seq_start + slot = tl.load(cache_indices_ptr + seq).to(tl.int64) + use_state = (tl.load(has_initial_state_ptr + seq) != 0) & (slot != pad_slot_id) + src_channels = channels.to(tl.int64) * src_stride_channel + state = conv_state_ptr + slot * state_stride_slot + channels.to(tl.int64) * state_stride_channel + + # The first WIDTH - 1 tokens. + row_mask = rows < tl.minimum(seq_len, WIDTH - 1) + acc = tl.zeros([BLOCK_EDGE, BLOCK_DIM], dtype=tl.float32) + for tap in tl.static_range(WIDTH): + shift = WIDTH - 1 - tap + weight = tl.load( + weight_ptr + channels * weight_stride_channel + tap * weight_stride_tap, + mask=dim_mask, + other=0.0, + ) + from_input = rows >= shift + x = tl.load( + src_ptr + + ((seq_start + rows - shift) * src_stride_token)[:, None] + + src_channels[None, :], + mask=(row_mask & from_input)[:, None] & dim_mask[None, :], + other=0.0, + ) + if shift > 0: + history = tl.load( + state[None, :] + ((WIDTH - 1 - shift + rows) * state_stride_tap)[:, None], + mask=(row_mask & (rows < shift) & use_state)[:, None] & dim_mask[None, :], + other=0.0, + ) + x = tl.where(from_input[:, None], x, history) + acc = _kda_conv_tap(acc, weight[None, :], x, MIXED_BF16_FMA) + _kda_store_conv_rows( + _kda_conv_silu(acc, v_out_ptr.dtype.element_ty), + (seq_start + rows)[:, None] * (NUM_HEADS * HEAD_DIM) + (head * HEAD_DIM + dims)[None, :], + row_mask[:, None] & dim_mask[None, :], + section, + q_out_ptr, + k_out_ptr, + v_out_ptr, + l2_norm_eps, + MIXED_BF16_FMA, + NORMALIZE_QK, + ) + + if slot != pad_slot_id: + tap_mask = rows < WIDTH - 1 + position = seq_len - (WIDTH - 1) + rows + from_input = position >= 0 + x = tl.load( + src_ptr + ((seq_start + position) * src_stride_token)[:, None] + src_channels[None, :], + mask=(tap_mask & from_input)[:, None] & dim_mask[None, :], + other=0.0, + ) + kept = tl.load( + state[None, :] + ((rows + seq_len) * state_stride_tap)[:, None], + mask=(tap_mask & (position < 0) & use_state)[:, None] & dim_mask[None, :], + other=0.0, + ) + new_state = tl.where(from_input[:, None], x, kept) + # Every conv-state read of this program comes before the stores. + tl.debug_barrier() + tl.store( + state[None, :] + (rows * state_stride_tap)[:, None], + new_state, + mask=tap_mask[:, None] & dim_mask[None, :], + ) + + +def fused_kda_conv_prefill( + qkv: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float = 1e-6, + padded_tokens: Optional[int] = None, + beta: Optional[torch.Tensor] = None, + pad_slot_id: int = PAD_SLOT_ID, + normalize_qk: bool = True, +): + """Short convolution, SiLU and Q/K L2 norm of packed KDA Q/K/V in one pass. + + Equivalent to ``causal_conv1d_fn`` (SiLU, varlen, indexed conv states) + on the channel-major packed input followed by ``fused_kda_post_conv``, + without materializing the packed tensor. + + ``qkv`` is the ``[tokens, 3 * num_heads * head_dim]`` convolution input + (Q, K, V sections), read in place; a unit channel stride (rows of the + QKVG projection output) gives vectorized loads. Sequences are the + ``query_start_loc`` ranges (``[batch + 1]``); sequence ``b`` reads its + initial conv state from ``conv_states[cache_indices[b]]`` when + ``has_initial_state[b]`` and its last ``W - 1`` inputs are written back + there. ``conv_states`` is ``[slots, 3 * num_heads * head_dim, W - 1]`` for + ``conv_weight`` ``[3 * num_heads * head_dim, W]``. A sequence whose slot is + ``pad_slot_id`` reads no conv state and leaves the pool untouched. + + Returns contiguous ``[1, out_tokens, num_heads, head_dim]`` Q, K, V (Q and + K L2-normalized unless ``normalize_qk`` is False, for cores that normalize + them in the kernel) and, if ``beta`` is given, its fp32 + ``[1, out_tokens, num_heads]`` copy. ``out_tokens`` is ``padded_tokens`` + when given (rows past ``tokens`` are exact zeros), else ``tokens``. + """ + projection_size = num_heads * head_dim + num_channels = 3 * projection_size + if qkv.ndim != 2 or qkv.shape[1] != num_channels: + raise ValueError( + f"Fused KDA conv expected input shape [tokens, {num_channels}], got {tuple(qkv.shape)}" + ) + if ( + conv_weight.ndim != 2 + or conv_weight.shape[0] != num_channels + or not 2 <= conv_weight.shape[1] <= 4 + ): + raise ValueError( + f"Fused KDA conv expected weight shape [{num_channels}, 2..4], " + f"got {tuple(conv_weight.shape)}" + ) + width = conv_weight.shape[1] + if ( + conv_states.ndim != 3 + or tuple(conv_states.shape[1:]) != (num_channels, width - 1) + or conv_states.dtype != qkv.dtype + ): + raise ValueError( + f"Fused KDA conv expected {qkv.dtype} conv states [slots, {num_channels}, " + f"{width - 1}], got {conv_states.dtype} {tuple(conv_states.shape)}" + ) + num_seqs = cache_indices.numel() + if ( + cache_indices.ndim != 1 + or tuple(query_start_loc.shape) != (num_seqs + 1,) + or tuple(has_initial_state.shape) != (num_seqs,) + or has_initial_state.dtype != torch.bool + ): + raise ValueError( + "Fused KDA conv expected query_start_loc [batch + 1], cache_indices [batch] and " + f"bool has_initial_state [batch], got {tuple(query_start_loc.shape)}, " + f"{tuple(cache_indices.shape)} and {has_initial_state.dtype} " + f"{tuple(has_initial_state.shape)}" + ) + tensors = (conv_weight, conv_states, query_start_loc, cache_indices, has_initial_state) + if any(tensor.device != qkv.device for tensor in tensors): + raise ValueError("Fused KDA conv tensors must be on one device") + + num_tokens = qkv.shape[0] + out_tokens = num_tokens if padded_tokens is None else padded_tokens + if out_tokens < num_tokens: + raise ValueError(f"padded_tokens={out_tokens} is smaller than tokens={num_tokens}") + block_tokens = _FUSED_CONV_BLOCK_TOKENS + output_shape = (1, out_tokens, num_heads, head_dim) + q_out = torch.empty(output_shape, dtype=qkv.dtype, device=qkv.device) + k_out = torch.empty_like(q_out) + v_out = torch.empty_like(q_out) + beta_2d, beta_out = _beta_rows_and_output(beta, num_tokens, out_tokens, num_heads, qkv.device) + block_dim = triton.next_power_of_2(head_dim) + mixed_bf16_fma = torch.cuda.get_device_capability(qkv.device)[0] >= 10 + + with torch.cuda.device(qkv.device): + if out_tokens > 0: + _fused_kda_conv_prefill_kernel[(triton.cdiv(out_tokens, block_tokens), 3 * num_heads)]( + qkv, + conv_weight, + q_out, + k_out, + v_out, + beta_2d if beta_2d is not None else q_out, + beta_out if beta_out is not None else q_out, + num_tokens, + out_tokens, + qkv.stride(0), + qkv.stride(1), + conv_weight.stride(0), + conv_weight.stride(1), + beta_2d.stride(0) if beta_2d is not None else 0, + l2_norm_eps, + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + WIDTH=width, + BLOCK_TOKENS=block_tokens, + BLOCK_DIM=block_dim, + CHANNEL_CONTIGUOUS=qkv.stride(1) == 1, + WEIGHT_CHANNEL_CONTIGUOUS=conv_weight.stride(0) == 1, + HAS_BETA=beta is not None, + MIXED_BF16_FMA=mixed_bf16_fma, + NORMALIZE_QK=normalize_qk, + num_warps=_FUSED_CONV_NUM_WARPS, + ) + if num_seqs > 0: + # Runs after the kernel above: it overwrites that kernel's sequence-edge + # tokens, and it updates the conv states only after reading them. + _kda_conv_prefill_edges_kernel[(num_seqs, 3 * num_heads)]( + qkv, + conv_weight, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + q_out, + k_out, + v_out, + qkv.stride(0), + qkv.stride(1), + conv_weight.stride(0), + conv_weight.stride(1), + conv_states.stride(0), + conv_states.stride(1), + conv_states.stride(2), + l2_norm_eps, + pad_slot_id, + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + WIDTH=width, + BLOCK_EDGE=max(2, triton.next_power_of_2(width - 1)), + BLOCK_DIM=block_dim, + MIXED_BF16_FMA=mixed_bf16_fma, + NORMALIZE_QK=normalize_qk, + num_warps=4, + ) + return (q_out, k_out, v_out) if beta is None else (q_out, k_out, v_out, beta_out) + + +@triton.jit +def _copy_kda_replay_conv_window_kernel( + conv_ptr, + q_cache_ptr, + k_cache_ptr, + v_cache_ptr, + state_indices_ptr, + conv_stride_slot, + conv_stride_dim, + conv_stride_window, + cache_stride_slot, + cache_stride_dim, + cache_stride_window, + PROJECTION_SIZE: tl.constexpr, + COMMITTED: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + request = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < PROJECTION_SIZE * COMMITTED + dim = offsets // COMMITTED + window = offsets % COMMITTED + slot = tl.load(state_indices_ptr + request).to(tl.int64) + + conv_offset = slot * conv_stride_slot + dim * conv_stride_dim + window * conv_stride_window + cache_offset = slot * cache_stride_slot + dim * cache_stride_dim + window * cache_stride_window + section_offset = PROJECTION_SIZE * conv_stride_dim + tl.store(q_cache_ptr + cache_offset, tl.load(conv_ptr + conv_offset, mask=mask), mask=mask) + tl.store( + k_cache_ptr + cache_offset, + tl.load(conv_ptr + conv_offset + section_offset, mask=mask), + mask=mask, + ) + tl.store( + v_cache_ptr + cache_offset, + tl.load(conv_ptr + conv_offset + 2 * section_offset, mask=mask), + mask=mask, + ) + + +def copy_kda_replay_conv_window( + conv_pool: torch.Tensor, + q_cache: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + state_indices: torch.Tensor, +) -> None: + """Copy selected packed ``W - 1`` rows into KDA replay caches. + + The live convolution pool is ``[slots, 3D, W - 1]``. Replay caches are + ``[slots, D, W - 1 + num_spec]`` with a dim-contiguous inner layout. + """ + if conv_pool.ndim != 3 or conv_pool.shape[1] % 3: + raise ValueError(f"Expected packed rank-3 KDA convolution pool, got {conv_pool.shape}") + projection_size = conv_pool.shape[1] // 3 + committed = conv_pool.shape[2] + caches = (q_cache, k_cache, v_cache) + expected_prefix = (conv_pool.shape[0], projection_size) + if any(cache.ndim != 3 or cache.shape[:2] != expected_prefix for cache in caches): + raise ValueError("KDA replay convolution caches do not match the live pool geometry") + if any(cache.shape[2] < committed for cache in caches): + raise ValueError("KDA replay convolution caches are shorter than the committed window") + if any(cache.stride() != q_cache.stride() for cache in caches[1:]): + raise ValueError("KDA replay convolution caches must share one layout") + if state_indices.ndim != 1 or state_indices.dtype not in (torch.int32, torch.int64): + raise ValueError("KDA replay state indices must be a rank-1 int32 or int64 tensor") + if any(tensor.device != conv_pool.device for tensor in (*caches, state_indices)): + raise ValueError("KDA replay convolution tensors must be on one device") + if state_indices.numel() == 0: + return + + block_size = 256 + grid = ( + state_indices.numel(), + triton.cdiv(projection_size * committed, block_size), + ) + with torch.cuda.device(conv_pool.device.index): + _copy_kda_replay_conv_window_kernel[grid]( + conv_pool, + q_cache, + k_cache, + v_cache, + state_indices, + conv_pool.stride(0), + conv_pool.stride(1), + conv_pool.stride(2), + q_cache.stride(0), + q_cache.stride(1), + q_cache.stride(2), + PROJECTION_SIZE=projection_size, + COMMITTED=committed, + BLOCK_SIZE=block_size, + ) + + +_TORCH_TO_TL_DTYPE = { + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +@triton.jit +def _seed_kda_replay_conv_layers_kernel( + conv_ptrs, + cache_ptrs, + state_indices_ptr, + conv_stride_slot, + conv_stride_dim, + conv_stride_window, + cache_stride_slot, + cache_stride_dim, + cache_stride_window, + CONV_DTYPE: tl.constexpr, + CACHE_DTYPE: tl.constexpr, + PROJECTION_SIZE: tl.constexpr, + COMMITTED: tl.constexpr, + EXTENDED: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Seed one slot's q/k/v replay conv windows in one layer. + + ``conv_ptrs[layer]`` is the base address of that layer's live conv pool and + ``cache_ptrs[3 * layer + section]`` of its replay conv cache for ``section``. + The committed window is copied and the draft tail zeroed. + """ + request = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + layer = tl.program_id(2) + mask = offsets < PROJECTION_SIZE * EXTENDED + dim = offsets // EXTENDED + window = offsets % EXTENDED + committed = mask & (window < COMMITTED) + slot = tl.load(state_indices_ptr + request).to(tl.int64) + + conv = tl.load(conv_ptrs + layer).to(tl.pointer_type(CONV_DTYPE)) + conv_offset = slot * conv_stride_slot + dim * conv_stride_dim + window * conv_stride_window + cache_offset = slot * cache_stride_slot + dim * cache_stride_dim + window * cache_stride_window + section_offset = PROJECTION_SIZE * conv_stride_dim + for section in tl.static_range(3): + cache = tl.load(cache_ptrs + 3 * layer + section).to(tl.pointer_type(CACHE_DTYPE)) + values = tl.load(conv + conv_offset + section * section_offset, mask=committed, other=0.0) + tl.store(cache + cache_offset, values.to(CACHE_DTYPE), mask=mask) + + +@triton.jit +def _zero_kda_replay_slots_kernel( + cache_ptrs, + state_indices_ptr, + stride_slot, + stride_row, + stride_col, + DTYPE: tl.constexpr, + ROWS: tl.constexpr, + COLS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Zero one slot's ``[ROWS, COLS]`` entry of the layer at ``cache_ptrs[layer]``.""" + request = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + layer = tl.program_id(2) + mask = offsets < ROWS * COLS + row = offsets // COLS + col = offsets % COLS + slot = tl.load(state_indices_ptr + request).to(tl.int64) + cache = tl.load(cache_ptrs + layer).to(tl.pointer_type(DTYPE)) + tl.store( + cache + slot * stride_slot + row * stride_row + col * stride_col, + tl.zeros([BLOCK_SIZE], dtype=DTYPE), + mask=mask, + ) + + +@dataclasses.dataclass(frozen=True) +class _KDAReplayZeroGroup: + """One replay buffer kind, ``[slots, ROWS, COLS]`` in every layer.""" + + ptrs: torch.Tensor + strides: Tuple[int, int, int] + dtype: torch.dtype + rows: int + cols: int + + +@dataclasses.dataclass(frozen=True) +class KDAReplaySeedPlan: + """Device pointer tables and the shared geometry of every layer's replay buffers.""" + + num_layers: int + conv_ptrs: torch.Tensor + conv_cache_ptrs: torch.Tensor + conv_strides: Tuple[int, int, int] + cache_strides: Tuple[int, int, int] + conv_dtype: torch.dtype + cache_dtype: torch.dtype + projection_size: int + committed: int + extended: int + zero_groups: Tuple[_KDAReplayZeroGroup, ...] + + +def _shared_layout(tensors: Sequence[torch.Tensor]) -> bool: + reference = tensors[0] + return ( + reference.ndim == 3 + and reference.is_cuda + and reference.dtype in _TORCH_TO_TL_DTYPE + and all( + tensor.shape == reference.shape + and tensor.stride() == reference.stride() + and tensor.dtype == reference.dtype + and tensor.device == reference.device + for tensor in tensors + ) + ) + + +def _pointer_table(tensors: Sequence[torch.Tensor], device: torch.device) -> torch.Tensor: + """Device int64 table of base addresses, copied without waiting for queued GPU work.""" + host = torch.tensor( + [tensor.data_ptr() for tensor in tensors], dtype=torch.int64, pin_memory=prefer_pinned() + ) + return host.to(device, non_blocking=True) + + +def plan_kda_replay_seed( + conv_pools: Sequence[torch.Tensor], + conv_caches: Sequence[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + zero_buffers: Sequence[Sequence[torch.Tensor]], +) -> Optional[KDAReplaySeedPlan]: + """Plan seeding every layer's replay caches with one launch per buffer kind. + + ``conv_pools[layer]`` is the live ``[slots, 3D, W - 1]`` conv pool view of a + layer and ``conv_caches[layer]`` its q/k/v ``[slots, D, W - 1 + num_spec]`` + replay conv caches. Each entry of ``zero_buffers`` lists one replay buffer + kind (``[slots, rows, cols]``) across the layers; seeding zeroes it. The + views may live anywhere (e.g. per-layer views into KV cache pool pages) as + long as each kind shares one shape, stride and dtype across layers. + + Returns None, so the caller seeds layer by layer, when they do not. + """ + num_layers = len(conv_pools) + if num_layers == 0 or len(conv_caches) != num_layers: + return None + flat_caches = [cache for caches in conv_caches for cache in caches] + if not _shared_layout(conv_pools) or not _shared_layout(flat_caches): + return None + conv_pool = conv_pools[0] + cache = flat_caches[0] + if conv_pool.shape[1] % 3 or cache.device != conv_pool.device: + return None + projection_size = conv_pool.shape[1] // 3 + committed = conv_pool.shape[2] + if cache.shape[:2] != (conv_pool.shape[0], projection_size) or cache.shape[2] < committed: + return None + zero_groups = [] + for buffers in zero_buffers: + if len(buffers) != num_layers or not _shared_layout(buffers): + return None + reference = buffers[0] + if reference.device != conv_pool.device or reference.shape[0] != conv_pool.shape[0]: + return None + zero_groups.append( + _KDAReplayZeroGroup( + ptrs=_pointer_table(buffers, conv_pool.device), + strides=tuple(reference.stride()), + dtype=reference.dtype, + rows=reference.shape[1], + cols=reference.shape[2], + ) + ) + return KDAReplaySeedPlan( + num_layers=num_layers, + conv_ptrs=_pointer_table(conv_pools, conv_pool.device), + conv_cache_ptrs=_pointer_table(flat_caches, conv_pool.device), + conv_strides=tuple(conv_pool.stride()), + cache_strides=tuple(cache.stride()), + conv_dtype=conv_pool.dtype, + cache_dtype=cache.dtype, + projection_size=projection_size, + committed=committed, + extended=cache.shape[2], + zero_groups=tuple(zero_groups), + ) + + +def seed_kda_replay_layers(plan: KDAReplaySeedPlan, state_indices: torch.Tensor) -> None: + """Seed the replay caches of ``state_indices`` in every planned layer. + + Copies each layer's committed conv window into its q/k/v replay conv + caches, zeroes their draft tails and zeroes the planned buffer kinds; the + same result as doing so layer by layer. + """ + if state_indices.ndim != 1 or state_indices.dtype not in (torch.int32, torch.int64): + raise ValueError("KDA replay state indices must be a rank-1 int32 or int64 tensor") + if state_indices.numel() == 0: + return + block_size = 256 + num_slots = state_indices.numel() + with torch.cuda.device(plan.conv_ptrs.device.index): + _seed_kda_replay_conv_layers_kernel[ + ( + num_slots, + triton.cdiv(plan.projection_size * plan.extended, block_size), + plan.num_layers, + ) + ]( + plan.conv_ptrs, + plan.conv_cache_ptrs, + state_indices, + *plan.conv_strides, + *plan.cache_strides, + CONV_DTYPE=_TORCH_TO_TL_DTYPE[plan.conv_dtype], + CACHE_DTYPE=_TORCH_TO_TL_DTYPE[plan.cache_dtype], + PROJECTION_SIZE=plan.projection_size, + COMMITTED=plan.committed, + EXTENDED=plan.extended, + BLOCK_SIZE=block_size, + ) + for group in plan.zero_groups: + _zero_kda_replay_slots_kernel[ + (num_slots, triton.cdiv(group.rows * group.cols, block_size), plan.num_layers) + ]( + group.ptrs, + state_indices, + *group.strides, + DTYPE=_TORCH_TO_TL_DTYPE[group.dtype], + ROWS=group.rows, + COLS=group.cols, + BLOCK_SIZE=block_size, + ) + + +@triton.jit +def _store_kda_replay_conv_window_kernel( + conv_ptr, + q_cache_ptr, + k_cache_ptr, + v_cache_ptr, + state_indices_ptr, + num_accepted_ptr, + conv_stride_slot, + conv_stride_dim, + conv_stride_window, + cache_stride_slot, + cache_stride_dim, + cache_stride_window, + PROJECTION_SIZE: tl.constexpr, + COMMITTED: tl.constexpr, + NUM_SPEC: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + request = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < PROJECTION_SIZE * COMMITTED + dim = offsets // COMMITTED + window = offsets % COMMITTED + slot = tl.load(state_indices_ptr + request).to(tl.int64) + # The window after the accepted drafts starts at their count, clamped + # like the verify kernels clamp it. + shift = tl.minimum(tl.maximum(tl.load(num_accepted_ptr + request), 0), NUM_SPEC) + + conv_offset = slot * conv_stride_slot + dim * conv_stride_dim + window * conv_stride_window + cache_offset = ( + slot * cache_stride_slot + dim * cache_stride_dim + (window + shift) * cache_stride_window + ) + section_offset = PROJECTION_SIZE * conv_stride_dim + tl.store(conv_ptr + conv_offset, tl.load(q_cache_ptr + cache_offset, mask=mask), mask=mask) + tl.store( + conv_ptr + conv_offset + section_offset, + tl.load(k_cache_ptr + cache_offset, mask=mask), + mask=mask, + ) + tl.store( + conv_ptr + conv_offset + 2 * section_offset, + tl.load(v_cache_ptr + cache_offset, mask=mask), + mask=mask, + ) + + +def store_kda_replay_conv_window( + conv_pool: torch.Tensor, + q_cache: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> None: + """Store each selected slot's conv window after its accepted drafts in + the live convolution pool. + + The replay caches hold the window before the golden token in columns + ``[0, W - 1)`` and the raw inputs of the round's drafts after it, so the + ``W - 1`` columns starting at the request's ``num_accepted_tokens`` entry + are the window the next verify would start from. They are raw + convolution inputs, so storing them in a bf16 pool is exact. + """ + projection_size = conv_pool.shape[1] // 3 + committed = conv_pool.shape[2] + caches = (q_cache, k_cache, v_cache) + if any(cache.shape[1] != projection_size or cache.shape[2] < committed for cache in caches): + raise ValueError("KDA replay convolution caches do not match the live pool geometry") + if any(cache.stride() != q_cache.stride() for cache in caches[1:]): + raise ValueError("KDA replay convolution caches must share one layout") + if num_accepted_tokens.shape != state_indices.shape: + raise ValueError("Expected one accepted-draft count per state index") + if state_indices.numel() == 0: + return + + block_size = 256 + grid = ( + state_indices.numel(), + triton.cdiv(projection_size * committed, block_size), + ) + with torch.cuda.device(conv_pool.device.index): + _store_kda_replay_conv_window_kernel[grid]( + conv_pool, + q_cache, + k_cache, + v_cache, + state_indices, + num_accepted_tokens, + conv_pool.stride(0), + conv_pool.stride(1), + conv_pool.stride(2), + q_cache.stride(0), + q_cache.stride(1), + q_cache.stride(2), + PROJECTION_SIZE=projection_size, + COMMITTED=committed, + NUM_SPEC=q_cache.shape[2] - committed, + BLOCK_SIZE=block_size, + ) + + +@triton.jit +def _replay_kda_token_records_kernel( + state_ptr, + records_ptr, + state_indices_ptr, + num_accepted_ptr, + state_stride_slot, + records_stride_slot, + NUM_SPEC: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + K: tl.constexpr, + BLOCK_V: tl.constexpr, +): + request = tl.program_id(0) + head = tl.program_id(1) + rows = tl.program_id(2) * BLOCK_V + tl.arange(0, BLOCK_V) + keys = tl.arange(0, K) + slot = tl.load(state_indices_ptr + request).to(tl.int64) + accepted = tl.minimum(tl.maximum(tl.load(num_accepted_ptr + request), 0), NUM_SPEC) + state_ptrs = ( + state_ptr + slot * state_stride_slot + (head * V + rows[:, None]) * K + keys[None, :] + ) + state = tl.load(state_ptrs) + records = records_ptr + slot * records_stride_slot + for t in range(0, accepted): + vn = tl.load(records + (t * H + head) * V + rows) + beta_k = tl.load(records + NUM_SPEC * H * V + (t * H + head) * K + keys) + decay = tl.load(records + NUM_SPEC * H * (V + K) + (t * H + head) * K + keys) + state = tl.fma(decay[None, :], state, vn[:, None] * beta_k[None, :]) + tl.store(state_ptrs, state) + + +def replay_kda_token_records( + recurrent_state: torch.Tensor, + state_tok: torch.Tensor, + state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> None: + """Replay each selected slot's accepted drafts onto its pool state from + the per-draft records the Kimi K3 verify kernels keep in ``state_tok``. + + A slot of ``state_tok`` (``[slots, 3, num_spec, H, K]`` fp32) holds the + records of the last round's drafts: the row innovations ``vn`` + ``[num_spec][H][V]``, then ``beta * k`` and the decay + ``[num_spec][H][K]`` (V = K). Draft ``t`` updates the state as the + kernels do, ``S = fma(decay, S, vn (beta * k))``, so the result is + bit-identical to the state their next verify starts from. + """ + _, num_spec, num_heads, k_dim = state_tok.shape[1:] + v_dim = recurrent_state.shape[-2] + if state_tok.shape[1] != 3 or v_dim != k_dim: + raise ValueError( + "replay_kda_token_records: expected the drafts' records " + "[slots, 3, num_spec, H, K] with V = K, got " + f"{tuple(state_tok.shape)} for a state of " + f"{tuple(recurrent_state.shape)}" + ) + if state_indices.numel() == 0: + return + block_v = 32 + grid = (state_indices.numel(), num_heads, v_dim // block_v) + with torch.cuda.device(recurrent_state.device.index): + _replay_kda_token_records_kernel[grid]( + recurrent_state, + state_tok, + state_indices, + num_accepted_tokens, + recurrent_state.stride(0), + state_tok.stride(0), + NUM_SPEC=num_spec, + H=num_heads, + V=v_dim, + K=k_dim, + BLOCK_V=block_v, + ) + + +def commit_kda_pending_drafts( + recurrent_state: torch.Tensor, + conv_pool: torch.Tensor, + conv_q: torch.Tensor, + conv_k: torch.Tensor, + conv_v: torch.Tensor, + k_cache: Optional[torch.Tensor], + g_cache: Optional[torch.Tensor], + v_cache: Optional[torch.Tensor], + beta_cache: Optional[torch.Tensor], + state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + state_tok: Optional[torch.Tensor] = None, +) -> None: + """Commit one layer's pending accepted drafts into its live pools. + + ``num_accepted_tokens[n]`` is the number of drafts of slot + ``state_indices[n]`` the last verify accepted, which the next verify + would take up first. Afterwards ``recurrent_state`` and ``conv_pool`` + hold the state and the conv window that verify would start from (the + window bit for bit), and the caller clears the slot's count: + + * with the Kimi K3 verify kernels' draft records (``state_tok``), the + drafts are replayed from those (``replay_kda_token_records``), and the + replay caches, which are then not allocated, may be None; + * otherwise they are replayed from the replay caches by the fused verify + kernel's own replay steps (``kda_mtp_commit_pending_drafts``). The + tensor-core verify computes a chunk in one pass, so this commit + rounds the state differently from a verify that replays the drafts + itself. + + The replay-only commit is main's (``REPLAY_ONLY`` in main's verify kernel); the branch has + none. It reads the same replay caches the private verify ops write. + """ + num_spec = conv_q.shape[2] - conv_pool.shape[2] + if state_tok is not None: + replay_kda_token_records(recurrent_state, state_tok, state_indices, num_accepted_tokens) + else: + assert ( + k_cache is not None + and g_cache is not None + and v_cache is not None + and beta_cache is not None + ) + _main_mtp_ops().kda_mtp_commit_pending_drafts( + recurrent_state, + conv_q, + conv_k, + conv_v, + k_cache, + g_cache, + v_cache, + beta_cache, + state_indices, + num_accepted_tokens, + num_spec, + conv_width=conv_pool.shape[2] + 1, + ) + store_kda_replay_conv_window( + conv_pool, conv_q, conv_k, conv_v, state_indices, num_accepted_tokens + ) + + +# --------------------------------------------------------------------------- +# Prefill layout and state-pool helpers: the channel-major convolution input of the unfused +# prefill convolution, and the reset of fresh requests' state rows. +# --------------------------------------------------------------------------- + + +@triton.jit +def _extract_transpose_prefill_kernel( + src_ptr, + dst_ptr, + num_prefill_tokens, + src_stride_seq, + d_inner, + conv_dim, + BLOCK_SEQ: tl.constexpr, + BLOCK_CONV: tl.constexpr, +): + """Extract src[0:num_prefill_tokens, d_inner:d_inner+conv_dim] and + transpose to dst[conv_dim, num_prefill_tokens].""" + pid_seq = tl.program_id(0) + pid_conv = tl.program_id(1) + + seq_offsets = pid_seq * BLOCK_SEQ + tl.arange(0, BLOCK_SEQ) + conv_offsets = pid_conv * BLOCK_CONV + tl.arange(0, BLOCK_CONV) + + seq_mask = seq_offsets < num_prefill_tokens + conv_mask = conv_offsets < conv_dim + mask = seq_mask[:, None] & conv_mask[None, :] + + # Cast to int64 to avoid overflow: seq_offsets * src_stride_seq can exceed + # INT32_MAX (e.g., 131071 * 22656 = 2,969,544,576 > 2,147,483,647) + src_offsets = ( + seq_offsets[:, None].to(tl.int64) * src_stride_seq + d_inner + conv_offsets[None, :] + ) + data = tl.load(src_ptr + src_offsets, mask=mask, other=0.0) + + dst_offsets = conv_offsets[:, None] * num_prefill_tokens + seq_offsets[None, :] + tl.store(dst_ptr + dst_offsets, tl.trans(data), mask=conv_mask[:, None] & seq_mask[None, :]) + + +def extract_transpose_prefill_slice( + src: torch.Tensor, + num_prefill_tokens: int, + start_col: int, + width: int, +) -> torch.Tensor: + """ + Extract and transpose a prefill slice for the channel-major causal convolution. + + Input: src[num_tokens, num_cols], rows contiguous (arbitrary row stride, + so column-slice views of a wider tensor work in place) + Output: [width, num_prefill_tokens] + """ + assert src.stride(1) == 1 + out = torch.empty(width, num_prefill_tokens, dtype=src.dtype, device=src.device) + + BLOCK_SEQ, BLOCK_CONV = 32, 128 + grid = (triton.cdiv(num_prefill_tokens, BLOCK_SEQ), triton.cdiv(width, BLOCK_CONV)) + + _extract_transpose_prefill_kernel[grid]( + src, + out, + num_prefill_tokens, + src.stride(0), + start_col, + width, + BLOCK_SEQ, + BLOCK_CONV, + ) + return out + + +@triton.jit +def _reset_recurrent_state_rows_kernel( + recurrent_states, + conv_states, + state_indices, + has_initial_states, + recurrent_state_stride, + conv_state_stride, + NUM_CACHE_LINES: tl.constexpr, + RECURRENT_STATE_SIZE: tl.constexpr, + CONV_STATE_SIZE: tl.constexpr, + RESET_CONV: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + request_idx = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state_idx = tl.load(state_indices + request_idx).to(tl.int64) + needs_reset = ~tl.load(has_initial_states + request_idx).to(tl.int1) + valid_state = (state_idx >= 0) & (state_idx < NUM_CACHE_LINES) + recurrent_row_offset = state_idx * recurrent_state_stride.to(tl.int64) + tl.store( + recurrent_states + recurrent_row_offset + offsets, + 0.0, + mask=needs_reset & valid_state & (offsets < RECURRENT_STATE_SIZE), + ) + if RESET_CONV: + conv_row_offset = state_idx * conv_state_stride.to(tl.int64) + tl.store( + conv_states + conv_row_offset + offsets, + 0.0, + mask=needs_reset & valid_state & (offsets < CONV_STATE_SIZE), + ) + + +def reset_recurrent_state_rows( + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + conv_states: Optional[torch.Tensor] = None, +) -> None: + """Zero fresh-request rows in recurrent and optional convolution pools. + + Invalid or negative slot indices are ignored, matching the cache-manager + reset contract used by GDN and KDA. The launch is ordered on PyTorch's + current stream, so a following indexed recurrence observes the zeroed rows. + """ + num_requests = state_indices.shape[0] + if num_requests == 0: + return + if recurrent_states.shape[0] == 0: + raise ValueError("recurrent state pool must have at least one cache line") + if conv_states is not None and conv_states.shape[0] != recurrent_states.shape[0]: + raise ValueError("recurrent and convolution pools must have the same slot count") + if has_initial_states.shape != state_indices.shape: + raise ValueError("has_initial_states and state_indices must have matching shapes") + recurrent_state_size = recurrent_states.numel() // recurrent_states.shape[0] + conv_state_size = conv_states.numel() // conv_states.shape[0] if conv_states is not None else 0 + if not recurrent_states[0].is_contiguous(): + raise ValueError("recurrent state rows must be contiguous") + if recurrent_states.stride(0) < recurrent_state_size: + raise ValueError("recurrent state rows must not overlap") + if conv_states is not None and not conv_states[0].is_contiguous(): + raise ValueError("convolution state rows must be contiguous") + if conv_states is not None and conv_states.stride(0) < conv_state_size: + raise ValueError("convolution state rows must not overlap") + + block_size = 256 + grid = ( + num_requests, + triton.cdiv(max(recurrent_state_size, conv_state_size), block_size), + ) + _reset_recurrent_state_rows_kernel[grid]( + recurrent_states, + conv_states if conv_states is not None else recurrent_states, + state_indices, + has_initial_states, + recurrent_states.stride(0), + conv_states.stride(0) if conv_states is not None else 0, + recurrent_states.shape[0], + recurrent_state_size, + conv_state_size, + conv_states is not None, + block_size, + ) + + +# --------------------------------------------------------------------------- +# The CuTe DSL KDA entries of the package: ops/ssm/kda_prefill (_kda_prefill.py over the CuTe DSL +# prefill kernels) and ops/ssm/kda_mtp_decode (mtp_ops.py over the v1 / v2 / two-CTA verify +# kernels). Each is imported when a dispatch is built on a supported GPU, else on first use; +# importing either raises ImportError without CuTe DSL. +# --------------------------------------------------------------------------- + + +@functools.cache +def _prefill_entry(): + """The prefill entry (``ops/ssm/kda_prefill``), imported once.""" + from ...ops.ssm.kda_prefill import kda_prefill + + return kda_prefill + + +@functools.cache +def _mtp_ops() -> ModuleType: + """The verify launcher module (``mtp_ops.py``), imported once.""" + from . import mtp_ops + + return mtp_ops + + +@functools.cache +def _mtp_entry(): + """The verify entry (``ops/ssm/kda_mtp_decode``), imported once.""" + from ...ops.ssm.kda_mtp_decode import kda_mtp_decode + + return kda_mtp_decode + + +@functools.cache +def _main_mtp_ops() -> ModuleType: + """Main's verify ops: the replay-only draft commit of commit_kda_pending_drafts.""" + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops + + return cute_dsl_kimi_k3_kda_mtp_ops + + +def kda_mtp_verify_ctas(num_heads: int, num_requests: int, device: torch.device) -> int: + """CTAs of the v2 verify kernel's grid for these tiles (cluster pairs up to + ``SPLIT_MAX_TILES`` tiles, or one CTA per tile).""" + return _mtp_ops().verify_ctas(num_heads, num_requests, device) + + +# --------------------------------------------------------------------------- +# Dispatch API used by the K3KDALinearAttention module. +# --------------------------------------------------------------------------- + + +class K3KDAKernelDispatch: + """The KDA kernel calls of one ``K3KDALinearAttention``: the prefill entry, the fused + single-token decode kernel and the fused multi-token verify entry.""" + + _selection_logged = False + + def __init__(self) -> None: + if is_kda_optimized_supported(): + # Import the entries at construction, before any forward (and any graph capture). + _prefill_entry() + _mtp_entry() + # One line per process so runs record which kernels the KDA layers call. + if not K3KDAKernelDispatch._selection_logged: + K3KDAKernelDispatch._selection_logged = True + from tensorrt_llm.logger import logger + + logger.info( + "KDA kernels: prefill ops/ssm/kda_prefill, decode CUDA fused decode, " + "verify ops/ssm/kda_mtp_decode" + ) + + def mtp_verify( + self, *, mxfp8_output: bool = False, **kwargs + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Fused KDA multi-token verify (``ops/ssm/kda_mtp_decode``): the BF16 core, or with + ``mxfp8_output`` the E4M3 core and its R128c4 scales, which always fuses the output norm + and launches without ``use_pdl`` (see ``mtp_ops.py`` for the argument and + state-management contract).""" + if mxfp8_output: + if not kwargs.pop("fuse_output_norm", True): + raise ValueError("the MXFP8-output KDA MTP op always fuses the output norm") + kwargs.pop("use_pdl", None) + return _mtp_entry()(**kwargs, fuse_output_norm=True, quantize_output=True) + return _mtp_entry()(**kwargs) + + def mtp_verify_fused_fb(self, **kwargs) -> torch.Tensor: + """Fused KDA multi-token verify that computes its forget gate from ``f_a`` + (``ops/ssm/kda_mtp_decode`` with ``fa`` / ``w_fb``: packed rows, fused gated RMSNorm, + BF16 output).""" + return _mtp_entry()( + **kwargs, + g=None, + cu_seqlens=None, + packed_token_layout=True, + fuse_output_norm=True, + ) + + @staticmethod + def can_use_indexed_prefill( + *, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_sequences: int, + ) -> bool: + """Whether ``ops/ssm/kda_prefill`` takes these rows: fp32 V-first + ``[slots, H, V, K]`` rows and int32 / int64 16-byte-aligned indices.""" + if cu_seqlens is not None: + if cu_seqlens.shape[0] - 1 != num_sequences: + return False + if state_pool.dtype != torch.float32 or state_pool.ndim != 4: + return False + _, H, V, K = state_pool.shape + if state_pool.stride()[1:] != (V * K, K, 1): + return False + if state_pool.stride(0) < H * V * K or state_pool.stride(0) % 4: + return False + if state_pool.data_ptr() % 16: + return False + if ( + state_indices.ndim != 1 + or state_indices.shape[0] != num_sequences + or state_indices.dtype not in (torch.int32, torch.int64) + or not state_indices.is_contiguous() + ): + return False + if state_indices.data_ptr() % 16: + return False + if ( + has_initial_states.ndim != 1 + or has_initial_states.shape[0] != num_sequences + or has_initial_states.dtype != torch.bool + or not has_initial_states.is_contiguous() + ): + return False + if ( + state_pool.device != state_indices.device + or state_pool.device != has_initial_states.device + ): + return False + return True + + @staticmethod + def prefill_chunk_kda( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + lower_bound: float, + cu_seqlens: Optional[torch.Tensor], + chunk_indices: Optional[torch.Tensor], + chunk_size: int, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + varlen_is_aligned: Optional[bool], + single_sequence_length: Optional[int], + snapshot_rows: Optional[torch.Tensor], + snapshot_chunks: Optional[torch.Tensor], + ) -> torch.Tensor: + """``ops/ssm/kda_prefill`` on L2-normalized Q/K, updating the indexed fp32 V-first + ``[slots, H, V, K]`` rows of ``state_pool`` in place; returns the core ``[1, T, H, V]``. + + ``chunk_indices``, ``varlen_is_aligned`` and ``single_sequence_length`` are prepared once + from the K3 runtime metadata so no layer reads ``cu_seqlens`` back from the GPU. + ``snapshot_rows`` / ``snapshot_chunks`` (int32 ``[num_sequences]``) also store each + selected sequence's state after its first ``chunk_size * snapshot_chunks[b]`` tokens into + ``state_pool[snapshot_rows[b]]``; ``-1`` skips a sequence. + + A single sequence that is not a chunk multiple is zero-padded to one (Q/K/V/beta; the op + rebuilds g's tail itself) unless the caller already padded it (inputs longer than + ``single_sequence_length``); the output is cropped back to the real length. + """ + snapshot_kwargs = ( + {} + if snapshot_rows is None + else {"snapshot_rows": snapshot_rows, "snapshot_chunks": snapshot_chunks} + ) + real_T = q.shape[1] + if ( + cu_seqlens is not None + and single_sequence_length is not None + and q.shape[1] != single_sequence_length + ): + # Caller already produced chunk-padded inputs (zero tails); + # only the output needs trimming back to the real length. + real_T = single_sequence_length + padded_T = q.shape[1] + if ( + cu_seqlens.shape[0] != 2 + or padded_T % chunk_size != 0 + or not 0 < padded_T - real_T < chunk_size + or any(t.shape[1] != padded_T for t in (k, v, beta)) + or g.shape[1] not in (real_T, padded_T) + ): + raise ValueError( + "Pre-padded KDA prefill inputs must be one sequence padded to the " + f"next {chunk_size}-token boundary (real_T={real_T}, " + f"q/k/v/g/beta lengths={[t.shape[1] for t in (q, k, v, g, beta)]})." + ) + elif cu_seqlens is not None: + # The op's varlen single-sequence path expects the packed tensors zero-padded to a + # chunk multiple while cu_seqlens keeps the real length; the op re-sentinels g's tail + # itself. Without this the op would run the mask-free kernel on a partial final + # chunk. Multi-sequence varlen runs the masked path and needs no pad. + if cu_seqlens.shape[0] == 2 and real_T % chunk_size != 0: + pad = chunk_size - real_T % chunk_size + q = F.pad(q, (0, 0, 0, 0, 0, pad)) + k = F.pad(k, (0, 0, 0, 0, 0, pad)) + v = F.pad(v, (0, 0, 0, 0, 0, pad)) + g = F.pad(g, (0, 0, 0, 0, 0, pad)) + beta = F.pad(beta, (0, 0, 0, pad)) + + out = _prefill_entry()( + q=q, + k=k, + v=v, + g=g, + beta=beta, + state_pool=state_pool, + state_indices=state_indices, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=True, + lower_bound=lower_bound, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + A_log=A_log.detach(), + dt_bias=dt_bias.detach(), + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, + **snapshot_kwargs, + ) + if out.shape[1] != real_T: + out = out[:, :real_T] + return out + + @staticmethod + def decode_kda(**kwargs) -> torch.Tensor: + """The fused CUDA single-token decode kernel (``_kda_decode``).""" + return _kda_decode.run_kda_decode_fusion_cuda(**kwargs) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/akk_inverse.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/akk_inverse.py new file mode 100644 index 000000000000..a8a96aa43bac --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/akk_inverse.py @@ -0,0 +1,1075 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Akk 64×64 Lower Triangular Block Inversion — mixed FP32/BF16. + +BF16 input/output and packed BF16 shared storage. The four 16×16 diagonal +blocks use FP32 forward substitution to avoid unstable BF16 Neumann +cancellation; the cross-block products retain BF16 MMA with FP32 accumulators. + +Architecture: + - 4 warps (128 threads) per CTA, each CTA processes one 64×64 Akk matrix + - sAkk [64,36] fp32 — each FP32 slot holds packed bf16x2 (stride 36) + Total SMEM: 64*36*4 = 9216 bytes (vs 64*72*4 = 18432 in FP32 version) + - sTemp [16, 24, 2] fp32: inter-warp FP32 accumulator communication + - cp.async: BF16 global → packed bf16x2 in FP32 SMEM (raw bytes align) + - FP32 forward substitution for diagonal blocks + - BF16 MMA m16n8k16 with FP32 accumulators for cross-block products + - movmatrix.sync.aligned.m8n8.trans.b16 for A→B layout conversion + - C→A chain: cvt.rn.bf16x2.f32 to pack FP32 accum → bf16x2 A-operand + +Block layout in sAkk (4×4 sub-blocks of 16×16, packed bf16x2): + Upper tri = INPUT, Diagonal = in-place inversion, Lower tri = OUTPUT + +Stages: + 0. cp.async load bf16 64×64 → sAkk (packed bf16x2) + 1. Invert 4 diagonal blocks via FP32 forward substitution (2 warps) + 2. Warps 0-2: Ai10, Ai21, Ai32 via chain MMA (C→A pack) + 3. Warps 0+2 → Ai20, warps 1+3 → Ai31 (parallel pairs, sTemp) + 4. Warps 0+1+2 → Ai30 (sTemp aggregation) + 5. All warps: store bf16x2 SMEM → bf16 global + +Inputs: A_in [B, T, H, BT] bf16 +Outputs: A_out [B, T, H, BT] bf16 +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import T, dsl_user_op + +# =========================================================================== +# Constants +# =========================================================================== +BS = 64 +SB = 16 +THREADS = 128 +TEMP_PAD = 8 +TEMP_COLS = SB + TEMP_PAD # 24 +NUM_TEMPS = 2 +AKK_PAD = 4 +AKK_STRIDE = BS // 2 + AKK_PAD # 36 (in FP32 units = 72 bf16 elements) + + +# =========================================================================== +# BF16 MMA m16n8k16 with FP32 accumulator +# =========================================================================== +@dsl_user_op +def mma_bf16_m16n8k16( + a0, + a1, + a2, + a3, + b0, + b1, + c0, + c1, + c2, + c3, + *, + loc=None, + ip=None, +): + """mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + A: 4×i32 (bf16x2 pairs), B: 2×i32, C/D: 4×f32 (FP32 accum).""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +# =========================================================================== +# movmatrix.sync.aligned.m8n8.trans.b16 — hardware A→B transpose (warp-level) +# Works on any 16-bit format including BF16. +# =========================================================================== +@dsl_user_op +def _movmatrix_trans(src, *, loc=None, ip=None): + src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [src_b], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +# =========================================================================== +# BF16x2 pack/unpack helpers +# =========================================================================== +@dsl_user_op +def _pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None): + """Pack two FP32 values into one i32 holding two BF16 values. + cvt.rn.bf16x2.f32 converts and packs in one instruction.""" + result = llvm.inline_asm( + T.i32(), + [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _mask_packed_ltri(packed, row, pair, *, loc=None, ip=None): + """Apply lower-triangular mask to packed bf16x2. + Zero out elements where row < col. col0=2*pair, col1=2*pair+1.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)], + """{ + .reg .b32 %c0, %c1, %mlo, %mhi, %mask; + .reg .pred %p0, %p1; + shl.b32 %c0, $3, 1; + add.u32 %c1, %c0, 1; + setp.ge.s32 %p0, $2, %c0; + setp.ge.s32 %p1, $2, %c1; + selp.b32 %mlo, 0xFFFF, 0, %p0; + selp.b32 %mhi, 0xFFFF0000, 0, %p1; + or.b32 %mask, %mlo, %mhi; + and.b32 $0, $1, %mask; + }""", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _unpack_bf16x2_lo(packed, *, loc=None, ip=None): + """Unpack lower BF16 from packed i32 to FP32. + BF16 is upper 16 bits of FP32, so shift left by 16.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + """{ + shl.b32 $0, $1, 16; + }""", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _unpack_bf16x2_hi(packed, *, loc=None, ip=None): + """Unpack upper BF16 from packed i32 to FP32. + Upper 16 bits are already in the right position for FP32.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + """{ + and.b32 $0, $1, 0xFFFF0000; + }""", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +# =========================================================================== +# Neumann diagonal 16×16 inversion using BF16 MMA m16n8k16 + FP32 accum +# +# (I+L)⁻¹ = (I-L)(I+L²)(I+L⁴)(I+L⁸) +# +# All intermediate values kept in FP32. Pack to bf16x2 only at MMA boundaries. +# sAkk is packed bf16x2: sAkk[row, pair] = {bf16[row, 2*pair], bf16[row, 2*pair+1]} +# =========================================================================== +@dsl_user_op +def _invert_diag_neumann(sAkk: cute.Tensor, block_idx, lane_id, *, loc=None, ip=None): + r_off = block_idx * 16 + c_off = block_idx * 8 # packed bf16x2: 16 cols → 8 pairs + gid = lane_id // 4 + tid = lane_id % 4 + + # --- Load A from packed bf16x2 SMEM → unpack to 8 FP32 values --- + packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid]) + packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid]) + packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid]) + packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid]) + + A_f0 = _unpack_bf16x2_lo(packed0) # A[gid, 2*tid] + A_f1 = _unpack_bf16x2_hi(packed0) # A[gid, 2*tid+1] + A_f2 = _unpack_bf16x2_lo(packed1) # A[gid+8, 2*tid] + A_f3 = _unpack_bf16x2_hi(packed1) # A[gid+8, 2*tid+1] + A_f4 = _unpack_bf16x2_lo(packed2) # A[gid, 8+2*tid] + A_f5 = _unpack_bf16x2_hi(packed2) # A[gid, 8+2*tid+1] + A_f6 = _unpack_bf16x2_lo(packed3) # A[gid+8, 8+2*tid] + A_f7 = _unpack_bf16x2_hi(packed3) # A[gid+8, 8+2*tid+1] + + # Build identity (FP32) + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid) + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + 1 + ) + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid) + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + 1 + ) + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + ) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + 1 + ) + + # L = A - I (FP32) + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + + # INV = I - L = 2I - A (FP32) + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + + # Pack L and INV → bf16x2 for MMA + L_a0 = _pack_bf16x2(L_f0, L_f1) + L_a1 = _pack_bf16x2(L_f2, L_f3) + L_a2 = _pack_bf16x2(L_f4, L_f5) + L_a3 = _pack_bf16x2(L_f6, L_f7) + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 1: L² = L × L, then INV = INV + INV × L² === + L_b0 = _movmatrix_trans(L_a0) + L_b1 = _movmatrix_trans(L_a1) + L_b2 = _movmatrix_trans(L_a2) + L_b3 = _movmatrix_trans(L_a3) + + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = mma_bf16_m16n8k16( + L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf + ) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = mma_bf16_m16n8k16( + L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf + ) + + Lp_a0 = _pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _pack_bf16x2(Lp_c6, Lp_c7) + + Lp_b0 = _movmatrix_trans(Lp_a0) + Lp_b1 = _movmatrix_trans(Lp_a1) + Lp_b2 = _movmatrix_trans(Lp_a2) + Lp_b3 = _movmatrix_trans(Lp_a3) + + # mm = INV × L² (FP32 accum output) + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + + # INV += mm (FP32) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 2: L⁴ = L² × L², then INV = INV + INV × L⁴ === + L4_c0, L4_c1, L4_c2, L4_c3 = mma_bf16_m16n8k16( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + L4_c4, L4_c5, L4_c6, L4_c7 = mma_bf16_m16n8k16( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + + L4_a0 = _pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _pack_bf16x2(L4_c6, L4_c7) + + L4_b0 = _movmatrix_trans(L4_a0) + L4_b1 = _movmatrix_trans(L4_a1) + L4_b2 = _movmatrix_trans(L4_a2) + L4_b3 = _movmatrix_trans(L4_a3) + + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 3: L⁸ = L⁴ × L⁴, then INV = INV + INV × L⁸ === + L8_c0, L8_c1, L8_c2, L8_c3 = mma_bf16_m16n8k16( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + L8_c4, L8_c5, L8_c6, L8_c7 = mma_bf16_m16n8k16( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + + L8_a0 = _pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _pack_bf16x2(L8_c6, L8_c7) + + L8_b0 = _movmatrix_trans(L8_a0) + L8_b1 = _movmatrix_trans(L8_a1) + L8_b2 = _movmatrix_trans(L8_a2) + L8_b3 = _movmatrix_trans(L8_a3) + + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf + ) + + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + # --- Store INV to packed bf16x2 SMEM --- + sAkk[r_off + gid, c_off + tid] = _pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _load_packed_bf16(sAkk: cute.Tensor, row, packed_col_base, col, *, loc=None, ip=None): + """Load one logical BF16 value from packed-bf16x2 shared storage.""" + valid_i = cutlass.Int32(col >= 0) + safe_col = col * valid_i + packed = cutlass.Float32(sAkk[row, packed_col_base + safe_col // 2]) + lo = _unpack_bf16x2_lo(packed) + hi = _unpack_bf16x2_hi(packed) + use_hi = cutlass.Float32(safe_col % 2) + value = lo * (cutlass.Float32(1.0) - use_hi) + hi * use_hi + return value * cutlass.Float32(valid_i) + + +@dsl_user_op +def _invert_diag_forward_fp32( + sAkk: cute.Tensor, + block_idx, + lane_id, + *, + loc=None, + ip=None, +): + """FP32 forward substitution with one final BF16 pack per inverse row.""" + row = lane_id % 16 + halfwarp_base = (lane_id // 16) * 16 + r_off = block_idx * 16 + c_off = block_idx * 8 + inv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32) + inv[0] = cutlass.Float32(1.0) + for d in range(1, 16): + inv[d] = cutlass.Float32(0.0) + + for d in range(1, 16): + col_d = row - d + valid = cutlass.Float32(col_d >= 0) + a_val = _load_packed_bf16(sAkk, r_off + row, c_off, col_d) + acc = cutlass.Float32(0.0) + for j in range(1, d): + a_re = _load_packed_bf16(sAkk, r_off + row, c_off, row - (d - j)) + inv_prev = cute.arch.shuffle_sync(inv[j], halfwarp_base + row - d + j) + acc = acc + a_re * inv_prev + inv[d] = (-a_val - acc) * valid + + for pair in range(8): + col0 = pair * 2 + col1 = col0 + 1 + d0 = row - col0 + d1 = row - col1 + v0 = cutlass.Float32(row == col0) + v1 = cutlass.Float32(row == col1) + for d in range(1, 16): + v0 = v0 + inv[d] * cutlass.Float32(d0 == d) + v1 = v1 + inv[d] * cutlass.Float32(d1 == d) + sAkk[r_off + row, c_off + pair] = _pack_bf16x2(v0, v1) + + +# =========================================================================== +# 16×16 matmul: load A & B from packed bf16x2 sAkk, BF16 MMA, return FP32 C +# =========================================================================== +@dsl_user_op +def _matmul_AB(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 # packed: 16 cols → 8 pairs + rB = br_B * 16 + cB = bc_B * 8 + + # Load A-operand (packed bf16x2, direct from SMEM) + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + + # Load B-operand: load in A-layout then movmatrix → B-layout + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _movmatrix_trans(bA0) + b1 = _movmatrix_trans(bA1) + b2 = _movmatrix_trans(bA2) + b3 = _movmatrix_trans(bA3) + + # 16×16 = 2 × m16n8k16 + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + + # Return 8 FP32 C-registers (C-layout of m16n8k16 FP32 accum) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# =========================================================================== +# Chain MMA: pre-loaded A (from C→A pack), load B from sAkk (Stage 2) +# A-operand already packed as bf16x2 from previous C result. +# =========================================================================== +@dsl_user_op +def _chain_mma_B(sAkk: cute.Tensor, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 8 + + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _movmatrix_trans(bA0) + b1 = _movmatrix_trans(bA1) + b2 = _movmatrix_trans(bA2) + b3 = _movmatrix_trans(bA3) + + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# =========================================================================== +# Chain MMA: load A from sAkk, pre-loaded B (from C→B shuffle) (Stages 3-4) +# B-operand from shuffle is still FP32 TF32-layout → need to convert. +# Actually in BF16 version we use movmatrix, so B comes from pack+movmatrix. +# This function takes B already in B-layout (packed bf16x2 after movmatrix). +# =========================================================================== +@dsl_user_op +def _chain_mma_A(sAkk: cute.Tensor, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# =========================================================================== +# Store negated C result (16×16, FP32 accum) to packed bf16x2 sAkk +# C-layout for m16n8k16 FP32 accum: +# cn0_0 = C[gid, 0..7 left half col0], cn0_1 = C[gid+8, left half col0] +# cn0_2 = C[gid, left half col1], cn0_3 = C[gid+8, left half col1] +# cn1_* = right half +# Packing: negate FP32 then pack pairs → bf16x2 +# =========================================================================== +@dsl_user_op +def _store_neg_C( + sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 8 # packed + + # Pack negated FP32 pairs → bf16x2, then store + sAkk[r + gid, c + tid] = _pack_bf16x2(-c0, -c1) + sAkk[r + gid + 8, c + tid] = _pack_bf16x2(-c2, -c3) + sAkk[r + gid, c + 4 + tid] = _pack_bf16x2(-c4, -c5) + sAkk[r + gid + 8, c + 4 + tid] = _pack_bf16x2(-c6, -c7) + + +# =========================================================================== +# Pack FP32 C-accum → bf16x2 A-operand for C→A chain +# C-layout (FP32 accum, m16n8k16): 8 floats → 4 bf16x2 A-regs +# c0,c1 → a0 (rows gid/gid+8, left-half k0..7) +# c2,c3 → a1 +# c4,c5 → a2 (right-half k8..15) +# c6,c7 → a3 +# =========================================================================== +@dsl_user_op +def _pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _pack_bf16x2(c0, c1) + a1 = _pack_bf16x2(c2, c3) + a2 = _pack_bf16x2(c4, c5) + a3 = _pack_bf16x2(c6, c7) + return a0, a1, a2, a3 + + +# =========================================================================== +# Convert FP32 C-accum → bf16x2 B-operand via pack + movmatrix +# =========================================================================== +@dsl_user_op +def _pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _pack_bf16x2(c0, c1) + a1 = _pack_bf16x2(c2, c3) + a2 = _pack_bf16x2(c4, c5) + a3 = _pack_bf16x2(c6, c7) + b0 = _movmatrix_trans(a0) + b1 = _movmatrix_trans(a1) + b2 = _movmatrix_trans(a2) + b3 = _movmatrix_trans(a3) + return b0, b1, b2, b3 + + +# =========================================================================== +# sTemp helpers (FP32, non-swizzled, for inter-warp accumulator exchange) +# =========================================================================== +@dsl_user_op +def _store_C_temp( + sT: cute.Tensor, + buf, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _load_C_temp(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +# =========================================================================== +# Main kernel +# =========================================================================== +@cute.kernel +def akk_inv_kernel( + g2s_copy: cute.TiledCopy, + gA_tensor: cute.Tensor, + mOut: cute.Tensor, + mBeta: cute.Tensor, + akk_smem_layout: cute.Layout, + temp_layout: cute.Layout, + NT: cutlass.Int32, + H: int, + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_id = tidx % 32 + h_idx, nt_idx, b_idx = cute.arch.block_idx() + + # ===== SMEM allocation ===== + smem = cutlass.utils.SmemAllocator() + sAkk = smem.allocate_tensor(cutlass.Float32, akk_smem_layout, 128) + sTemp = smem.allocate_tensor(cutlass.Float32, temp_layout, 128) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout(BS, stride=1), 128) + + # ===== Stage 0: cp.async load bf16 global → packed bf16x2 sAkk ===== + ld_bnt = nt_idx + if IS_VARLEN: + ld_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + ld_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + ld_bos = cutlass.Int32(mCuSeqlens[ld_seq_id]) + ld_bnt = ld_bos + ld_local * BS + gA_batch = gA_tensor[(None, None, h_idx, ld_bnt, b_idx)] + + thr_g2s = g2s_copy.get_slice(tidx) + thr_gSrc = thr_g2s.partition_S(gA_batch) + thr_sDst = thr_g2s.partition_D(sAkk) + cute.copy(g2s_copy, thr_gSrc, thr_sDst) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + # Zero out-of-bounds rows for varlen partial chunks. + if IS_VARLEN: + _z_seq = cutlass.Int32(mChunkIndices[nt_idx, 0]) + _z_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + _z_bos = cutlass.Int32(mCuSeqlens[_z_seq]) + _z_eos = cutlass.Int32(mCuSeqlens[_z_seq + 1]) + _z_cs = _z_bos + _z_local * BS + _z_vr = _z_eos - _z_cs + _zr_start = warp_idx * SB + for ri in cutlass.range_constexpr(SB): + row = _zr_start + ri + if row >= _z_vr: + # Zero packed bf16x2 slots (each slot = 2 bf16 elements) + c0 = lane_id + if c0 < AKK_STRIDE: + sAkk[row, c0] = cutlass.Float32(0.0) + cute.arch.barrier() + + # ===== Stage 0b: Load 64 per-token betas into sBeta (fp32) ===== + _b_chunk_start = nt_idx * BS + _b_eos = cutlass.Int32(_b_chunk_start + BS) + if IS_VARLEN: + _b_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + _b_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + _b_bos = cutlass.Int32(mCuSeqlens[_b_seq_id]) + _b_eos = cutlass.Int32(mCuSeqlens[_b_seq_id + 1]) + _b_chunk_start = _b_bos + _b_local * BS + + if warp_idx == 0: + _bcol_lo = lane_id * 2 + _bcol_hi = _bcol_lo + 1 + if _bcol_lo < BS: + _bt_lo = _b_chunk_start + _bcol_lo + _bt_hi = _b_chunk_start + _bcol_hi + if IS_VARLEN: + if _bt_lo < _b_eos: + sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx]) + else: + sBeta[_bcol_lo] = cutlass.Float32(0.0) + if _bt_hi < _b_eos: + sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx]) + else: + sBeta[_bcol_hi] = cutlass.Float32(0.0) + else: + sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx]) + sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx]) + cute.arch.barrier() + + # ===== Stage 1: FP32 forward solve, two diagonal blocks per warp ===== + if warp_idx == 0: + _invert_diag_forward_fp32(sAkk, lane_id // 16, lane_id) + if warp_idx == 1: + _invert_diag_forward_fp32(sAkk, 2 + lane_id // 16, lane_id) + + cute.arch.barrier() + + # ===== Stage 2: First batch — Ai10, Ai21, Ai32 ===== + # C→A chain: pack FP32 C-accum → bf16x2 A-operand via _pack_C_to_A + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 1, 0, 1, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 0, 0, a0, a1, a2, a3, lane_id) + _store_neg_C(sAkk, 1, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + if warp_idx == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 2, 2, 1, 2, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 1, 1, a0, a1, a2, a3, lane_id) + _store_neg_C(sAkk, 2, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + if warp_idx == 2: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 3, 3, 2, 3, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 2, 2, a0, a1, a2, a3, lane_id) + _store_neg_C(sAkk, 3, 2, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 3: Second batch — Ai20, Ai31 (warp pairs via sTemp) ===== + _z = cutlass.Float32(0.0) + t0 = _z + t1 = _z + t2 = _z + t3 = _z + t4 = _z + t5 = _z + t6 = _z + t7 = _z + + # --- Ai20 = -Ai22 @ (Akk20 @ Ai00 + Akk21 @ Ai10) --- + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 2, 0, 0, lane_id) + + if warp_idx == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 2, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + # --- Ai31 = -Ai33 @ (Akk31 @ Ai11 + Akk32 @ Ai21) --- + if warp_idx == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 3, 1, 1, lane_id) + + if warp_idx == 3: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 1, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + cute.arch.barrier() + + # Warp 0: accumulate T1+T2, pack→B, multiply by Ai22 + if warp_idx == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 2, 2, b0, b1, b2, b3, lane_id) + _store_neg_C(sAkk, 2, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + # Warp 1: accumulate T1'+T2', pack→B, multiply by Ai33 + if warp_idx == 1: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 3, 3, b0, b1, b2, b3, lane_id) + _store_neg_C(sAkk, 3, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 4: Third batch — Ai30 ===== + t0 = _z + t1 = _z + t2 = _z + t3 = _z + t4 = _z + t5 = _z + t6 = _z + t7 = _z + + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 3, 0, 0, lane_id) + + if warp_idx == 1: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 3, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + if warp_idx == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 0, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + cute.arch.barrier() + + # Warp 0: accumulate all three, pack→B, multiply by Ai33 + if warp_idx == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 3, 3, b0, b1, b2, b3, lane_id) + _store_neg_C(sAkk, 3, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 5: Store packed bf16x2 sAkk → global (b32 with direct bit-mask) ===== + vl_chunk_start = nt_idx * BS + vl_eos = cutlass.Int32(vl_chunk_start + BS) + if IS_VARLEN: + vl_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + vl_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + vl_bos = cutlass.Int32(mCuSeqlens[vl_seq_id]) + vl_eos = cutlass.Int32(mCuSeqlens[vl_seq_id + 1]) + vl_chunk_start = vl_bos + vl_local * BS + + row_start = warp_idx * SB + + # PDL: hint downstream K4 to pre-launch so its setup work overlaps with our + # gmem writes below. K4 has griddepcontrol_wait after its setup but before + # reading any gmem (TMA loads / gk_last_exp), so ordering is correct. + cute.arch.griddepcontrol_launch_dependents() + + for ri in cutlass.range_constexpr(SB): + row = row_start + ri + pair = lane_id + if pair < BS // 2: + packed = cutlass.Float32(sAkk[row, pair]) + masked = _mask_packed_ltri(packed, cutlass.Int32(row), cutlass.Int32(pair)) + + _beta_lo = cutlass.Float32(sBeta[2 * pair]) + _beta_hi = cutlass.Float32(sBeta[2 * pair + 1]) + _lo_f32 = _unpack_bf16x2_lo(masked) * _beta_lo + _hi_f32 = _unpack_bf16x2_hi(masked) * _beta_hi + final = _pack_bf16x2(_lo_f32, _hi_f32) + + t_row = vl_chunk_start + row + if IS_VARLEN: + if t_row < vl_eos: + mOut[b_idx, t_row, h_idx, pair] = final + else: + mOut[b_idx, t_row, h_idx, pair] = final + + +# =========================================================================== +# Host JIT function +# =========================================================================== +@cute.jit +def akk_inv_host( + A_in: cute.Tensor, + A_out: cute.Tensor, + Beta_in: cute.Tensor, + # B / NT / T_VAL are runtime scalars (shape-independent compile): + # baking them recompiled the kernel per batch shape — at eval traffic + # (unique total T per prefill batch) ~100 s of JIT per batch. Only H + # and IS_VARLEN remain compile-time specializers. + B: cutlass.Int32, + NT: cutlass.Int32, + H: cutlass.Constexpr[int], + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + T_VAL: cutlass.Int32, + # Launch stream — runtime argument; launching on the DSL default stream + # races with the executor's non-blocking execution stream. + stream: cuda.CUstream, +): + # BF16 input: view as FP32 (packed bf16x2). Each FP32 element = 2 BF16. + # Original BF16 shape: [B, T, H, BS] with stride per-element in BF16. + # As FP32 packed bf16x2: shape [BS, BS//2, H, dim3, B] + # The row stride in BF16 units is H*BS. In FP32 (bf16x2) units = H*BS/2. + # The col stride in BF16 is 1. In bf16x2 pairs: pair stride = 1 (adjacent pairs in memory). + _dim3_size = IS_VARLEN * T_VAL + (1 - IS_VARLEN) * (T_VAL // BS) + _dim3_stride_bf16 = IS_VARLEN * (H * BS) + (1 - IS_VARLEN) * (BS * H * BS) + # All strides in FP32 units (each FP32 = 2 BF16): + _row_stride = H * BS // 2 # row stride (BF16 row stride = H*BS, /2 for bf16x2) + _col_stride = 1 # adjacent bf16x2 pairs + _h_stride = BS // 2 # head stride (BF16 = BS, /2 for bf16x2) + _dim3_stride = _dim3_stride_bf16 // 2 + # T_VAL is a runtime Int32, so this stride is dynamic. Mathematically it is + # T_VAL * (H * BS // 2) = T_VAL * 32 * H, i.e. always a multiple of 4 fp32 + # elements (16 B). The IR verifier can't derive that on its own, and without + # it the b_idx slice offset breaks the 128-bit alignment proof required by + # the cp.async G2S atom below (ICE: "src ptr alignment (32 bits) does not + # meet requirement (128 bits)"). Assert divisibility explicitly. + _batch_stride = cute.assume(T_VAL * (H * BS // 2), divby=H * BS // 2) + + view_layout = cute.make_layout( + (BS, BS // 2, H, _dim3_size, B), + stride=(_row_stride, _col_stride, _h_stride, _dim3_stride, _batch_stride), + ) + gA_view = cute.make_tensor(A_in.iterator, view_layout) + + # sAkk: packed bf16x2, logical shape (64, 32) with stride 36 for bank-conflict-free access + akk_smem_2d = cute.make_layout((BS, BS // 2), stride=(AKK_STRIDE, 1)) + + # cp.async G→S copy: 128-bit (4×fp32 = 8 bf16) vectorised + copy_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + cutlass.Float32, + num_bits_per_copy=128, + ) + # Source layout: BS rows × BS//2 pairs = 64×32 + # 128 threads, each copies 4 FP32 (128 bits) per iteration + # Need 64*32/128/4 = 4 iterations → but (16,8) × (1,4) = 128 threads × 4 vals = 512 per iter + # 64×32 = 2048 elements / 512 = 4 iterations. But SMEM is 64×36, need to handle padding. + # Actually: the copy maps source shape to SMEM shape. + # Source is (BS, BS//2) = (64, 32), SMEM is (BS, AKK_STRIDE) = (64, 36). + # The copy handles the actual data region (64, 32); padding cols 32-35 stay zero. + g2s_copy = cute.make_tiled_copy_tv( + copy_atom, + thr_layout=cute.make_layout((16, 8), stride=(8, 1)), + val_layout=cute.make_layout((1, 4)), + ) + + # sTemp layout (non-swizzled) + temp_layout = cute.make_layout( + (SB, TEMP_COLS, NUM_TEMPS), stride=(TEMP_COLS, 1, SB * TEMP_COLS) + ) + + smem_bytes = BS * AKK_STRIDE * 4 + SB * TEMP_COLS * NUM_TEMPS * 4 + BS * 4 + 256 + + out_layout = cute.make_layout( + (B, T_VAL, H, BS // 2), stride=(T_VAL * H * BS // 2, H * BS // 2, BS // 2, 1) + ) + out_view = cute.make_tensor(A_out.iterator, out_layout) + + akk_inv_kernel( + g2s_copy, + gA_view, + out_view, + Beta_in, + akk_smem_2d, + temp_layout, + NT, + H, + mCuSeqlens, + mChunkIndices, + IS_VARLEN, + ).launch( + grid=(H, NT, B), + block=(THREADS, 1, 1), + smem=smem_bytes, + stream=stream, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/fused_k123.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/fused_k123.py new file mode 100644 index 000000000000..9f89fbc82ca6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/fused_k123.py @@ -0,0 +1,2357 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Persistent Fused K1+K2+K3 Kernel for KDA. + +Fuses gate activation + cumsum + scaling (K1), intra sub-chunk Aqk/Akk (K2), +and inter sub-chunk solve + merged inverse (K3) into a single persistent kernel. + +Grid: (grid_blocks, 1, 1) — one persistent block per usable SM, each loops over work units + Total work units = (NT/4) * H * B, distributed round-robin across SMs + Block i processes work units i, i+grid_blocks, i+2*grid_blocks, ... +Block: 992 threads (31 warps), warp-specialized with setmaxnreg: + Warps 0-15: TMA+K1 fused (8×2, vec2, prefetch pipeline) – 4 WGs, 56 regs + Warps 16-26: K2 MMA compute (10 active + warp 26 as TMA producer) – 72 regs + Warps 27-30: Store/Inversion warps – 24 regs + +Pipeline (single for_generate, warp groups separated by if-blocks): + per work unit: + Warps 0-15: prefetch chunk 0→stage 0 (warp 0), then loop: + TMA next chunk (warp 0), wait cur chunk, K1 compute, arrive(k1_done) + Warps 16-26: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse) + Warps 27-30: wait(mma_done), store sAqk/sAkk→GMEM, arrive(store_done) + All warp-group invariants are computed inside each group's if-block (not hoisted) + to eliminate cross-group register pressure — same budget as the _all version. + Mbarrier phases self-reset after 4 iterations (2 stages × 2 phases). + +Mbarriers: + tma_mbars[2]: count=1, warp 0 lane 0 → K1+MMA wait for TMA data + stage_reuse_mbars[2]: count=832, K1(16)+MMA(10 warps) → TMA waits before reuse + k1_done_mbars[2]: count=512, K1(16 warps) → MMA waits for g_cumsum ready + mma_done_mbars[2]: count=320, MMA(10 warps) → Store waits for sAqk/sAkk ready + store_done_mbars[2]: count=128, Store(4 warps) → MMA waits for sAqk/sAkk stage free + +SMEM: ~215KB (q+k+g × [64,128] bf16 × 2 stages + g_cumsum [64,136] fp32 × 2 stages + + sPartialLast [8,132] fp32 + sAqk [16,168,2] bf16 + + sAkk [64,72,2] fp32 block-transposed upper-tri layout) + +Inputs: + g [B,T,H,K] bf16 raw gate + k [B,T,H,K] bf16 normalized by caller + q [B,T,H,K] bf16 normalized by caller + A_log [H] fp32 per-head log decay + beta [B,T,H] bf16/fp32, activated in-kernel when requested + scale fp32 1/sqrt(K) + +Outputs (g_cumsum stays in SMEM, not written to GMEM): + k_scaled [B,T,H,K] bf16 + q_scaled [B,T,H,K] bf16 + kg [B,T,H,K] bf16 + gk_last_exp[B,NT,H,K] fp32 + A_qk [B,T,H,BT] bf16 full merged (diagonal + off-diagonal) + A_kk [B,T,H,BT] fp32 block-transposed upper-tri (input to akk_inv) +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import for_generate, yield_out +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op + +BT = 64 +BC = 16 +K_DIM = 128 +K_PAD = 8 +K_STRIDE = K_DIM + K_PAD # 136, padded row stride to avoid bank conflicts +CHUNKS_PER_BLOCK = 4 +NUM_SMS = 148 # Default persistent grid (B200); callers pass the device's usable SM count + +NUM_K1_TMA_WARPS = 16 # Warps 0-15: K1 compute (4 warpgroups, 8×2) -- TMA offloaded +NUM_MMA_WARPS = 11 # Warps 16-26: MMA (10 active + 1 TMA producer, dropped idle warp 27) +NUM_MMA_ACTIVE = 10 # mma_warp 0..9: actual MMA work +TMA_WARP_ID = NUM_K1_TMA_WARPS + NUM_MMA_ACTIVE # warp 26 = dedicated TMA producer +NUM_STORE_WARPS = 4 # Warps 27-30: Store/Inversion (1 warpgroup) +NUM_WARPS = NUM_K1_TMA_WARPS + NUM_MMA_WARPS + NUM_STORE_WARPS # 31 +THREADS = NUM_WARPS * 32 # 992 + +NUM_SUB_CHUNKS = BT // BC # 4 +NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10 lower-tri tiles +MMA_K_TILE = 16 +NUM_MMA_K_TILES = K_DIM // MMA_K_TILE # 8 (bf16 m16n8k16) +AQK_TILE_PAD = 8 +AQK_TILE_STRIDE = BT + AQK_TILE_PAD # 72 — sAqk now 64x72 row-major (same shape as sAkk) + +AKK_PAD = 8 +AKK_STRIDE = BT + AKK_PAD # 72 + +# For fused akk_inv: sAkk_pkd viewed as fp32 packed bf16x2, stride = BT/2 + small pad +AKK_PKD_PAD = 4 +AKK_PKD_STRIDE = BT // 2 + AKK_PKD_PAD # 36 fp32 units = 72 bf16 elements per row +AKK_TEMP_PAD = 8 +AKK_TEMP_COLS = 16 + AKK_TEMP_PAD # 24 +AKK_TEMP_BUFS = 2 + +K1_ROW_GROUPS = 8 +K1_COL_GROUPS = 2 +ROWS_PER_K1_WARP = BT // K1_ROW_GROUPS # 8 +K1_COLS_PER_WARP = K_DIM // K1_COL_GROUPS # 64 +ROWS_PER_STORE_WARP = BT // NUM_STORE_WARPS # 16 + +VEC = K1_COLS_PER_WARP // 32 # 2 +K_VEC = K_DIM // VEC # 64 +NUM_STAGES = 2 +PARTIAL_COLS = K_DIM + 4 # 132 +PARTIAL_COLS_PER_WARP = K_DIM // NUM_K1_TMA_WARPS # 8 + +_TILE_IQ = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3] +_TILE_IK = [0, 0, 1, 0, 1, 2, 0, 1, 2, 3] + +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +RCP_LN2 = LOG2E + + +@dsl_user_op +def k1_internal_barrier(*, loc=None, ip=None): + """Named barrier for K1+TMA warps (0-15, 512 threads). barrier_id=2.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 2, 512; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def pack_bf16x2_f32(hi_f32, lo_f32, *, loc=None, ip=None): + """Pack two fp32 values into a bf16x2 u32 register. + Returns u32 with bf16(hi) in bits [31:16] and bf16(lo) in bits [15:0]. + PTX: cvt.rn.bf16x2.f32 d, a, b -> d[31:16]=bf16(a), d[15:0]=bf16(b) + """ + result = llvm.inline_asm( + T.i32(), + [hi_f32.ir_value(loc=loc, ip=ip), lo_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@dsl_user_op +def mma_bf16_m16n8k16( + a0, + a1, + a2, + a3, # 4 u32 (bf16x2 packed) A operands + b0, + b1, # 2 u32 (bf16x2 packed) B operands + c0, + c1, + c2, + c3, # 4 fp32 accumulators + *, + loc=None, + ip=None, +): + """bf16 MMA with fp32 accumulator, shape m16n8k16. + D_fp32 = A_bf16 * B_bf16 + C_fp32 + """ + # a/b already i32 (from pack_bf16x2_f32) -> no bitcast needed + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0.ir_value(loc=loc, ip=ip), + a1.ir_value(loc=loc, ip=ip), + a2.ir_value(loc=loc, ip=ip), + a3.ir_value(loc=loc, ip=ip), + b0.ir_value(loc=loc, ip=ip), + b1.ir_value(loc=loc, ip=ip), + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +SHFL_W8_CLAMP = 0x1800 + + +@dsl_user_op +def fast_rcp(x, *, loc=None, ip=None): + """Hardware fast reciprocal: rcp.approx.ftz.f32 (~2 cycles vs ~20 for div).""" + result = llvm.inline_asm( + T.f32(), + [x.ir_value(loc=loc, ip=ip)], + "rcp.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(result) + + +# ============================================================ +# Akk inverse helpers (fused from Akk_inverse_lower_triangle_bf16.py) +# ============================================================ +@dsl_user_op +def store_internal_barrier(*, loc=None, ip=None): + """Named barrier for Store warps (4 warps, 128 threads). barrier_id=3.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 3, 128; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ak_pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None): + """cvt.rn.bf16x2.f32 -- pack two fp32 into bf16x2 (as fp32 bitcast view).""" + result = llvm.inline_asm( + T.i32(), + [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_movmatrix_trans(src, *, loc=None, ip=None): + src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [src_b], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_mask_packed_ltri(packed, row, pair, *, loc=None, ip=None): + """Mask packed bf16x2 to lower-tri: zero elements where row < col.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)], + """{ + .reg .b32 %c0, %c1, %mlo, %mhi, %mask; + .reg .pred %p0, %p1; + shl.b32 %c0, $3, 1; + add.u32 %c1, %c0, 1; + setp.ge.s32 %p0, $2, %c0; + setp.ge.s32 %p1, $2, %c1; + selp.b32 %mlo, 0xFFFF, 0, %p0; + selp.b32 %mhi, 0xFFFF0000, 0, %p1; + or.b32 %mask, %mlo, %mhi; + and.b32 $0, $1, %mask; + }""", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_unpack_bf16x2_lo(packed, *, loc=None, ip=None): + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + "shl.b32 $0, $1, 16;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_unpack_bf16x2_hi(packed, *, loc=None, ip=None): + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + "and.b32 $0, $1, 0xFFFF0000;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_mma(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """BF16 MMA m16n8k16, args are Float32 (bitcast-viewed as i32 for bf16x2).""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def _ak_invert_diag_neumann(sAkk, block_idx, lane_id, *, loc=None, ip=None): + """Invert 16x16 diag block via Neumann series (I+L)^-1 = (I-L)(I+L^2)(I+L^4)(I+L^8). + Reads/writes sAkk as fp32 packed bf16x2 (stride 36). Each block handled by 1 warp.""" + r_off = block_idx * 16 + c_off = block_idx * 8 # 16 cols = 8 pairs + gid = lane_id // 4 + tid = lane_id % 4 + + packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid]) + packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid]) + packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid]) + packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid]) + + A_f0 = _ak_unpack_bf16x2_lo(packed0) + A_f1 = _ak_unpack_bf16x2_hi(packed0) + A_f2 = _ak_unpack_bf16x2_lo(packed1) + A_f3 = _ak_unpack_bf16x2_hi(packed1) + A_f4 = _ak_unpack_bf16x2_lo(packed2) + A_f5 = _ak_unpack_bf16x2_hi(packed2) + A_f6 = _ak_unpack_bf16x2_lo(packed3) + A_f7 = _ak_unpack_bf16x2_hi(packed3) + + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid) + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + 1 + ) + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid) + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + 1 + ) + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + ) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + 1 + ) + + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + L_a0 = _ak_pack_bf16x2(L_f0, L_f1) + L_a1 = _ak_pack_bf16x2(L_f2, L_f3) + L_a2 = _ak_pack_bf16x2(L_f4, L_f5) + L_a3 = _ak_pack_bf16x2(L_f6, L_f7) + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 1: L^2, INV += INV*L^2 + L_b0 = _ak_movmatrix_trans(L_a0) + L_b1 = _ak_movmatrix_trans(L_a1) + L_b2 = _ak_movmatrix_trans(L_a2) + L_b3 = _ak_movmatrix_trans(L_a3) + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf) + Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7) + Lp_b0 = _ak_movmatrix_trans(Lp_a0) + Lp_b1 = _ak_movmatrix_trans(Lp_a1) + Lp_b2 = _ak_movmatrix_trans(Lp_a2) + Lp_b3 = _ak_movmatrix_trans(Lp_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 2: L^4, INV += INV*L^4 + L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7) + L4_b0 = _ak_movmatrix_trans(L4_a0) + L4_b1 = _ak_movmatrix_trans(L4_a1) + L4_b2 = _ak_movmatrix_trans(L4_a2) + L4_b3 = _ak_movmatrix_trans(L4_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 3: L^8, INV += INV*L^8 + L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7) + L8_b0 = _ak_movmatrix_trans(L8_a0) + L8_b1 = _ak_movmatrix_trans(L8_a1) + L8_b2 = _ak_movmatrix_trans(L8_a2) + L8_b3 = _ak_movmatrix_trans(L8_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _ak_invert_diag_neumann_inreg( + sAkk, + block_idx, + lane_id, + raw_f0, + raw_f1, + raw_f2, + raw_f3, + raw_f4, + raw_f5, + raw_f6, + raw_f7, + *, + loc=None, + ip=None, +): + """In-register variant: invert 16x16 diag block from K2 acc fragment values + (no SMEM read). Applies I+L mask, runs Neumann iterations, writes INV to sAkk. + + raw_f0..raw_f7 are K2 fp32 acc*beta values in C-fragment layout for m16n8k16: + raw_f0,raw_f1 = (gid, 2*tid), (gid, 2*tid+1) + raw_f2,raw_f3 = (gid+8, 2*tid), (gid+8, 2*tid+1) + raw_f4,raw_f5 = (gid, 2*tid+8), (gid, 2*tid+9) + raw_f6,raw_f7 = (gid+8, 2*tid+8), (gid+8, 2*tid+9) + """ + r_off = block_idx * 16 + c_off = block_idx * 8 + gid = lane_id // 4 + tid = lane_id % 4 + + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + + # Identity pattern (matches the SMEM-read variant's I_f computation) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + ) # always 0 + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + 1 + ) # always 0 + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + ) # always 0 + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + 1 + ) # always 0 + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + ) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + 1 + ) + + # Apply I+L mask to raw K2 values: + # diag (row==col): replace with 1 + # strict upper (rowcol within block): keep raw value (= L value) + m_gt_a = cutlass.Float32(gid > 2 * tid) # for cols 2*tid (with row=gid) + m_gt_b = cutlass.Float32(gid > 2 * tid + 1) # for cols 2*tid+1 + # (gid, 2*tid) and (gid, 2*tid+1): mask conditional + A_f0 = m_gt_a * raw_f0 + I_f0 + A_f1 = m_gt_b * raw_f1 + I_f1 + # (gid+8, 2*tid) and (gid+8, 2*tid+1): always strict lower (gid+8 > 2*tid+1 always) + A_f2 = raw_f2 + A_f3 = raw_f3 + # (gid, 2*tid+8) and (gid, 2*tid+9): always strict upper (gid <= 7 < 8 <= 2*tid+8) + A_f4 = _zero + A_f5 = _zero + # (gid+8, 2*tid+8): row-col offset = gid+8 - (2*tid+8) = gid - 2*tid -> same as A_f0 mask + # (gid+8, 2*tid+9): similar -> same as A_f1 mask + A_f6 = m_gt_a * raw_f6 + I_f6 + A_f7 = m_gt_b * raw_f7 + I_f7 + + # L = A - I, INV = I - L (first two Neumann terms) + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + L_a0 = _ak_pack_bf16x2(L_f0, L_f1) + L_a1 = _ak_pack_bf16x2(L_f2, L_f3) + L_a2 = _ak_pack_bf16x2(L_f4, L_f5) + L_a3 = _ak_pack_bf16x2(L_f6, L_f7) + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 1: L^2, INV += INV*L^2 + L_b0 = _ak_movmatrix_trans(L_a0) + L_b1 = _ak_movmatrix_trans(L_a1) + L_b2 = _ak_movmatrix_trans(L_a2) + L_b3 = _ak_movmatrix_trans(L_a3) + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf) + Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7) + Lp_b0 = _ak_movmatrix_trans(Lp_a0) + Lp_b1 = _ak_movmatrix_trans(Lp_a1) + Lp_b2 = _ak_movmatrix_trans(Lp_a2) + Lp_b3 = _ak_movmatrix_trans(Lp_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 2: L^4, INV += INV*L^4 + L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7) + L4_b0 = _ak_movmatrix_trans(L4_a0) + L4_b1 = _ak_movmatrix_trans(L4_a1) + L4_b2 = _ak_movmatrix_trans(L4_a2) + L4_b3 = _ak_movmatrix_trans(L4_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 3: L^8, INV += INV*L^8 + L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7) + L8_b0 = _ak_movmatrix_trans(L8_a0) + L8_b1 = _ak_movmatrix_trans(L8_a1) + L8_b2 = _ak_movmatrix_trans(L8_a2) + L8_b3 = _ak_movmatrix_trans(L8_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _ak_matmul_AB(sAkk, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + rB = br_B * 16 + cB = bc_B * 8 + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _ak_movmatrix_trans(bA0) + b1 = _ak_movmatrix_trans(bA1) + b2 = _ak_movmatrix_trans(bA2) + b3 = _ak_movmatrix_trans(bA3) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_chain_mma_B(sAkk, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 8 + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _ak_movmatrix_trans(bA0) + b1 = _ak_movmatrix_trans(bA1) + b2 = _ak_movmatrix_trans(bA2) + b3 = _ak_movmatrix_trans(bA3) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_chain_mma_A(sAkk, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_store_neg_C(sAkk, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 8 + sAkk[r + gid, c + tid] = _ak_pack_bf16x2(-c0, -c1) + sAkk[r + gid + 8, c + tid] = _ak_pack_bf16x2(-c2, -c3) + sAkk[r + gid, c + 4 + tid] = _ak_pack_bf16x2(-c4, -c5) + sAkk[r + gid + 8, c + 4 + tid] = _ak_pack_bf16x2(-c6, -c7) + + +@dsl_user_op +def _ak_pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _ak_pack_bf16x2(c0, c1) + a1 = _ak_pack_bf16x2(c2, c3) + a2 = _ak_pack_bf16x2(c4, c5) + a3 = _ak_pack_bf16x2(c6, c7) + return a0, a1, a2, a3 + + +@dsl_user_op +def _ak_pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _ak_pack_bf16x2(c0, c1) + a1 = _ak_pack_bf16x2(c2, c3) + a2 = _ak_pack_bf16x2(c4, c5) + a3 = _ak_pack_bf16x2(c6, c7) + b0 = _ak_movmatrix_trans(a0) + b1 = _ak_movmatrix_trans(a1) + b2 = _ak_movmatrix_trans(a2) + b3 = _ak_movmatrix_trans(a3) + return b0, b1, b2, b3 + + +@dsl_user_op +def _ak_store_C_temp(sT, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _ak_load_C_temp(sT, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +@dsl_user_op +def opaque_zero_from_work_id(*, loc=None, ip=None): + """ + Return 0 via opaque side-effectful asm (no inputs needed). + + Because has_side_effects=True, MLIR LICM treats this as having memory + effects and will NOT hoist it outside the for_generate loop. Any value + computed from the result (_oz) therefore appears loop-variant to LICM, + preventing get_slice() and scalar layout-invariant computations from being + hoisted to the kernel prologue. This keeps prologue register pressure < 64 + and eliminates the 440-byte stack frame that caused ~300-cycle L2 LDL + penalties per iteration. + """ + result = llvm.inline_asm( + T.i32(), + [], + "mov.b32 $0, 0;", # output = 0 (opaque to compiler constant-folding) + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@cute.kernel +def fused_kernel123( + tma_atom_Q: cute.CopyAtom, + tma_tensor_Q: cute.Tensor, + tma_atom_K: cute.CopyAtom, + tma_tensor_K: cute.Tensor, + tma_atom_G: cute.CopyAtom, + tma_tensor_G: cute.Tensor, + mA_log: cute.Tensor, + mBeta: cute.Tensor, + mBetaActivated: cute.Tensor, + scale: cutlass.Float32, + mKscaled: cute.Tensor, + mKg: cute.Tensor, + mQscaled: cute.Tensor, + mGkLast: cute.Tensor, + mAqk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAkk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAqk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + mAkk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_layout, + g_smem_layout, + g_cumsum_layout, + num_chunks: cutlass.Int32, # runtime — shape-independent compile + num_heads: int, # baked (single variant in practice) + batch_size: cutlass.Int32, # runtime — shape-independent compile + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + mDtBias: cute.Tensor, + lower_bound: cutlass.Float32, + HAS_BIAS: cutlass.Constexpr[int], + USE_SAFE_GATE: cutlass.Constexpr[int], + valid_tokens: cutlass.Int32, + USE_BETA_SIGMOID: cutlass.Constexpr[int], + VARLEN_PURE: cutlass.Constexpr[int] = 0, + GRID_BLOCKS: cutlass.Constexpr[int] = NUM_SMS, +): + block_id, _, _ = cute.arch.block_idx() + tidx = cute.arch.thread_idx()[0] + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_id = tidx % 32 + + total_cgs_per_head = cutlass.Int32(0) + cgs_per_head = cutlass.Int32(0) + total_cgs = cutlass.Int32(0) + if IS_VARLEN: + total_cgs_per_head = (num_chunks + CHUNKS_PER_BLOCK - 1) // CHUNKS_PER_BLOCK + total_cgs = total_cgs_per_head * num_heads + else: + cgs_per_head = num_chunks // CHUNKS_PER_BLOCK + total_cgs = cgs_per_head * num_heads * batch_size + + # ===================================================================== + # SMEM allocation + # ===================================================================== + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor( + cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner + ) + sK = smem.allocate_tensor( + cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner + ) + sG = smem.allocate_tensor(cutlass.BFloat16, g_smem_layout, 128) + sGcum = smem.allocate_tensor(cutlass.Float32, g_cumsum_layout, 128) + beta_layout = cute.make_layout((BT, NUM_STAGES), stride=(1, BT)) + sBeta = smem.allocate_tensor(cutlass.BFloat16, beta_layout, 128) + partial_last_layout = cute.make_layout((K1_ROW_GROUPS, PARTIAL_COLS), stride=(PARTIAL_COLS, 1)) + sPartialLast = smem.allocate_tensor(cutlass.Float32, partial_last_layout, 128) + + # sAqk: 64x72 row-major (BT rows, BT+pad cols), same shape as sAkk. + # Each sub-tile (i_q, i_k) sits at SMEM rows [i_q*BC, (i_q+1)*BC) cols + # [i_k*BC, (i_k+1)*BC) — directly mirrors the 64x64 attention matrix. + aqk_tile_layout = cute.make_layout( + (BT, AQK_TILE_STRIDE, NUM_STAGES), stride=(AQK_TILE_STRIDE, 1, BT * AQK_TILE_STRIDE) + ) + sAqk = smem.allocate_tensor(cutlass.BFloat16, aqk_tile_layout, 128) + + akk_tile_layout = cute.make_layout( + (BT, AKK_STRIDE, NUM_STAGES), stride=(AKK_STRIDE, 1, BT * AKK_STRIDE) + ) + sAkk = smem.allocate_tensor(cutlass.BFloat16, akk_tile_layout, 128) + # sAkk_pkd / sTemp removed: akk_inv runs as a separate kernel call (chained back-to-back). + + # ===================================================================== + # Mbarrier allocation & init + # ===================================================================== + tma_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + stage_reuse_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + k1_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + mma_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + store_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + + bytes_per_stage = BT * K_DIM * 2 * 3 + + if tidx == 0: + for s in range(NUM_STAGES): + cute.arch.mbarrier_init(tma_mbars + s, 1) + # Both K1 Pass 2b and MMA read Q/K before TMA can reuse a stage. + # The dedicated TMA warp does not arrive on either barrier. + cute.arch.mbarrier_init( + stage_reuse_mbars + s, (NUM_MMA_WARPS - 1 + NUM_K1_TMA_WARPS) * 32 + ) + cute.arch.mbarrier_init(k1_done_mbars + s, NUM_K1_TMA_WARPS * 32) + cute.arch.mbarrier_init(mma_done_mbars + s, (NUM_MMA_WARPS - 1) * 32) + cute.arch.mbarrier_init(store_done_mbars + s, NUM_STORE_WARPS * 32) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + # ===================================================================== + # SMEM init: zero out sAqk and sAkk valid 64x64 region (cols 64..71 are + # padding for SMEM bank-conflict avoidance, never read or written by MMA + # or store warps). Required for downstream row-major store optimizations + # — positions outside MMA-written sub-tiles stay at 0. + # + # Cooperative pattern (31 warps × 32 lanes = 992 threads): + # - Rows strided by warp count: warp w owns rows w, w+31, w+62 (< BT) + # - Each lane owns 2 contiguous bf16 cols (lane*2, lane*2+1) + # - Adjacent (lane*2, lane*2+1) bf16 pairs are 4-byte aligned → + # ptxas should fuse into STS.32 wide stores. + # ===================================================================== + _warp_id_in_cta = tidx >> 5 # tidx // 32, range 0..30 + _lane_id_warp = tidx & 31 # tidx % 32, range 0..31 + _col_lo = _lane_id_warp * 2 # this lane owns cols [_col_lo, _col_lo+1] + _col_hi = _col_lo + 1 + for _s in cutlass.range_constexpr(NUM_STAGES): + for _ri in cutlass.range_constexpr((BT + NUM_WARPS - 1) // NUM_WARPS): + _row = _warp_id_in_cta + _ri * NUM_WARPS + if _row < BT: + sAqk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) + sAqk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) + sAkk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) + sAkk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) + cute.arch.barrier() + + # ===================================================================== + # Pre-arrive before the first TMA load and MMA store + # stage_reuse_mbars: pre-arrive all 16 K1 and 10 MMA reader warps + # store_done_mbars: MMA waits before Store arrives → pre-arrive first 4 MMA warps + # ===================================================================== + if ( + warp_idx >= NUM_K1_TMA_WARPS + and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS + and warp_idx != TMA_WARP_ID + ): + mma_warp_tmp = warp_idx - NUM_K1_TMA_WARPS + for s in range(NUM_STAGES): + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + if mma_warp_tmp < NUM_STORE_WARPS: + cute.arch.mbarrier_arrive(store_done_mbars + s) + + if warp_idx < NUM_K1_TMA_WARPS: + for s in range(NUM_STAGES): + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + # ================================================================= + # Persistent outer loop. Single for_generate at top level (required). + # Opaque asm barrier on work_id prevents MLIR LICM from hoisting + # get_slice() and scalar layout invariants to the kernel prologue, + # keeping register pressure < 64 and eliminating prologue spill. + # ================================================================= + for work_id in for_generate(block_id, total_cgs, GRID_BLOCKS): + i_cg = cutlass.Int32(0) + i_h = cutlass.Int32(0) + i_b = cutlass.Int32(0) + chunk_base = cutlass.Int32(0) + if IS_VARLEN: + i_cg = work_id % total_cgs_per_head + i_h = work_id // total_cgs_per_head + i_b = cutlass.Int32(0) + chunk_base = i_cg * CHUNKS_PER_BLOCK + else: + i_cg = work_id % cgs_per_head + i_h = (work_id // cgs_per_head) % num_heads + i_b = work_id // (cgs_per_head * num_heads) + chunk_base = i_cg * CHUNKS_PER_BLOCK + + # Anti-LICM barrier: _oz is always 0 but appears to depend on work_id. + # Because this asm has side_effects=True, it stays inside the loop. + # Any value computed from _oz/_lane/_warp is also loop-variant from + # LICM's perspective → get_slice() and scalar invariants stay in-loop. + _oz = opaque_zero_from_work_id() + _lane = lane_id + _oz + _warp = warp_idx + _oz + + # ============================================================= + # Warps 0-15: Fused TMA + K1 + # ============================================================= + if warp_idx < NUM_K1_TMA_WARPS: + # Warp-layout invariants (scope-local → no cross-group register spill) + k1_warp = _warp + warp_row_group = k1_warp % K1_ROW_GROUPS + warp_col_group = k1_warp // K1_ROW_GROUPS + k1_row_start = warp_row_group * ROWS_PER_K1_WARP + col_base = warp_col_group * K1_COLS_PER_WARP + _lane * VEC + col_vec_idx = warp_col_group * (K1_COLS_PER_WARP // VEC) + _lane + cumsum_scale = cutlass.Float32(RCP_LN2) + thr_copy_k1 = tiled_copy_qk_k1.get_slice(_lane) + + rAcc = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rPrefix = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rGkLast = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rKsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rQsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rKgOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rGkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + + # exp_A depends on i_h (changes per work unit) + exp_A = cute.exp(mA_log[i_h], fastmath=True) + + # Load dt_bias per (head, col) — broadcast across all rows + rBias = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + if HAS_BIAS: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = mDtBias[i_h, col_base + vi] + else: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = cutlass.Float32(0.0) + + # 3D TMA head slices (fixed for this work unit's head) + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + cur_stage = chunk_iter % NUM_STAGES + cur_phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + ci_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + ci_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(tma_mbars + cur_stage, cur_phase) + + csG = sG[(None, None, cur_stage)] + csGcum = sGcum[(None, None, cur_stage)] + csQ = sQ[(None, None, cur_stage)] + csK = sK[(None, None, cur_stage)] + rGact = cute.make_rmem_tensor( + cute.make_layout((ROWS_PER_K1_WARP, VEC)), cutlass.Float32 + ) + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = cutlass.Float32(0.0) + + # Activate beta once per token/head and retain the bf16 + # result in this stage for K2. The activated GMEM scratch is + # consumed by akk_inv, beta's second pipeline consumer. + if k1_warp < 2: + beta_row = k1_warp * 32 + lane_id + beta_t = chunk_start + beta_row + beta_value = cutlass.Float32(0.0) + if IS_VARLEN: + if chunk_idx < num_chunks and beta_t < ci_eos: + beta_value = mBeta[i_b, beta_t, i_h].to(cutlass.Float32) + if USE_BETA_SIGMOID: + beta_value = fast_rcp( + cutlass.Float32(1.0) + + cute.exp2( + -beta_value * LOG2E, + fastmath=True, + ) + ) + mBetaActivated[i_b, beta_t, i_h] = beta_value.to(cutlass.BFloat16) + else: + if beta_t < valid_tokens: + beta_value = mBeta[i_b, beta_t, i_h].to(cutlass.Float32) + if USE_BETA_SIGMOID: + beta_value = fast_rcp( + cutlass.Float32(1.0) + + cute.exp2( + -beta_value * LOG2E, + fastmath=True, + ) + ) + if USE_BETA_SIGMOID: + # Eqlen padding is allocated storage, so publish + # explicit zeros for its invalid tail. Applying + # sigmoid to the padded zero logits would produce + # 0.5 and corrupt the recurrent state. + mBetaActivated[i_b, beta_t, i_h] = beta_value.to(cutlass.BFloat16) + sBeta[beta_row, cur_stage] = beta_value.to(cutlass.BFloat16) + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + c = col_base + vi + g_val = csG[row, c].to(cutlass.Float32) + if HAS_BIAS: + g_val = g_val + rBias[vi] + g_activated = cutlass.Float32(0.0) + if USE_SAFE_GATE: + sigmoid_g = fast_rcp( + cutlass.Float32(1.0) + + cute.exp2(-exp_A * g_val * LOG2E, fastmath=True) + ) + g_activated = lower_bound * sigmoid_g + else: + softplus_g = ( + cute.log2( + cutlass.Float32(1.0) + cute.exp2(g_val * LOG2E, fastmath=True), + fastmath=True, + ) + * LN2 + ) + g_activated = -exp_A * softplus_g + # Varlen: zero gate for out-of-bounds rows so cumsum + # stays flat beyond the last valid position. + # VARLEN_PURE=1 elides this at compile time — caller + # guarantees all seq lengths are multiples of BT so no + # chunk has OOB rows. + if IS_VARLEN and not VARLEN_PURE: + if chunk_start + row >= ci_eos: + g_activated = cutlass.Float32(0.0) + rGact[ri, vi] = g_activated + rAcc[vi] = rAcc[vi] + g_activated + + for vi in cutlass.range_constexpr(VEC): + sPartialLast[warp_row_group, col_base + vi] = rAcc[vi] + + k1_internal_barrier() + + prefix_col_start = k1_warp * PARTIAL_COLS_PER_WARP + row_in_prefix = lane_id % K1_ROW_GROUPS + col_in_group = lane_id // K1_ROW_GROUPS + + for j in cutlass.range_constexpr(PARTIAL_COLS_PER_WARP // 4): + col = prefix_col_start + j * 4 + col_in_group + val = cutlass.Float32(sPartialLast[row_in_prefix, col]) + tmp = cute.arch.shuffle_sync_up(val, 1, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 1: + val = val + tmp + tmp = cute.arch.shuffle_sync_up(val, 2, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 2: + val = val + tmp + tmp = cute.arch.shuffle_sync_up(val, 4, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 4: + val = val + tmp + sPartialLast[row_in_prefix, col] = val + + k1_internal_barrier() + + for vi in cutlass.range_constexpr(VEC): + rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi] + + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = cutlass.Float32(0.0) + if warp_row_group > 0: + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = sPartialLast[warp_row_group - 1, col_base + vi] + + # The next chunk's partial sums overwrite sPartialLast, so every + # K1 warp must have read its prefix and the chunk total first. + k1_internal_barrier() + + # ---- Pass 2a: ONLY cumsum + write csGcum (critical path, minimal work) ---- + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rPrefix[vi] + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rAcc[vi] + rGact[ri, vi] + csGcum[row, col_base + vi] = rAcc[vi] * cumsum_scale + + # Signal MMA early: csGcum is ready + cute.arch.mbarrier_arrive(k1_done_mbars + cur_stage) + + # ---- Pass 2b: recompute + write GMEM (overlaps with MMA, off critical path) ---- + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rPrefix[vi] + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + t = chunk_start + row + + sK_tile = cute.local_tile( + csK, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group) + ) + tCsK = thr_copy_k1.partition_S(sK_tile) + tCrK = cute.make_fragment_like(tCsK) + cute.copy(tiled_copy_qk_k1, tCsK, thr_copy_k1.retile(tCrK)) + + sQ_tile = cute.local_tile( + csQ, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group) + ) + tCsQ = thr_copy_k1.partition_S(sQ_tile) + tCrQ = cute.make_fragment_like(tCsQ) + cute.copy(tiled_copy_qk_k1, tCsQ, thr_copy_k1.retile(tCrQ)) + + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rAcc[vi] + rGact[ri, vi] + cs = rAcc[vi] * cumsum_scale + + k_val = tCrK[vi].to(cutlass.Float32) + q_val = tCrQ[vi].to(cutlass.Float32) + + exp2_cs = cute.exp2(cs, fastmath=True) + gk_last_cs = rGkLast[vi] * cumsum_scale + exp2_kg = cute.exp2(gk_last_cs - cs, fastmath=True) + + rKsOut[vi] = (k_val * exp2_cs).to(cutlass.BFloat16) + rQsOut[vi] = (q_val * exp2_cs * scale).to(cutlass.BFloat16) + rKgOut[vi] = (k_val * exp2_kg).to(cutlass.BFloat16) + + if IS_VARLEN and not VARLEN_PURE: + if t < ci_eos: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + elif IS_VARLEN and VARLEN_PURE: + if chunk_idx < num_chunks: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + else: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + + if warp_row_group == 0: + for vi in cutlass.range_constexpr(VEC): + rGkOut[vi] = cute.exp2(rGkLast[vi] * cumsum_scale, fastmath=True) + if IS_VARLEN: + if ci_eos > cutlass.Int32(0): + cute.autovec_copy( + rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None] + ) + else: + cute.autovec_copy(rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None]) + + # K1 still reads Q/K in Pass 2b after signaling k1_done. + # TMA may reuse this stage only once those reads finish too. + cute.arch.mbarrier_arrive(stage_reuse_mbars + cur_stage) + + # ============================================================= + # Warp 26 (TMA_WARP_ID): dedicated TMA producer. + # Waits stage_reuse (K1 and MMA readers), issues TMA for Q/K/G, + # signals tma_mbar. Decouples MMA -> TMA dependency from K1 compute. + # ============================================================= + if warp_idx == TMA_WARP_ID: + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + # Prefetch chunk 0 -> stage 0 (stage_reuse[0] pre-arrived) + pf_cs = cutlass.Int32(0) + if IS_VARLEN: + pf_seq_id_0 = cutlass.Int32(mChunkIndices[chunk_base, 0]) + pf_local_0 = cutlass.Int32(mChunkIndices[chunk_base, 1]) + pf_bos_0 = cutlass.Int32(mCuSeqlens[pf_seq_id_0]) + pf_cs = pf_bos_0 + pf_local_0 * BT + else: + pf_cs = i_b * num_chunks * BT + chunk_base * BT + cute.arch.mbarrier_wait(stage_reuse_mbars, 0) + if lane_id == 0: + cute.arch.mbarrier_expect_tx(tma_mbars, bytes_per_stage) + sQ_pf = sQ[(None, None, 0)] + gQ_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gQ_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ_pf, 0, 2), + cute.group_modes(gQ_pf, 0, 2), + ) + cute.copy(tma_atom_Q, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + sK_pf = sK[(None, None, 0)] + gK_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gK_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_K, + 0, + cute.make_layout(1), + cute.group_modes(sK_pf, 0, 2), + cute.group_modes(gK_pf, 0, 2), + ) + cute.copy(tma_atom_K, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + sG_pf = sG[(None, None, 0)] + gG_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gG_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_G, + 0, + cute.make_layout(1), + cute.group_modes(sG_pf, 0, 2), + cute.group_modes(gG_pf, 0, 2), + ) + cute.copy(tma_atom_G, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars) + + # Issue TMAs for chunks 1..CHUNKS_PER_BLOCK-1 + for next_i in cutlass.range_constexpr(1, CHUNKS_PER_BLOCK): + next_stage = next_i % NUM_STAGES + next_phase = next_i // NUM_STAGES % 2 + next_cs = cutlass.Int32(0) + if IS_VARLEN: + next_chunk_idx = chunk_base + next_i + if next_chunk_idx < num_chunks: + _nsid = cutlass.Int32(mChunkIndices[next_chunk_idx, 0]) + next_cs = ( + cutlass.Int32(mCuSeqlens[_nsid]) + + cutlass.Int32(mChunkIndices[next_chunk_idx, 1]) * BT + ) + else: + next_cs = i_b * num_chunks * BT + (chunk_base + next_i) * BT + cute.arch.mbarrier_wait(stage_reuse_mbars + next_stage, next_phase) + if lane_id == 0: + cute.arch.mbarrier_expect_tx(tma_mbars + next_stage, bytes_per_stage) + sQ_ns = sQ[(None, None, next_stage)] + gQ_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gQ_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ_ns, 0, 2), + cute.group_modes(gQ_ns, 0, 2), + ) + cute.copy(tma_atom_Q, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + sK_ns = sK[(None, None, next_stage)] + gK_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gK_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_K, + 0, + cute.make_layout(1), + cute.group_modes(sK_ns, 0, 2), + cute.group_modes(gK_ns, 0, 2), + ) + cute.copy(tma_atom_K, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + sG_ns = sG[(None, None, next_stage)] + gG_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gG_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_G, + 0, + cute.make_layout(1), + cute.group_modes(sG_ns, 0, 2), + cute.group_modes(gG_ns, 0, 2), + ) + cute.copy(tma_atom_G, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars + next_stage) + + # ============================================================= + # Warps 16-26 (excluding TMA_WARP_ID=26): K2 MMA Compute + # ============================================================= + if ( + warp_idx >= NUM_K1_TMA_WARPS + and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS + and warp_idx != TMA_WARP_ID + ): + # Warp-layout invariants (scope-local → no cross-group register spill) + _tid_in_group = _lane % 4 + _group_id = _lane // 4 + mma_warp = _warp - NUM_K1_TMA_WARPS + my_i_q = cutlass.Int32(0) + my_i_k = cutlass.Int32(0) + if mma_warp < 1: + my_i_q = cutlass.Int32(0) + my_i_k = mma_warp + elif mma_warp < 3: + my_i_q = cutlass.Int32(1) + my_i_k = mma_warp - 1 + elif mma_warp < 6: + my_i_q = cutlass.Int32(2) + my_i_k = mma_warp - 3 + elif mma_warp < NUM_MMA_ACTIVE: + my_i_q = cutlass.Int32(3) + my_i_k = mma_warp - 6 + q_row_base = my_i_q * BC + k_row_base = my_i_k * BC + akk_row_base = k_row_base + akk_col_base = q_row_base + norm_row = q_row_base + if my_i_q == my_i_k: + norm_row = q_row_base + cutlass.Int32(BC // 2) + row0 = _group_id + row1 = _group_id + 8 + col0 = _tid_in_group * 2 + col1 = _tid_in_group * 2 + 1 + col2 = 8 + _tid_in_group * 2 + col3 = 8 + _tid_in_group * 2 + 1 + thr_mma = tiled_mma_k2.get_slice(_lane) + thr_copy_A = tiled_copy_mma_A.get_slice(_lane) + thr_copy_B = tiled_copy_mma_B.get_slice(_lane) + thr_copy_Gn = tiled_copy_Gcum_norm.get_slice(_tid_in_group) + thr_copy_Ggate = tiled_copy_Gcum_gate.get_slice(_lane) + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(k1_done_mbars + s, phase) + cute.arch.mbarrier_wait(store_done_mbars + s, phase) + + if mma_warp < NUM_MMA_ACTIVE: + csQ = sQ[(None, None, s)] + csK = sK[(None, None, s)] + csGcum = sGcum[(None, None, s)] + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + + _z = cutlass.Float32(0.0) + + # K1 stages activated beta with invalid rows set to zero. + # Reading it from SMEM also avoids repeating sigmoid for + # every lower-triangular K2 tile that shares a query row. + beta_row0 = sBeta[q_row_base + row0, s].to(cutlass.Float32) + beta_row1 = sBeta[q_row_base + row1, s].to(cutlass.Float32) + + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = _z, _z, _z, _z + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = _z, _z, _z, _z + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z + + # bf16 m16n8k16 MMA: each k_block covers k=16. + for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES): + # ---- Load Q/Kq bf16 fragments (16x16, 8 bf16/thread) ---- + sQ_tile = cute.local_tile(csQ, tiler=(16, 16), coord=(my_i_q, k_block)) + tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile)) + cute.copy( + tiled_copy_mma_A, + thr_copy_A.partition_S(sQ_tile), + thr_copy_A.retile(tCrQ), + ) + + sKq_tile = cute.local_tile(csK, tiler=(16, 16), coord=(my_i_q, k_block)) + tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile)) + cute.copy( + tiled_copy_mma_A, + thr_copy_A.partition_S(sKq_tile), + thr_copy_A.retile(tCrKq), + ) + + # ---- Issue K n0/n1 LDSMs early for better ILP ---- + sK_tile_n0 = cute.local_tile( + csK, tiler=(8, 16), coord=(my_i_k * 2, k_block) + ) + tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0)) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n0), + thr_copy_B.retile(tCrK_n0), + ) + + sK_tile_n1 = cute.local_tile( + csK, tiler=(8, 16), coord=(my_i_k * 2 + 1, k_block) + ) + tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1)) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n1), + thr_copy_B.retile(tCrK_n1), + ) + + # ---- Gate norm (2x k=8 covers k=16) ---- + sGn_a = cute.local_tile(csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2)) + tCsGn_a = thr_copy_Gn.partition_S(sGn_a) + tCrGn_a = cute.make_fragment_like(tCsGn_a, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn_a, thr_copy_Gn.retile(tCrGn_a)) + gn_a0 = tCrGn_a[0] + gn_a1 = tCrGn_a[1] + + sGn_b = cute.local_tile( + csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2 + 1) + ) + tCsGn_b = thr_copy_Gn.partition_S(sGn_b) + tCrGn_b = cute.make_fragment_like(tCsGn_b, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn_b, thr_copy_Gn.retile(tCrGn_b)) + gn_b0 = tCrGn_b[0] + gn_b1 = tCrGn_b[1] + + # ---- Gate Q (2x (16,8) partition_C covers m=16,k=16) ---- + sGq_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2)) + tCrGq_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_a)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_a), + thr_copy_Ggate.retile(tCrGq_a), + ) + + sGq_b = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2 + 1) + ) + tCrGq_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_b)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_b), + thr_copy_Ggate.retile(tCrGq_b), + ) + + # 8 Q gate values per thread (matching A bf16 m16n8k16 layout): + # first half k=0..7 (tCrGq_a): a0=(r0,c0) a1=(r0,c0+1) a2=(r0+8,c0) a3=(r0+8,c0+1) + # second half k=8..15 (tCrGq_b): a4=(r0,c0+8) a5=(r0,c0+9) a6=(r0+8,c0+8) a7=(r0+8,c0+9) + gate_q_0 = cute.exp2(tCrGq_a[0] - gn_a0, fastmath=True) + gate_q_1 = cute.exp2(tCrGq_a[1] - gn_a1, fastmath=True) + gate_q_2 = cute.exp2(tCrGq_a[2] - gn_a0, fastmath=True) + gate_q_3 = cute.exp2(tCrGq_a[3] - gn_a1, fastmath=True) + gate_q_4 = cute.exp2(tCrGq_b[0] - gn_b0, fastmath=True) + gate_q_5 = cute.exp2(tCrGq_b[1] - gn_b1, fastmath=True) + gate_q_6 = cute.exp2(tCrGq_b[2] - gn_b0, fastmath=True) + gate_q_7 = cute.exp2(tCrGq_b[3] - gn_b1, fastmath=True) + + # qa fp32 = Q*gate (8 per thread). tCrQ indexing assumption: + # [0..3] first k-chunk (k=0..7), [4..7] second k-chunk (k=8..15). + qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0 + qa1 = tCrQ[1].to(cutlass.Float32) * gate_q_1 + qa2 = tCrQ[2].to(cutlass.Float32) * gate_q_2 + qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3 + qa4 = tCrQ[4].to(cutlass.Float32) * gate_q_4 + qa5 = tCrQ[5].to(cutlass.Float32) * gate_q_5 + qa6 = tCrQ[6].to(cutlass.Float32) * gate_q_6 + qa7 = tCrQ[7].to(cutlass.Float32) * gate_q_7 + + # Pack fp32 qa -> 4 u32 bf16x2 for MMA A (k-adjacent pairs per u32) + qa_u32_0 = pack_bf16x2_f32(qa1, qa0) # reg0 = [qa0 lo | qa1 hi] + qa_u32_1 = pack_bf16x2_f32(qa3, qa2) + qa_u32_2 = pack_bf16x2_f32(qa5, qa4) + qa_u32_3 = pack_bf16x2_f32(qa7, qa6) + + ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0 + ka1 = tCrKq[1].to(cutlass.Float32) * gate_q_1 + ka2 = tCrKq[2].to(cutlass.Float32) * gate_q_2 + ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3 + ka4 = tCrKq[4].to(cutlass.Float32) * gate_q_4 + ka5 = tCrKq[5].to(cutlass.Float32) * gate_q_5 + ka6 = tCrKq[6].to(cutlass.Float32) * gate_q_6 + ka7 = tCrKq[7].to(cutlass.Float32) * gate_q_7 + + ka_u32_0 = pack_bf16x2_f32(ka1, ka0) + ka_u32_1 = pack_bf16x2_f32(ka3, ka2) + ka_u32_2 = pack_bf16x2_f32(ka5, ka4) + ka_u32_3 = pack_bf16x2_f32(ka7, ka6) + + # ---- Gate K (2x (16,8) partition_C covers m=16,k=16) ---- + sGk_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2)) + tCrGk_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_a)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_a), + thr_copy_Ggate.retile(tCrGk_a), + ) + + sGk_b = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2 + 1) + ) + tCrGk_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_b)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_b), + thr_copy_Ggate.retile(tCrGk_b), + ) + + # n0 uses rows 0..7 of (16,*) tile (tCrGk_*[0,1]) + # n1 uses rows 8..15 of (16,*) tile (tCrGk_*[2,3]) + gk_n0_0 = cute.exp2(gn_a0 - tCrGk_a[0], fastmath=True) + gk_n0_1 = cute.exp2(gn_a1 - tCrGk_a[1], fastmath=True) + gk_n0_2 = cute.exp2(gn_b0 - tCrGk_b[0], fastmath=True) + gk_n0_3 = cute.exp2(gn_b1 - tCrGk_b[1], fastmath=True) + + gk_n1_0 = cute.exp2(gn_a0 - tCrGk_a[2], fastmath=True) + gk_n1_1 = cute.exp2(gn_a1 - tCrGk_a[3], fastmath=True) + gk_n1_2 = cute.exp2(gn_b0 - tCrGk_b[2], fastmath=True) + gk_n1_3 = cute.exp2(gn_b1 - tCrGk_b[3], fastmath=True) + + # tCrK_n0 bf16 fragment: 4 elems/thread at (n_row, c0), (n_row, c0+1), + # (n_row, c0+8), (n_row, c0+9) — k-adjacent pairs + k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0 + k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1 + k_n0_b2 = tCrK_n0[2].to(cutlass.Float32) * gk_n0_2 + k_n0_b3 = tCrK_n0[3].to(cutlass.Float32) * gk_n0_3 + + k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0 + k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1 + k_n1_b2 = tCrK_n1[2].to(cutlass.Float32) * gk_n1_2 + k_n1_b3 = tCrK_n1[3].to(cutlass.Float32) * gk_n1_3 + + # Pack fp32 k -> 2 u32 bf16x2 for MMA B + k_n0_u32_0 = pack_bf16x2_f32(k_n0_b1, k_n0_b0) + k_n0_u32_1 = pack_bf16x2_f32(k_n0_b3, k_n0_b2) + k_n1_u32_0 = pack_bf16x2_f32(k_n1_b1, k_n1_b0) + k_n1_u32_1 = pack_bf16x2_f32(k_n1_b3, k_n1_b2) + + # ---- 4 bf16 MMA calls ---- + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_bf16_m16n8k16( + qa_u32_0, + qa_u32_1, + qa_u32_2, + qa_u32_3, + k_n0_u32_0, + k_n0_u32_1, + acc_aqk_n0_0, + acc_aqk_n0_1, + acc_aqk_n0_2, + acc_aqk_n0_3, + ) + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_bf16_m16n8k16( + qa_u32_0, + qa_u32_1, + qa_u32_2, + qa_u32_3, + k_n1_u32_0, + k_n1_u32_1, + acc_aqk_n1_0, + acc_aqk_n1_1, + acc_aqk_n1_2, + acc_aqk_n1_3, + ) + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_bf16_m16n8k16( + ka_u32_0, + ka_u32_1, + ka_u32_2, + ka_u32_3, + k_n0_u32_0, + k_n0_u32_1, + acc_akk_n0_0, + acc_akk_n0_1, + acc_akk_n0_2, + acc_akk_n0_3, + ) + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_bf16_m16n8k16( + ka_u32_0, + ka_u32_1, + ka_u32_2, + ka_u32_3, + k_n1_u32_0, + k_n1_u32_1, + acc_akk_n1_0, + acc_akk_n1_1, + acc_akk_n1_2, + acc_akk_n1_3, + ) + + # sQ/sK/sG reads done, signal TMA before SMEM writes + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + # Dual-path MMA write (constexpr-gated): + # - non-pure: apply causal + diag=1 inline so SMEM matches + # final GMEM layout (pairs with row-major vec autovec + # store warp that does no causal). + # - pure: write all 16x16 unconditionally (pairs with + # per-tile store warp that applies causal + diag in + # store; this is the original baseline behavior — no + # extra MMA-write cost). + _z16 = cutlass.BFloat16(0.0) + _one16 = cutlass.BFloat16(1.0) + if IS_VARLEN and not VARLEN_PURE: + if my_i_q == my_i_k: + # Diagonal sub-tile: causal-mask sAqk and write + # diag=1 / strict-lower=MMA*beta / strict-upper=0 + # to sAkk so SMEM is in final form. + _v_q00 = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16) + if row0 < col0: + _v_q00 = _z16 + csAqk[q_row_base + row0, k_row_base + col0] = _v_q00 + _v_q01 = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16) + if row0 < col1: + _v_q01 = _z16 + csAqk[q_row_base + row0, k_row_base + col1] = _v_q01 + _v_q02 = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16) + if row1 < col0: + _v_q02 = _z16 + csAqk[q_row_base + row1, k_row_base + col0] = _v_q02 + _v_q03 = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16) + if row1 < col1: + _v_q03 = _z16 + csAqk[q_row_base + row1, k_row_base + col1] = _v_q03 + _v_q04 = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16) + if row0 < col2: + _v_q04 = _z16 + csAqk[q_row_base + row0, k_row_base + col2] = _v_q04 + _v_q05 = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16) + if row0 < col3: + _v_q05 = _z16 + csAqk[q_row_base + row0, k_row_base + col3] = _v_q05 + _v_q06 = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16) + if row1 < col2: + _v_q06 = _z16 + csAqk[q_row_base + row1, k_row_base + col2] = _v_q06 + _v_q07 = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16) + if row1 < col3: + _v_q07 = _z16 + csAqk[q_row_base + row1, k_row_base + col3] = _v_q07 + + _v_k00 = (acc_akk_n0_0 * beta_row0).to(cutlass.BFloat16) + if row0 == col0: + _v_k00 = _one16 + if row0 < col0: + _v_k00 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col0] = _v_k00 + _v_k01 = (acc_akk_n0_1 * beta_row0).to(cutlass.BFloat16) + if row0 == col1: + _v_k01 = _one16 + if row0 < col1: + _v_k01 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col1] = _v_k01 + _v_k02 = (acc_akk_n0_2 * beta_row1).to(cutlass.BFloat16) + if row1 == col0: + _v_k02 = _one16 + if row1 < col0: + _v_k02 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col0] = _v_k02 + _v_k03 = (acc_akk_n0_3 * beta_row1).to(cutlass.BFloat16) + if row1 == col1: + _v_k03 = _one16 + if row1 < col1: + _v_k03 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col1] = _v_k03 + _v_k04 = (acc_akk_n1_0 * beta_row0).to(cutlass.BFloat16) + if row0 == col2: + _v_k04 = _one16 + if row0 < col2: + _v_k04 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col2] = _v_k04 + _v_k05 = (acc_akk_n1_1 * beta_row0).to(cutlass.BFloat16) + if row0 == col3: + _v_k05 = _one16 + if row0 < col3: + _v_k05 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col3] = _v_k05 + _v_k06 = (acc_akk_n1_2 * beta_row1).to(cutlass.BFloat16) + if row1 == col2: + _v_k06 = _one16 + if row1 < col2: + _v_k06 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col2] = _v_k06 + _v_k07 = (acc_akk_n1_3 * beta_row1).to(cutlass.BFloat16) + if row1 == col3: + _v_k07 = _one16 + if row1 < col3: + _v_k07 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col3] = _v_k07 + else: + # Non-diag (i_q > i_k): write all 16x16 unchanged. + csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row0, akk_col_base + col0] = ( + acc_akk_n0_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col1] = ( + acc_akk_n0_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col0] = ( + acc_akk_n0_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col1] = ( + acc_akk_n0_3 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col2] = ( + acc_akk_n1_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col3] = ( + acc_akk_n1_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col2] = ( + acc_akk_n1_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col3] = ( + acc_akk_n1_3 * beta_row1 + ).to(cutlass.BFloat16) + else: + # PURE: write all 16x16 unconditionally — store warp + # applies causal+diag in its per-tile loop. This is + # the original baseline behavior (no extra MMA cost). + csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row0, akk_col_base + col0] = ( + acc_akk_n0_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col1] = ( + acc_akk_n0_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col0] = ( + acc_akk_n0_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col1] = ( + acc_akk_n0_3 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col2] = ( + acc_akk_n1_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col3] = ( + acc_akk_n1_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col2] = ( + acc_akk_n1_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col3] = ( + acc_akk_n1_3 * beta_row1 + ).to(cutlass.BFloat16) + else: + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + cute.arch.mbarrier_arrive(mma_done_mbars + s) + + # ============================================================= + # Warps 27-30: Store/Inversion warps + # ============================================================= + if warp_idx >= NUM_K1_TMA_WARPS + NUM_MMA_WARPS: + store_warp = warp_idx - (NUM_K1_TMA_WARPS + NUM_MMA_WARPS) + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + st_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + st_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(mma_done_mbars + s, phase) + + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + + # Dual-path row-major store. SMEM is in final GMEM layout + # already (causal mask + diag=1 applied at MMA write; upper-tri + # SMEM is zero from CTA-startup init). Both paths use vec2 + # autovec_copy → STG.E.32 (4-byte coalesced). + # + # mAqk_v2 / mAkk_v2 shape (B, T, H, BT/2, 2): last dim is vec2. + # + # NON-PURE: full 64-col row-major + row mask (chunk may + # overflow seq end). 32 lanes/warp × 1 vec2 = full row. + # PURE: reduced cols. Warp s writes (s+1)*16 cols/row (no row + # mask, all rows in seq). Saves ~37% GMEM bandwidth vs full. + col_lo = lane_id * 2 + col_hi = col_lo + 1 + col_vec_idx = lane_id + rAqkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rAkkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + + if IS_VARLEN and not VARLEN_PURE: + # NON-PURE: row-major full-row vec autovec + row mask. SMEM + # upper-tri is zero (MMA-masked at write), so writing all 64 + # cols is correct. STG.E.32 (32 lanes × 4 bytes per row). + row_base_warp = store_warp * (BT // NUM_STORE_WARPS) + for ri in cutlass.range_constexpr(BT // NUM_STORE_WARPS): + local_row = row_base_warp + ri + abs_row = chunk_start + local_row + rAqkOut[0] = csAqk[local_row, col_lo] + rAqkOut[1] = csAqk[local_row, col_hi] + rAkkOut[0] = csAkk[local_row, col_lo] + rAkkOut[1] = csAkk[local_row, col_hi] + if abs_row < st_eos: + cute.autovec_copy( + rAqkOut, mAqk_v2[i_b, abs_row, i_h, col_vec_idx, None] + ) + cute.autovec_copy( + rAkkOut, mAkk_v2[i_b, abs_row, i_h, col_vec_idx, None] + ) + else: + # PURE/eqlen: reduced lower-triangular store. Eqlen folds + # tile_valid to True; pure varlen checks the scheduler's + # potentially padded tail slot once per store warp. + tile_valid = True if not IS_VARLEN else chunk_idx < num_chunks + if tile_valid: + for tile_idx in cutlass.range_constexpr(NUM_TILES): + i_q = _TILE_IQ[tile_idx] + i_k = _TILE_IK[tile_idx] + is_diag = _TILE_IQ[tile_idx] == _TILE_IK[tile_idx] + gmem_aqk_row_base = chunk_start + i_q * BC + gmem_aqk_col_base = i_k * BC + gmem_akk_row_base = chunk_start + i_k * BC + gmem_akk_col_base = i_q * BC + smem_aqk_row_base = i_q * BC + smem_aqk_col_base = i_k * BC + smem_akk_row_base = i_k * BC + smem_akk_col_base = i_q * BC + for ri in cutlass.range_constexpr(BC // NUM_STORE_WARPS): + local_row = store_warp * (BC // NUM_STORE_WARPS) + ri + if lane_id < BC: + local_col = lane_id + aqk_val = csAqk[ + smem_aqk_row_base + local_row, + smem_aqk_col_base + local_col, + ] + akk_val = csAkk[ + smem_akk_row_base + local_row, + smem_akk_col_base + local_col, + ] + if is_diag and local_row < local_col: + aqk_val = cutlass.BFloat16(0.0) + if is_diag and local_row < local_col: + akk_val = cutlass.BFloat16(0.0) + if is_diag and local_row == local_col: + akk_val = cutlass.BFloat16(1.0) + mAqk[ + i_b, + gmem_aqk_row_base + local_row, + i_h, + gmem_aqk_col_base + local_col, + ] = aqk_val + mAkk[ + i_b, + gmem_akk_row_base + local_row, + i_h, + gmem_akk_col_base + local_col, + ] = akk_val + + cute.arch.mbarrier_arrive(store_done_mbars + s) + + yield_out() + + +# ========================================================================= +# Host function +# ========================================================================= +def make_host_function( + B, + NT, + H, + is_varlen=False, + T_padded=None, + has_bias=False, + use_safe_gate=False, + use_beta_sigmoid=False, + varlen_pure=False, + grid_blocks=NUM_SMS, +): + """ + `grid_blocks` is the persistent grid size (one resident block per SM); the + caller passes the SMs this kernel may occupy. + + `varlen_pure=True` asserts that all seq lengths in the batch are multiples + of BT (= 64). Under that assumption every chunk has 64 valid rows so the + four mask sites (K1 row mask, K1 store mask, MMA accumulator zero-fill, + Store row mask) are guaranteed to never fire and are dead-code eliminated + at compile time. Caller's data layout is unchanged — this is a hint only. + """ + _H = H + _IS_VARLEN = 1 if is_varlen else 0 + _HAS_BIAS = 1 if has_bias else 0 + _USE_SAFE_GATE = 1 if use_safe_gate else 0 + _USE_BETA_SIGMOID = 1 if use_beta_sigmoid else 0 + _VARLEN_PURE = 1 if (is_varlen and varlen_pure) else 0 + if is_varlen: + assert B == 1, "Varlen requires B=1" + assert T_padded is not None, "T_padded required for varlen" + + # B / NT / T are RUNTIME arguments of host_fn (rt_b / rt_nt / + # rt_t_total below), not closure constants: baking them keyed every + # compiled kernel to the batch shape, and eval traffic (a distinct + # total token count per prefill batch) recompiled ~100 s per batch — + # observed as GSM8K "hangs" (0/1319 responses). Only genuine constexpr + # specializers (H, is_varlen, bias/gate/pure variants) stay baked. + s_row = _H * K_DIM + s_col = 1 + s_h = K_DIM + + @cute.jit + def host_fn( + mQ, + mK, + mG, + mA_log, + mBeta, + mBetaActivated, + scale, + mKscaled, + mKg, + mQscaled, + mGkLast, + mAqk, + mAkk, + mCuSeqlens, + mChunkIndices, + mDtBias, + lower_bound_val, + rt_nt: cutlass.Int32, + rt_b: cutlass.Int32, + rt_t_total: cutlass.Int32, + valid_tokens: cutlass.Int32, + # Launch stream — a runtime argument (the executor runs the model on + # a dedicated non-blocking torch stream; launching on the DSL default + # stream races with the caller's stream). See _launch_fused_k123_inv. + stream: cuda.CUstream, + ): + # 3D TMA view: (B*T, K_DIM, H). domain_offset in kernel shifts to + # arbitrary chunk_start — no BT alignment required for varlen. + # Dynamic T extent: the TMA tensormap is materialized per call from + # the runtime shape, so one compiled host_fn serves every batch. + view_layout_3d = cute.make_layout( + (rt_t_total, K_DIM, _H), + stride=(s_row, s_col, s_h), + ) + mQ_view = cute.make_tensor(mQ.iterator, view_layout_3d) + mK_view = cute.make_tensor(mK.iterator, view_layout_3d) + mG_view = cute.make_tensor(mG.iterator, view_layout_3d) + + smem_atom_qk = tcgen05.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + qk_smem_2d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM), order=(0, 1)) + qk_smem_3d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM, NUM_STAGES), order=(0, 1, 2)) + + g_smem_2d = cute.make_layout((BT, K_DIM), stride=(K_DIM, 1)) + g_smem_3d = cute.make_layout((BT, K_DIM, NUM_STAGES), stride=(K_DIM, 1, BT * K_DIM)) + + tma_op = cpasync.CopyBulkTensorTileG2SOp(cpasync.CtaGroup.ONE) + tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom( + tma_op, mQ_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + tma_op, mK_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + tma_atom_G, tma_tensor_G = cpasync.make_tiled_tma_atom( + tma_op, mG_view, g_smem_2d, cute.product_each(g_smem_2d.shape), num_multicast=1 + ) + + g_cumsum_layout = cute.make_layout( + (BT, K_DIM, NUM_STAGES), stride=(K_STRIDE, 1, BT * K_STRIDE) + ) + + copy_atom_qk_k1 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=32 + ) + tiled_copy_qk_k1 = cute.make_tiled_copy_tv( + copy_atom_qk_k1, + thr_layout=cute.make_layout((1, 32)), + val_layout=cute.make_layout((1, 2)), + ) + + # The v2 views are indexed [i_b, per-batch row, ...] by the kernel, + # so their T extent / batch stride must be PER-BATCH T, not the flat + # B*T used by the 3D TMA view above. Using rt_t_total here sent every + # i_b > 0 store rt_b× too far (OOB writes + parity failure for the + # eqlen B=2 path; varlen is B=1 so it never noticed). + rt_t_pb = rt_t_total // rt_b + out_v2_layout = cute.make_layout( + (rt_b, rt_t_pb, _H, K_VEC, VEC), + stride=(rt_t_pb * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mKscaled_v2 = cute.make_tensor(mKscaled.iterator, out_v2_layout) + mQscaled_v2 = cute.make_tensor(mQscaled.iterator, out_v2_layout) + mKg_v2 = cute.make_tensor(mKg.iterator, out_v2_layout) + + # vec2 views for mAqk / mAkk (BT dimension instead of K_DIM). + # Shape: (B, T, H, BT/2, 2). Each VEC=2 slot = 2 contiguous bf16 = 1 + # fp32-aligned 4-byte unit. Used by store warp's autovec_copy → + # STG.E.32 with 32 lanes coalesced to 1 cache line per row. + BT_VEC = BT // VEC # 32 + akk_v2_layout = cute.make_layout( + (rt_b, rt_t_pb, _H, BT_VEC, VEC), + stride=(rt_t_pb * _H * BT, _H * BT, BT, VEC, 1), + ) + mAqk_v2 = cute.make_tensor(mAqk.iterator, akk_v2_layout) + mAkk_v2 = cute.make_tensor(mAkk.iterator, akk_v2_layout) + + gklast_v2_layout = cute.make_layout( + (rt_b, rt_nt, _H, K_VEC, VEC), + stride=(rt_nt * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mGkLast_v2 = cute.make_tensor(mGkLast.iterator, gklast_v2_layout) + + mma_op = cute.nvgpu.warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma_k2 = cute.make_tiled_mma( + mma_op, cute.make_layout((1, 1, 1)), permutation_mnk=(16, 8, 16) + ) + + tiled_copy_mma_A = cute.make_tiled_copy_A( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16), + tiled_mma_k2, + ) + tiled_copy_mma_B = cute.make_tiled_copy_B( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16), + tiled_mma_k2, + ) + + copy_atom_Gcum = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64 + ) + tiled_copy_Gcum_norm = cute.make_tiled_copy_tv( + copy_atom_Gcum, thr_layout=cute.make_layout((1, 4)), val_layout=cute.make_layout((1, 2)) + ) + tiled_copy_Gcum_gate = cute.make_tiled_copy_C(copy_atom_Gcum, tiled_mma_k2) + + smem_size = ( + BT * K_DIM * 2 * 2 * NUM_STAGES + + BT * K_DIM * 2 * NUM_STAGES + + BT * K_STRIDE * 4 * NUM_STAGES + + BT * NUM_STAGES * 2 # sBeta bf16 + + K1_ROW_GROUPS * PARTIAL_COLS * 4 + + BT * AQK_TILE_STRIDE * 2 * NUM_STAGES # sAqk bf16 (64x72 row-major) + + BT * AKK_STRIDE * 2 * NUM_STAGES # sAkk bf16 + + 512 + ) + + # Constant grid: the persistent for_generate loop bounds work by the + # runtime total_cgs, so blocks past the work count exit immediately. + # (A shape-dependent min() here would re-bake the grid per NT.) + _grid_x = grid_blocks + + fused_kernel123( + tma_atom_Q, + tma_tensor_Q, + tma_atom_K, + tma_tensor_K, + tma_atom_G, + tma_tensor_G, + mA_log, + mBeta, + mBetaActivated, + scale, + mKscaled_v2, + mKg_v2, + mQscaled_v2, + mGkLast_v2, + mAqk, + mAkk, + mAqk_v2, + mAkk_v2, + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_3d, + g_smem_3d, + g_cumsum_layout, + rt_nt, + _H, + rt_b, + mCuSeqlens, + mChunkIndices, + _IS_VARLEN, + mDtBias, + lower_bound_val, + _HAS_BIAS, + _USE_SAFE_GATE, + valid_tokens, + _USE_BETA_SIGMOID, + _VARLEN_PURE, + grid_blocks, + ).launch( + grid=(_grid_x, 1, 1), + block=(THREADS, 1, 1), + smem=smem_size, + stream=stream, + ) + + return host_fn diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/k4_persistent.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/k4_persistent.py new file mode 100644 index 000000000000..0f983ed9aa46 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/k4_persistent.py @@ -0,0 +1,1637 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""K4-only fused KDA kernel: persistent variant with varlen support. + +Persistent scheduler (GDNTileScheduler) with 3D tensor layout, TensorMapManager +for per-tile TMA descriptor updates, and domain_offset+flat_divide for per-chunk +addressing. Supports variable-length sequences via cu_seqlens. + +Precondition: every sequence in cu_seqlens must have length > 0. A zero-length +sequence yields num_chunks = 0, so its scheduler tile skips the chunk loop and +never commits the per-chunk mbarriers the other warp groups wait on (deadlock). +The prefill runtime never produces zero-length sequences, so this is not +checked on the device or in the launch path (cu_seqlens lives on the GPU; a +host-side check would force a sync on the hot path). + +K4 chunk loop with 6 MMAs per chunk: + MMA1: W = AB @ KS (K-MN, K=64) + MMA2: U = AB @ V (K-MN, K=64) + MMA3: NV = U - W_bf16 @ S (SS-mode: A=SMEM sO, B=SMEM sST, K=128, accumulate; + W is negated in its readout, FLA: b_v = u - w @ h) + MMA4: OI = QS @ S (K-MN, K=128) + MMA5: O = OI + AQC @ NV (K-MN, K=64, accumulate) + MMA6: State += NV^T @ KG (MN-MN, K=64, accumulate on decayed state) + +Execution order: MMA1->MMA2->MMA4->MMA3->MMA5->MMA6 + +Warp assignment (3 warpgroups, 12 warps, 384 threads): + WG0 (W0-3): W0=MMA issue, W2=TMA load/store, W1/W3=idle + WG1 (W4-7): State 2-pass: TMEM->bf16->sST, gk(SMEM)->decay->TMEM + WG2 (W8-11): GDN readout: W/NV/O TMEM->bf16->SMEM + +TMA pipelines (all PipelineTmaUmma, 1-stage): + - KS, V, QS: prefetch c+1 while MMA uses c + - AB, AQC, KG: 1-stage with explicit consumer release (no prefetch) + - TMA warp decoupled from O readout (no store_nbar wait) + +State management: + - tCtState [128,128] fp32 @ TMEM offset 256: persistent + - 2-pass per chunk: bf16->sST early, gk decay->TMEM late + - gk_last preloaded to SMEM (coalesced), read in Pass 2 + - single state_ready_mbar after both passes (TMEM no-overlap constraint) + +Mid-sequence state snapshot (host function built with ``has_snapshot=True``): + per-sequence int32 ``snapshot_chunks[b]`` / ``snapshot_rows[b]``. When + ``snapshot_chunks[b] == c`` (``0 <= c < num_chunks``), the state warpgroup + also stores the state at the top of chunk ``c`` -- the state after the + sequence's first ``64 * c`` tokens, before chunk ``c``'s decay -- into + ``s_fp32[snapshot_rows[b]]``. ``-1`` disables the snapshot for a sequence. + The snapshot row must not alias any sequence's ``state_indices`` row. The + default host function passes no snapshot tensors and compiles without this + store. +""" + +from typing import Optional + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.cutlass_dsl as _dsl_mod +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05 import Field, OperandMajorMode, OperandSource +from cutlass.cutlass_dsl import Int32, T, dsl_user_op +from cutlass.utils import TensorMapManager, TensorMapUpdateMode + +# Compatibility shim, global effect: flashinfer's gated_delta_net_tile_scheduler +# imports CuteExperimentalDSL from cutlass.cutlass_dsl, which older CUTLASS DSL +# releases do not define. Install an inert placeholder so the import below +# succeeds; other modules in the process see the same attribute. +if not hasattr(_dsl_mod, "CuteExperimentalDSL"): + + class _DummyExperimentalDSL: + jit = None + kernel = None + compile = None + + _dsl_mod.CuteExperimentalDSL = _DummyExperimentalDSL + +from flashinfer.gdn_kernels.blackwell.gated_delta_net_tile_scheduler import ( + GDNTileScheduler, + GDNTileSchedulerParams, +) + +SB = 16 +AKK_PAD = 8 +AKK_STRIDE = 64 + AKK_PAD # 72 +TEMP_COLS = SB + AKK_PAD # 24 +NUM_TEMPS = 2 + + +# =========================================================================== +# Inverse dsl_user_op functions (TF32 MMA m16n8k8, barrier 7) +# =========================================================================== +@dsl_user_op +def mma_tf32_m16n8k8(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def inv_internal_barrier(*, loc=None, ip=None): + llvm.inline_asm( + T.i32(), + [], + "bar.sync 7, 128; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _invert_diag(sAkk: cute.Tensor, block_rc, lane_id, *, loc=None, ip=None): + my_row = lane_id % 16 + halfwarp_base = (lane_id // 16) * 16 + r_off = block_rc * 16 + c_off = block_rc * 16 + rInv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32) + rInv[0] = cutlass.Float32(1.0) + for x in range(1, 16): + rInv[x] = cutlass.Float32(0.0) + for d in range(1, 16): + col_d = my_row - d + valid = cutlass.Float32(col_d >= 0) + a_val = cutlass.Float32(sAkk[r_off + my_row, c_off + col_d]) * valid + acc = cutlass.Float32(0.0) + for j in range(1, d): + a_re = cutlass.Float32(sAkk[r_off + my_row, c_off + my_row - (d - j)]) + inv_shfl = cute.arch.shuffle_sync(rInv[j], halfwarp_base + my_row - d + j) + acc = acc + a_re * inv_shfl + rInv[d] = (-a_val - acc) * valid + rInv[0] = cutlass.Float32(1.0) + sAkk[r_off + my_row, c_off + my_row] = rInv[0] + for d in range(1, 16): + sAkk[r_off + my_row, c_off + (my_row + 16 - d) % 16] = rInv[d] * cutlass.Float32( + my_row >= d + ) + + +@dsl_user_op +def _matmul_AB_inv(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 16 + rB = br_B * 16 + cB = bc_B * 16 + a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1]) + b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, _z, _z, _z, _z) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, _z, _z, _z, _z) + a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1]) + b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _chain_mma_B_inv( + sAkk: cute.Tensor, + br_B, + bc_B, + a0k0, + a1k0, + a2k0, + a3k0, + a0k1, + a1k1, + a2k1, + a3k1, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 16 + b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0k0, a1k0, a2k0, a3k0, b0n0, b1n0, _z, _z, _z, _z + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0k0, a1k0, a2k0, a3k0, b0n1, b1n1, _z, _z, _z, _z + ) + b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0k1, a1k1, a2k1, a3k1, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0k1, a1k1, a2k1, a3k1, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _chain_mma_A_inv( + sAkk: cute.Tensor, + br_A, + bc_A, + b0_k0n0, + b1_k0n0, + b0_k0n1, + b1_k0n1, + b0_k1n0, + b1_k1n0, + b0_k1n1, + b1_k1n1, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 16 + a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n0, b1_k0n0, _z, _z, _z, _z) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n1, b1_k0n1, _z, _z, _z, _z) + a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0_k1n0, b1_k1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0_k1n1, b1_k1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _store_neg_C_inv( + sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 16 + sAkk[r + gid, c + 2 * tid] = -c0 + sAkk[r + gid, c + 2 * tid + 1] = -c1 + sAkk[r + gid + 8, c + 2 * tid] = -c2 + sAkk[r + gid + 8, c + 2 * tid + 1] = -c3 + sAkk[r + gid, c + 8 + 2 * tid] = -c4 + sAkk[r + gid, c + 8 + 2 * tid + 1] = -c5 + sAkk[r + gid + 8, c + 8 + 2 * tid] = -c6 + sAkk[r + gid + 8, c + 8 + 2 * tid + 1] = -c7 + + +@dsl_user_op +def _shuffle_C_to_B_inv(c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + src_a = 8 * tid + gid // 2 + src_b = src_a + 4 + f_odd = cutlass.Float32(gid % 2) + f_even = cutlass.Float32(1) - f_odd + c0_a = cute.arch.shuffle_sync(c0, src_a) + c1_a = cute.arch.shuffle_sync(c1, src_a) + c2_a = cute.arch.shuffle_sync(c2, src_a) + c3_a = cute.arch.shuffle_sync(c3, src_a) + c4_a = cute.arch.shuffle_sync(c4, src_a) + c5_a = cute.arch.shuffle_sync(c5, src_a) + c6_a = cute.arch.shuffle_sync(c6, src_a) + c7_a = cute.arch.shuffle_sync(c7, src_a) + c0_b = cute.arch.shuffle_sync(c0, src_b) + c1_b = cute.arch.shuffle_sync(c1, src_b) + c2_b = cute.arch.shuffle_sync(c2, src_b) + c3_b = cute.arch.shuffle_sync(c3, src_b) + c4_b = cute.arch.shuffle_sync(c4, src_b) + c5_b = cute.arch.shuffle_sync(c5, src_b) + c6_b = cute.arch.shuffle_sync(c6, src_b) + c7_b = cute.arch.shuffle_sync(c7, src_b) + b0_00 = c0_a * f_even + c1_a * f_odd + b1_00 = c0_b * f_even + c1_b * f_odd + b0_10 = c2_a * f_even + c3_a * f_odd + b1_10 = c2_b * f_even + c3_b * f_odd + b0_01 = c4_a * f_even + c5_a * f_odd + b1_01 = c4_b * f_even + c5_b * f_odd + b0_11 = c6_a * f_even + c7_a * f_odd + b1_11 = c6_b * f_even + c7_b * f_odd + return b0_00, b1_00, b0_10, b1_10, b0_01, b1_01, b0_11, b1_11 + + +@dsl_user_op +def _store_C_temp_inv( + sT: cute.Tensor, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _load_C_temp_inv(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +def transform_partitioned_tensor_layout(tensor): + layout = tensor.layout + stored_layout = layout + if isinstance(stored_layout, cute.ComposedLayout): + layout = layout.outer + shape = layout.shape + stride = layout.stride + new_shape = ((shape[0][0], shape[1]), (shape[0][1], shape[2]), *shape[3:]) + new_stride = ((stride[0][0], stride[1]), (stride[0][1], stride[2]), *stride[3:]) + new_layout = cute.make_layout(shape=new_shape, stride=new_stride) + if isinstance(stored_layout, cute.ComposedLayout): + new_layout = cute.make_composed_layout( + stored_layout.inner, stored_layout.offset, new_layout + ) + return cute.make_tensor(tensor.iterator, new_layout) + + +mma_dtype = cutlass.BFloat16 +acc_dtype = cutlass.Float32 +out_dtype = cutlass.BFloat16 + +M = 64 +N = 128 +K = 64 +K3 = 128 +M6 = 128 +N6 = 128 +K6 = 64 + +threads_per_cta = 384 +warp_threads = 32 +warpgroup_threads = 128 + +BYTES_PER_TENSORMAP = 128 +NUM_TENSORMAPS = 7 # a, b, v, q, aqc, kg, o + +MMA_WARP = 0 +O_STORE_WARP = 1 +TMA_WARP = 2 +STATE_WG = 1 +READOUT_WG = 2 + +NUM_REGS_WG0 = 40 +NUM_REGS_WG1 = 232 +NUM_REGS_WG2 = 232 +MAX_REGS = 168 + +try: + from tensorrt_llm.logger import logger as _logger +except ImportError: # standalone kernel use outside the full package + import logging as _logging + + _logger = _logging.getLogger(__name__) + +# Monkey-patch, GLOBAL effect: CuTeDSL._get_pipeline is class-level state, so +# every CuTe DSL kernel compiled in this process after this import (not just +# K4) gets ptxas --uumn (unified uniform register allocation, needed here to +# keep WG0 under its 40-register budget). The DSL exposes no per-compile +# ptx-options hook at the cute.compile call sites, hence the patch. The extra +# flag is benign for the other kernels in this package. Idempotent: skips +# pipelines that already carry ptx-options. +try: + from cutlass.cutlass_dsl.cutlass import CuTeDSL as _CuTeDSL + + _orig_get_pipeline = _CuTeDSL._get_pipeline + _patch_applied = False + + def _patched_get_pipeline(self, _pipeline): + global _patch_applied + result = _orig_get_pipeline(self, _pipeline) + if result and "ptx-options=" not in result: + if "cubin-format=bin" in result: + result = result.replace("cubin-format=bin", "cubin-format=bin ptx-options='--uumn'") + _patch_applied = True + else: + _logger.warning( + f"k4_persistent ptx-options patch: 'cubin-format=bin' not found " + f"in pipeline: {result[:200]}" + ) + elif result and "ptx-options=" in result: + _logger.debug("k4_persistent ptx-options patch: ptx-options already present") + _patch_applied = True + return result + + _CuTeDSL._get_pipeline = _patched_get_pipeline +# Workaround for a CuTe DSL parser bug: the nvidia-cutlass-dsl 4.5.0 AST +# preprocessor cannot parse tuple except handlers anywhere in a kernel +# module ("'Tuple' object has no attribute 'id'"), which breaks every +# cute.compile of this file. Keep a single Exception clause until the DSL +# pin moves past the bug; the isinstance re-raise below narrows it to the +# expected failure kinds without putting a tuple in the except clause. +except Exception as e: + if not isinstance(e, (AttributeError, ImportError)): + # Not a known-benign "DSL not present / API surface moved" case: + # swallowing it would defer the blow-up to the first K4 compile + # (ptxas over WG0's 40-register budget without --uumn) with no + # trail back to this patch. Fail here, where the cause is visible. + raise + _logger.warning(f"k4_persistent ptx-options patch failed: {e}") + _patch_applied = False + + +@cute.kernel +def k4_persistent_kernel( + tiled_mma_kmn: cute.TiledMma, + tiled_mma_mn_mn: cute.TiledMma, + tma_a, + a_sl: cute.ComposedLayout, + tma_b, + b_sl: cute.ComposedLayout, + tma_v, + v_sl: cute.ComposedLayout, + s_sl: cute.ComposedLayout, + tma_q, + q_sl: cute.ComposedLayout, + tma_aqc, + aqc_sl: cute.ComposedLayout, + tma_kg, + kg_sl: cute.ComposedLayout, + tma_o, + store_sl: cute.ComposedLayout, + readout_k_sl: cute.ComposedLayout, + nv_b_sl: cute.ComposedLayout, + nv_a_sl: cute.ComposedLayout, + kg_a_sl: cute.ComposedLayout, + nv_b_mn_sl: cute.ComposedLayout, + mGkLastExp: cute.Tensor, + mS_fp32: cute.Tensor, + state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + snapshot_rows: Optional[cute.Tensor], + snapshot_chunks: Optional[cute.Tensor], + mA: cute.Tensor, + mB: cute.Tensor, + mV_g: cute.Tensor, + mQ: cute.Tensor, + mAQC: cute.Tensor, + mKG: cute.Tensor, + mO: cute.Tensor, + tensormap_workspace: cute.Tensor, + scheduler_params: GDNTileSchedulerParams, +): + bidx, bidy, bidz = cute.arch.block_idx() + grid_dim = cute.arch.grid_dim() + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(tidx // warp_threads) + warpgroup_idx = cute.arch.make_warp_uniform(tidx // warpgroup_threads) + warpgroup_tidx = tidx % warpgroup_threads + thr_kmn = tiled_mma_kmn.get_slice(0) + thr_mn = tiled_mma_mn_mn.get_slice(0) + dice = (None, None, None) + + cta_linear_idx = bidz * grid_dim[1] * grid_dim[0] + bidy * grid_dim[0] + bidx + tensormap_manager = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) + tm_ws = cute.make_tensor( + tensormap_workspace.iterator, + cute.make_layout( + (grid_dim[0] * grid_dim[1] * grid_dim[2], NUM_TENSORMAPS, BYTES_PER_TENSORMAP), + stride=(NUM_TENSORMAPS * BYTES_PER_TENSORMAP, BYTES_PER_TENSORMAP, 1), + ), + ) + tm_a_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 0, None)].iterator) + tm_b_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 1, None)].iterator) + tm_v_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 2, None)].iterator) + tm_q_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 3, None)].iterator) + tm_aqc_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 4, None)].iterator) + tm_kg_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 5, None)].iterator) + tm_o_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 6, None)].iterator) + + smem = cutlass.utils.SmemAllocator() + AL = 128 + sA = smem.allocate_tensor(mma_dtype, a_sl.outer, AL, a_sl.inner) + sB = smem.allocate_tensor(mma_dtype, b_sl.outer, AL, b_sl.inner) + sV = smem.allocate_tensor(mma_dtype, v_sl.outer, AL, v_sl.inner) + sST = smem.allocate_tensor(mma_dtype, s_sl.outer, AL, s_sl.inner) + sQ = smem.allocate_tensor(mma_dtype, q_sl.outer, AL, q_sl.inner) + sAQC = smem.allocate_tensor(mma_dtype, aqc_sl.outer, AL, aqc_sl.inner) + sKG = smem.allocate_tensor(mma_dtype, kg_sl.outer, AL, kg_sl.inner) + sNV = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sO = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sO_out = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sGk_buf = smem.allocate_array(acc_dtype, N) + + sNV_b = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_b_sl.inner, mma_dtype), nv_b_sl.outer) + # sNV_a = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_a_sl.inner, mma_dtype), nv_a_sl.outer) + sKG_a = cute.make_tensor(cute.recast_ptr(sKG.iterator, kg_a_sl.inner, mma_dtype), kg_a_sl.outer) + sNV_b_mn = cute.make_tensor( + cute.recast_ptr(sNV.iterator, nv_b_mn_sl.inner, mma_dtype), nv_b_mn_sl.outer + ) + sO_out_st = cute.make_tensor( + cute.recast_ptr(sO_out.iterator, store_sl.inner, out_dtype), store_sl.outer + ) + + tmem_smem = smem.allocate_array(cutlass.Int32, 1) + if warp_idx == 0: + cute.arch.alloc_tmem(512, tmem_smem) + + sNV_ready_nbar = pipeline.NamedBarrier(3, warpgroup_threads + warp_threads) + sW_ready_nbar = pipeline.NamedBarrier(2, warpgroup_threads + warp_threads) + gk_load_nbar = pipeline.NamedBarrier(5, warpgroup_threads) + + elect_one = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + wg_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads) + warp_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warp_threads) + + mma6_done_mbar = smem.allocate_array(cutlass.Int64, 1) + gmem_done_mbar = smem.allocate_array(cutlass.Int64, 1) + state_ready_mbar = smem.allocate_array(cutlass.Int64, 1) + final_state_done_mbar = smem.allocate_array(cutlass.Int64, 1) + + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mma6_done_mbar, warp_threads) + cute.arch.mbarrier_init(gmem_done_mbar, warpgroup_threads) + cute.arch.mbarrier_init(state_ready_mbar, warpgroup_threads) + cute.arch.mbarrier_init(final_state_done_mbar, warpgroup_threads) + + def _make_tma_pipe(total_byte_count, num_stages=1): + ptr = smem.allocate_array(cutlass.Int64, 2 * num_stages) + return pipeline.PipelineTmaUmma.create( + barrier_storage=ptr, + num_stages=num_stages, + producer_group=elect_one, + consumer_group=elect_one, + tx_count=total_byte_count // num_stages, + defer_sync=True, + ).make_participants() + + b_prod, b_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, b_sl), num_stages=1) + v_prod, v_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, v_sl), num_stages=1) + q_prod, q_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, q_sl), num_stages=1) + a_prod, a_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, a_sl), num_stages=1) + aqc_prod, aqc_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, aqc_sl), num_stages=1) + kg_prod, kg_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, kg_sl), num_stages=1) + + def _make_umma_pipe(): + ptr = smem.allocate_array(cutlass.Int64, 2) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=ptr, + num_stages=1, + producer_group=elect_one, + consumer_group=wg_coop, + defer_sync=True, + ).make_participants() + + w_prod, w_cons = _make_umma_pipe() + nv_prod, nv_cons = _make_umma_pipe() + o_prod, o_cons = _make_umma_pipe() + + o_store_mbar = smem.allocate_array(cutlass.Int64, 2) + o_store_prod, o_store_cons = pipeline.PipelineAsync.create( + barrier_storage=o_store_mbar, + num_stages=1, + producer_group=wg_coop, + consumer_group=warp_coop, + defer_sync=True, + ).make_participants() + + cute.arch.sync_threads() + + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Int32, 16, tmem_smem) + + tCtW_shape = tiled_mma_kmn.partition_shape_C((M, N)) + tCtW_fake = tiled_mma_kmn.make_fragment_C(tCtW_shape) + tCtW = cute.make_tensor(cute.recast_ptr(tmem_ptr + 0, dtype=acc_dtype), tCtW_fake.layout) + + tCtNV_shape = tiled_mma_kmn.partition_shape_C((M, N)) + tCtNV_fake = tiled_mma_kmn.make_fragment_C(tCtNV_shape) + tCtNV = cute.make_tensor(cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtNV_fake.layout) + tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + 384, dtype=acc_dtype), tCtNV_fake.layout) + + tCtS_shape = tiled_mma_mn_mn.partition_shape_C((M6, N6)) + tCtS_fake = tiled_mma_mn_mn.make_fragment_C(tCtS_shape) + tCtState = cute.make_tensor(cute.recast_ptr(tmem_ptr + 256, dtype=acc_dtype), tCtS_fake.layout) + + # PDL: setup is done; now wait for upstream akk_inv to commit before reading + # gmem inputs (TMA loads, gk_last_exp reads). PDL allows the entire setup phase + # above (~tmem alloc, smem layout, pipeline init) to overlap with akk_inv's tail. + cute.arch.griddepcontrol_wait() + + # ==== WG1: State readout + decay ==== + if warpgroup_idx == STATE_WG: + cId_128 = cute.make_identity_tensor((M6, N6)) + tCtState_mn = transform_partitioned_tensor_layout(tCtState) + + atom_state_t2r = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), acc_dtype) + tiled_state_t2r = tcgen05.make_tmem_copy(atom_state_t2r, tCtState[(None, None), 0, 0]) + thr_state_t2r = tiled_state_t2r.get_slice(warpgroup_tidx) + tTR_tCtState = thr_state_t2r.partition_S(tCtState_mn) + tTR_tCcState = thr_state_t2r.partition_D(cId_128) + tRrState = cute.make_rmem_tensor_like(tTR_tCcState, acc_dtype) + + atom_state_r2t = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition(32)), acc_dtype) + tiled_state_r2t = tcgen05.make_tmem_copy(atom_state_r2t, tCtState[(None, None), 0, 0]) + thr_state_r2t = tiled_state_r2t.get_slice(warpgroup_tidx) + tRT_tCtState = thr_state_r2t.partition_D(tCtState_mn) + + atom_state_g2r = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32 + ) + tiled_state_g2r = cute.make_tiled_copy_S(atom_state_g2r, tiled_state_r2t) + thr_state_g2r = tiled_state_g2r.get_slice(warpgroup_tidx) + + atom_state_r2g = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32 + ) + tiled_state_r2g = cute.make_tiled_copy_D(atom_state_r2g, tiled_state_t2r) + thr_state_r2g = tiled_state_r2g.get_slice(warpgroup_tidx) + + atom_state_r2s = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), mma_dtype, num_bits_per_copy=128 + ) + tiled_state_r2s = cute.make_tiled_copy_D(atom_state_r2s, tiled_state_t2r) + thr_state_r2s = tiled_state_r2s.get_slice(warpgroup_tidx) + sST_vk_view = transform_partitioned_tensor_layout(sST) + sST_kv_view = cute.make_tensor( + sST.iterator, cute.select(sST_vk_view.layout, mode=[1, 0, 2]) + ) + tCsState_inp = thr_state_r2s.partition_D(sST_kv_view) + tRrState_bf16 = cute.make_rmem_tensor_like(tTR_tCcState, mma_dtype) + tCrState_bf16 = tiled_state_r2s.retile(tRrState_bf16) + + sGk = cute.make_tensor(sGk_buf, cute.make_layout(((N,), (1,)))) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + + while work.is_valid_tile: + batch_idx, head_idx, _ = work.tile_idx + state_idx = state_indices[batch_idx] + batch_start = cu_seqlens[batch_idx] + batch_end = cu_seqlens[batch_idx + 1] + num_chunks = cute.ceil_div(batch_end - batch_start, M) + # chunk_base from cumulative chunk-offsets array (handles non-64-aligned varlen). + # batch_start // M only works when all seq lengths are multiples of M. + chunk_base = chunk_offsets[batch_idx] + + state_vk = mS_fp32[state_idx, head_idx, None, None] + state_kv = cute.make_tensor( + state_vk.iterator, + cute.make_layout((M6, N6), stride=(1, M6)), + ) + gS_init = cute.flat_divide(state_kv, (M6, N6))[None, None, 0, 0] + tGR_tCgState_in = thr_state_g2r.partition_S(gS_init) + tGR_tCrState_in = thr_state_g2r.retile(tRrState) + cute.copy(tiled_state_g2r, tGR_tCgState_in, tGR_tCrState_in) + num_state_subs = tRrState.shape[2] + for sub in cutlass.range(num_state_subs): + cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub]) + cute.arch.fence_view_async_tmem_store() + + sGk[warpgroup_tidx] = cutlass.Float32(mGkLastExp[chunk_base, head_idx, warpgroup_tidx]) + gk_load_nbar.arrive_and_wait() + + if cutlass.const_expr(snapshot_rows is not None): + snapshot_row = snapshot_rows[batch_idx] + snapshot_chunk = snapshot_chunks[batch_idx] + + for chunk_c in cutlass.range(num_chunks): + num_state_subs = tRrState.shape[2] + _sub_tile_size = cute.size(tRrState.shape[0]) + + if chunk_c > 0: + cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2) + gk_load_nbar.arrive_and_wait() + + cute.copy(tiled_state_t2r, tTR_tCtState, tRrState) + tRrState_bf16.store(tRrState.load().to(mma_dtype)) + cute.copy(tiled_state_r2s, tCrState_bf16, tCsState_inp[None, None, None, 0]) + cute.arch.fence_view_async_shared() + cute.arch.mbarrier_arrive(gmem_done_mbar) + + if cutlass.const_expr(snapshot_rows is not None): + # tRrState still holds the undecayed state at the top of + # this chunk: the state after the first 64 * chunk_c tokens. + if chunk_c == snapshot_chunk: + cute.arch.fence_view_async_tmem_load() + snap_vk = mS_fp32[snapshot_row, head_idx, None, None] + snap_kv = cute.make_tensor( + snap_vk.iterator, + cute.make_layout((M6, N6), stride=(1, M6)), + ) + gS_snap = cute.flat_divide(snap_kv, (M6, N6))[None, None, 0, 0] + tGR_tCgState_snap = thr_state_r2g.partition_D(gS_snap) + tGR_tCrState_snap = thr_state_r2g.retile(tRrState) + for sub in cutlass.range(num_state_subs): + cute.copy( + tiled_state_r2g, + tGR_tCrState_snap[None, 0, sub], + tGR_tCgState_snap[None, 0, sub], + ) + + for sub in cutlass.range(num_state_subs): + cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) + cute.arch.fence_view_async_tmem_load() + for sub in cutlass.range(num_state_subs): + for i in cutlass.range(_sub_tile_size): + coord = tTR_tCcState[i, 0, sub] + # FIX: probe shows kernel state TMEM[m=K, n=V] (m corresponds to K-dim); + # OLD K4 / fla wants K-axis decay (state[k, v] *= gk[k]) → use coord[0]=m=K-idx. + k_idx = coord[0] + gk_val = cutlass.Float32(sGk[k_idx]) + tRrState[i, 0, sub] = tRrState[i, 0, sub] * gk_val + cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub]) + cute.arch.fence_view_async_tmem_store() + cute.arch.mbarrier_arrive(state_ready_mbar) + + if chunk_c + 1 < num_chunks: + sGk[warpgroup_tidx] = cutlass.Float32( + mGkLastExp[chunk_base + chunk_c + 1, head_idx, warpgroup_tidx] + ) + global_chunk = global_chunk + 1 + + cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2) + num_state_subs_final = tRrState.shape[2] + for sub in cutlass.range(num_state_subs_final): + cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) + cute.arch.fence_view_async_tmem_load() + state_vk = mS_fp32[state_idx, head_idx, None, None] + state_kv = cute.make_tensor( + state_vk.iterator, + cute.make_layout((M6, N6), stride=(1, M6)), + ) + gS_out = cute.flat_divide(state_kv, (M6, N6))[None, None, 0, 0] + tGR_tCgState_out = thr_state_r2g.partition_D(gS_out) + tGR_tCrState_out = thr_state_r2g.retile(tRrState) + for sub in cutlass.range(num_state_subs_final): + cute.copy( + tiled_state_r2g, + tGR_tCrState_out[None, 0, sub], + tGR_tCgState_out[None, 0, sub], + ) + + cute.arch.mbarrier_arrive(final_state_done_mbar) + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== MMA warp (warp 0): 6 MMAs ==== + elif warp_idx == MMA_WARP: + fA_kmn = thr_kmn.make_fragment_A(sA) + fB_ks = thr_kmn.make_fragment_B(sB) + fB_v = thr_kmn.make_fragment_B(sV) + fB_s = thr_kmn.make_fragment_B(sST) + fA_q = thr_kmn.make_fragment_A(sQ) + fA_aqc = thr_kmn.make_fragment_A(sAQC) + fB_nv = thr_kmn.make_fragment_B(sNV_b) + fA_kg = thr_mn.make_fragment_A(sKG_a) + fB_nv_mn = thr_mn.make_fragment_B(sNV_b_mn) + fA_w = thr_kmn.make_fragment_A(sO) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + tile_count = Int32(0) + + while work.is_valid_tile: + batch_idx_m, head_idx_m, _ = work.tile_idx + num_chunks_m = cute.ceil_div(cu_seqlens[batch_idx_m + 1] - cu_seqlens[batch_idx_m], M) + + for chunk_c in cutlass.range(num_chunks_m): + c_phase = global_chunk % 2 + + ah = a_cons.wait_and_advance() + bh = b_cons.wait_and_advance() + w_h = w_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sB.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtW, + fA_kmn[dice + (0,)][None, None, k], + fB_ks[dice + (bh.index,)][None, None, k], + tCtW, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + bh.release() + w_h.commit() + + vh = v_cons.wait_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sV.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_kmn[dice + (0,)][None, None, k], + fB_v[dice + (vh.index,)][None, None, k], + tCtNV, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + vh.release() + ah.release() + + cute.arch.mbarrier_wait(gmem_done_mbar, phase=c_phase) + sW_ready_nbar.arrive_and_wait() + + nv_h = nv_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_w[dice + (0,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtNV, + ) + nv_h.commit() + + qh = q_cons.wait_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_q[dice + (qh.index,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtO, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + qh.release() + + sNV_ready_nbar.arrive_and_wait() + + aqch = aqc_cons.wait_and_advance() + o_h = o_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sNV_b.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_aqc[dice + (0,)][None, None, k], + fB_nv[dice + (0,)][None, None, k], + tCtO, + ) + o_h.commit() + aqch.release() + + cute.arch.mbarrier_wait(state_ready_mbar, phase=c_phase) + kgh = kg_cons.wait_and_advance() + tiled_mma_mn_mn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sKG_a.shape[2])): + cute.gemm( + tiled_mma_mn_mn, + tCtState, + fA_kg[dice + (0,)][None, None, k], + fB_nv_mn[dice + (0,)][None, None, k], + tCtState, + ) + kgh.release() + + w_prod.tail() + nv_prod.tail() + o_prod.tail() + tcgen05.commit(mma6_done_mbar) + global_chunk = global_chunk + 1 + + cute.arch.mbarrier_wait(final_state_done_mbar, phase=tile_count % 2) + tile_count = tile_count + 1 + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, 512) + + # ==== TMA warp (warp 2) ==== + elif warp_idx == TMA_WARP: + cta_layout = cute.make_layout(1) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + + if work.is_valid_tile: + tensormap_manager.init_tensormap_from_atom(tma_a[0], tm_a_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_b[0], tm_b_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_v[0], tm_v_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_q[0], tm_q_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_aqc[0], tm_aqc_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_kg[0], tm_kg_ptr, TMA_WARP) + tensormap_manager.fence_tensormap_initialization() + + while work.is_valid_tile: + batch_idx_t, head_idx_t, _ = work.tile_idx + batch_start_t = cu_seqlens[batch_idx_t] + batch_end_t = cu_seqlens[batch_idx_t + 1] + num_chunks_t = cute.ceil_div(batch_end_t - batch_start_t, M) + + bounded_a = cute.make_tensor( + mA.iterator, + cute.make_layout( + (batch_end_t, mA.shape[1], mA.shape[2]), + stride=(mA.stride[0], mA.stride[1], mA.stride[2]), + ), + ) + bounded_b = cute.make_tensor( + mB.iterator, + cute.make_layout( + (mB.shape[0], batch_end_t, mB.shape[2]), + stride=(mB.stride[0], mB.stride[1], mB.stride[2]), + ), + ) + bounded_v = cute.make_tensor( + mV_g.iterator, + cute.make_layout( + (mV_g.shape[0], batch_end_t, mV_g.shape[2]), + stride=(mV_g.stride[0], mV_g.stride[1], mV_g.stride[2]), + ), + ) + bounded_q = cute.make_tensor( + mQ.iterator, + cute.make_layout( + (batch_end_t, mQ.shape[1], mQ.shape[2]), + stride=(mQ.stride[0], mQ.stride[1], mQ.stride[2]), + ), + ) + bounded_aqc = cute.make_tensor( + mAQC.iterator, + cute.make_layout( + (batch_end_t, mAQC.shape[1], mAQC.shape[2]), + stride=(mAQC.stride[0], mAQC.stride[1], mAQC.stride[2]), + ), + ) + bounded_kg = cute.make_tensor( + mKG.iterator, + cute.make_layout( + (mKG.shape[0], batch_end_t, mKG.shape[2]), + stride=(mKG.stride[0], mKG.stride[1], mKG.stride[2]), + ), + ) + tensormap_manager.update_tensormap( + (bounded_a, bounded_b, bounded_v, bounded_q, bounded_aqc, bounded_kg), + (tma_a[0], tma_b[0], tma_v[0], tma_q[0], tma_aqc[0], tma_kg[0]), + (tm_a_ptr, tm_b_ptr, tm_v_ptr, tm_q_ptr, tm_aqc_ptr, tm_kg_ptr), + TMA_WARP, + (None, None, None, None, None, None), + ) + + for chunk_c in cutlass.range(num_chunks_t): + chunk_offset = batch_start_t + chunk_c * M + + # AB: A-operand (T, K, H) → domain_offset on (tokens, 0) + mA_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_a[1][None, None, head_idx_t] + ) + gA = cute.flat_divide(mA_c, (M, K)) + tCgA = thr_kmn.partition_A(gA) + tAsA, tAgA = cpasync.tma_partition( + tma_a[0], + 0, + cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + ah = a_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_a_ptr) + cute.copy( + tma_a[0], + tAgA[(None, 0, 0)], + tAsA[(None, ah.index)], + tma_bar_ptr=ah.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_a_ptr, cute.AddressSpace.generic + ), + ) + + # KS: B-operand (N, T, H) → domain_offset on (0, tokens) + mB_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_b[1][None, None, head_idx_t] + ) + gB_c = cute.flat_divide(mB_c, (N, K)) + tCgB = thr_kmn.partition_B(gB_c) + tBsB, tBgB = cpasync.tma_partition( + tma_b[0], + 0, + cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + bh = b_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_b_ptr) + cute.copy( + tma_b[0], + tBgB[(None, 0, 0)], + tBsB[(None, bh.index)], + tma_bar_ptr=bh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_b_ptr, cute.AddressSpace.generic + ), + ) + + # V: B-operand (N, T, H) → domain_offset on (0, tokens) + mV_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_v[1][None, None, head_idx_t] + ) + gV_c = cute.flat_divide(mV_c, (N, K)) + tCgV = thr_kmn.partition_B(gV_c) + tBsV, tBgV = cpasync.tma_partition( + tma_v[0], + 0, + cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tCgV, 0, 3), + ) + vh = v_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_v_ptr) + cute.copy( + tma_v[0], + tBgV[(None, 0, 0)], + tBsV[(None, vh.index)], + tma_bar_ptr=vh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_v_ptr, cute.AddressSpace.generic + ), + ) + + # QS: A-operand (T, N, H) → domain_offset on (tokens, 0) + mQ_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_q[1][None, None, head_idx_t] + ) + gQ_c = cute.flat_divide(mQ_c, (M, K3)) + tCgQ = thr_kmn.partition_A(gQ_c) + tAsQ, tAgQ = cpasync.tma_partition( + tma_q[0], + 0, + cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgQ, 0, 3), + ) + qh = q_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_q_ptr) + cute.copy( + tma_q[0], + tAgQ[(None, 0, 0)], + tAsQ[(None, qh.index)], + tma_bar_ptr=qh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_q_ptr, cute.AddressSpace.generic + ), + ) + + # AQC: A-operand (T, K, H) → domain_offset on (tokens, 0) + mAQC_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_aqc[1][None, None, head_idx_t] + ) + gAQC_c = cute.flat_divide(mAQC_c, (M, K)) + tCgAQC = thr_kmn.partition_A(gAQC_c) + tAsAQC, tAgAQC = cpasync.tma_partition( + tma_aqc[0], + 0, + cta_layout, + cute.group_modes(sAQC, 0, 3), + cute.group_modes(tCgAQC, 0, 3), + ) + aqch = aqc_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_aqc_ptr) + cute.copy( + tma_aqc[0], + tAgAQC[(None, 0, 0)], + tAsAQC[(None, aqch.index)], + tma_bar_ptr=aqch.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_aqc_ptr, cute.AddressSpace.generic + ), + ) + + # KG: B-operand (N6, T, H) → domain_offset on (0, tokens) + mKG_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_kg[1][None, None, head_idx_t] + ) + gKG_c = cute.flat_divide(mKG_c, (N6, K6)) + tCgKG = thr_mn.partition_B(gKG_c) + tBsKG, tBgKG = cpasync.tma_partition( + tma_kg[0], + 0, + cta_layout, + cute.group_modes(sKG, 0, 3), + cute.group_modes(tCgKG, 0, 3), + ) + kgh = kg_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_kg_ptr) + cute.copy( + tma_kg[0], + tBgKG[(None, 0, 0)], + tBsKG[(None, kgh.index)], + tma_bar_ptr=kgh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_kg_ptr, cute.AddressSpace.generic + ), + ) + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== Warp 1: O TMA store ==== + elif warp_idx == O_STORE_WARP: + cta_layout_o = cute.make_layout(1) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + + if work.is_valid_tile: + tensormap_manager.init_tensormap_from_atom(tma_o[0], tm_o_ptr, O_STORE_WARP) + tensormap_manager.fence_tensormap_initialization() + + while work.is_valid_tile: + batch_idx_o, head_idx_o, _ = work.tile_idx + batch_start_o = cu_seqlens[batch_idx_o] + batch_end_o = cu_seqlens[batch_idx_o + 1] + num_chunks_o = cute.ceil_div(batch_end_o - batch_start_o, M) + + bounded_o = cute.make_tensor( + mO.iterator, + cute.make_layout( + (batch_end_o, mO.shape[1], mO.shape[2]), + stride=(mO.stride[0], mO.stride[1], mO.stride[2]), + ), + ) + tensormap_manager.update_tensormap( + (bounded_o,), (tma_o[0],), (tm_o_ptr,), O_STORE_WARP, (None,) + ) + tensormap_manager.fence_tensormap_update(tm_o_ptr) + + for chunk_c in cutlass.range(num_chunks_o): + os_h = o_store_cons.wait_and_advance() + chunk_offset_o = batch_start_o + chunk_c * M + mO_c = cute.domain_offset( + (chunk_offset_o, Int32(0)), tma_o[1][None, None, head_idx_o] + ) + gOo = cute.flat_divide(mO_c, (M, N)) + sOt, gOt = cpasync.tma_partition( + tma_o[0], + 0, + cta_layout_o, + cute.group_modes(sO_out_st, 0, 2), + cute.group_modes(gOo, 0, 2), + ) + cute.copy( + tma_o[0], + sOt[None], + gOt[(None, 0, 0)], + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_o_ptr, cute.AddressSpace.generic + ), + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + os_h.release() + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== WG2: W/NV/O readout (SS-mode, no Phase 2) ==== + elif warpgroup_idx == READOUT_WG: + tCtW_mn = transform_partitioned_tensor_layout(tCtW) + tCtNV_mn = transform_partitioned_tensor_layout(tCtNV) + tCtO_mn = transform_partitioned_tensor_layout(tCtO) + + atom_t2r = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition(1)), acc_dtype) + tiled_t2r = tcgen05.make_tmem_copy(atom_t2r, tCtW[(None, None), 0, 0]) + thr_t2r = tiled_t2r.get_slice(warpgroup_tidx) + + tTR_W = thr_t2r.partition_S(tCtW_mn) + tTR_NV = thr_t2r.partition_S(tCtNV_mn) + tTR_O = thr_t2r.partition_S(tCtO_mn) + + atom_r2s_k = sm100_utils.get_smem_store_op( + utils.LayoutEnum.ROW_MAJOR, mma_dtype, acc_dtype, tiled_t2r + ) + tiled_r2s_k = cute.make_tiled_copy_D(atom_r2s_k, tiled_t2r) + thr_r2s_k = tiled_r2s_k.get_slice(warpgroup_tidx) + tCsO = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO)) + tCsO_out = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO_out)) + tCsNV = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sNV)) + + cId = cute.make_identity_tensor((M, N)) + tTR_cId = thr_t2r.partition_D(cId) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + + while work.is_valid_tile: + batch_idx_r, head_idx_r, _ = work.tile_idx + num_chunks_r = cute.ceil_div(cu_seqlens[batch_idx_r + 1] - cu_seqlens[batch_idx_r], M) + + for chunk_c in cutlass.range(num_chunks_r): + c_phase = global_chunk % 2 + tRrR = cute.make_rmem_tensor_like(tTR_cId, acc_dtype) + tRrR_out = cute.make_rmem_tensor_like(tRrR, mma_dtype) + tCrR_k = tiled_r2s_k.retile(tRrR_out) + num_subs = tRrR.shape[2] + + wh = w_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_W[None, 0, sub], tRrR[None, 0, sub]) + # Negate W in its readout so MMA3 accumulates NV = U - W @ S + # (FLA: b_v = u - w @ h). Folded into the bf16 cast — free. + tRrR_out[None, 0, sub].store((-(tRrR[None, 0, sub].load())).to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + wh.release() + cute.arch.fence_view_async_shared() + sW_ready_nbar.arrive_and_wait() + + nvh = nv_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_NV[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsNV[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + nvh.release() + cute.arch.fence_view_async_shared() + sNV_ready_nbar.arrive_and_wait() + + os_h = o_store_prod.acquire_and_advance() + + oh = o_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_O[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO_out[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + oh.release() + cute.arch.fence_view_async_shared() + os_h.commit() + global_chunk = global_chunk + 1 + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + o_store_prod.tail() + + +def make_host_fn(num_sm=148, has_snapshot=False): + """Create one K4 host function for all runtime context batch sizes. + + ``num_seqs`` is a runtime scalar. ``H`` and the token extent come from + dynamic tensor layouts, while ``num_sm`` is the fixed persistent-grid + ceiling. The scheduler computes ``min(num_seqs * H, num_sm)`` at launch. + + With ``has_snapshot`` the host function takes two more int32 ``[num_seqs]`` + tensors after ``chunk_offsets``, ``snapshot_rows`` and ``snapshot_chunks`` + (see the module docstring). The default host function and its kernel carry + no snapshot operands. + """ + _num_sm = num_sm + + def launch( + a_raw: cute.Tensor, + b_raw: cute.Tensor, + v_raw: cute.Tensor, + q_raw: cute.Tensor, + aqc_raw: cute.Tensor, + kg_raw: cute.Tensor, + o_raw: cute.Tensor, + gk_last_exp: cute.Tensor, + s_fp32: cute.Tensor, + state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + snapshot_rows: Optional[cute.Tensor], + snapshot_chunks: Optional[cute.Tensor], + tm_workspace: cute.Tensor, + num_seqs: cutlass.Int32, + stream: cuda.CUstream, + ): + """Build the K4 operands and launch the kernel (traced by the host function).""" + # Raw tensors from from_dlpack have PyTorch layout (T, H, dim) + # stride (H*dim, dim, 1). Reshape to 3D for TMA: + # A-operands: (T, dim, H), stride (dim*H, 1, dim) + # B-operands: (dim, T, H), stride (1, dim*H, dim) + T_tok = a_raw.shape[0] + H = a_raw.shape[1] + + a = cute.make_tensor( + a_raw.iterator, + cute.make_layout( + (T_tok, a_raw.shape[2], H), + stride=(a_raw.stride[0], a_raw.stride[2], a_raw.stride[1]), + ), + ) + q = cute.make_tensor( + q_raw.iterator, + cute.make_layout( + (T_tok, q_raw.shape[2], H), + stride=(q_raw.stride[0], q_raw.stride[2], q_raw.stride[1]), + ), + ) + aqc = cute.make_tensor( + aqc_raw.iterator, + cute.make_layout( + (T_tok, aqc_raw.shape[2], H), + stride=(aqc_raw.stride[0], aqc_raw.stride[2], aqc_raw.stride[1]), + ), + ) + o_out = cute.make_tensor( + o_raw.iterator, + cute.make_layout( + (T_tok, o_raw.shape[2], H), + stride=(o_raw.stride[0], o_raw.stride[2], o_raw.stride[1]), + ), + ) + + b = cute.make_tensor( + b_raw.iterator, + cute.make_layout( + (b_raw.shape[2], T_tok, H), + stride=(b_raw.stride[2], b_raw.stride[0], b_raw.stride[1]), + ), + ) + v = cute.make_tensor( + v_raw.iterator, + cute.make_layout( + (v_raw.shape[2], T_tok, H), + stride=(v_raw.stride[2], v_raw.stride[0], v_raw.stride[1]), + ), + ) + kg = cute.make_tensor( + kg_raw.iterator, + cute.make_layout( + (kg_raw.shape[2], T_tok, H), + stride=(kg_raw.stride[2], kg_raw.stride[0], kg_raw.stride[1]), + ), + ) + + tile1 = (M, N, K) + tile3 = (M, N, K3) + tile6 = (M6, N6, K6) + + mma_kmn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.K, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M, N), + OperandSource.SMEM, + ) + mma_mn_mn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.MN, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M6, N6), + OperandSource.SMEM, + ) + sl_a = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_v = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_s = sm100_utils.make_smem_layout_b(mma_kmn, tile3, mma_dtype, 1) + sl_q = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_aqc = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_kg = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1) + + sl_readout_k = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_nv_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_nv_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1) + sl_kg_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1) + sl_nv_b_mn = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1) + + tma_ld = cpasync.CopyBulkTensorTileG2SOp() + + ta_a = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, a, cute.select(sl_a, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_b = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, b, cute.select(sl_b, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, v, cute.select(sl_v, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, q, cute.select(sl_q, mode=[0, 1, 2]), tile3, mma_kmn + ) + ta_aqc = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, aqc, cute.select(sl_aqc, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_kg = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, kg, cute.select(sl_kg, mode=[0, 1, 2]), tile6, mma_mn_mn + ) + sk = sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.K, out_dtype, (M, N)) + sl_store = cute.tile_to_shape( + sm100_utils.make_smem_layout_atom(sk, out_dtype), (M, N), order=(0, 1) + ) + tma_st = cpasync.CopyBulkTensorTileS2GOp() + ta_o = cpasync.make_tiled_tma_atom(tma_st, o_out, sl_store, (M, N)) + + scheduler_params = GDNTileSchedulerParams( + num_seqs=num_seqs, + num_q_heads=H, + num_v_heads=H, + is_GQA=False, + is_persistent=True, + ) + grid_shape = GDNTileScheduler.get_grid_shape(scheduler_params, _num_sm) + + k4_persistent_kernel( + mma_kmn, + mma_mn_mn, + ta_a, + sl_a, + ta_b, + sl_b, + ta_v, + sl_v, + sl_s, + ta_q, + sl_q, + ta_aqc, + sl_aqc, + ta_kg, + sl_kg, + ta_o, + sl_store, + sl_readout_k, + sl_nv_b, + sl_nv_a, + sl_kg_a, + sl_nv_b_mn, + gk_last_exp, + s_fp32, + state_indices, + cu_seqlens, + chunk_offsets, + snapshot_rows, + snapshot_chunks, + a, + b, + v, + q, + aqc, + kg, + o_out, + tm_workspace, + scheduler_params, + ).launch(grid=grid_shape, block=(threads_per_cta, 1, 1), use_pdl=True, stream=stream) + + if has_snapshot: + + @cute.jit + def host_fn_snapshot( + a_raw: cute.Tensor, + b_raw: cute.Tensor, + v_raw: cute.Tensor, + q_raw: cute.Tensor, + aqc_raw: cute.Tensor, + kg_raw: cute.Tensor, + o_raw: cute.Tensor, + gk_last_exp: cute.Tensor, + s_fp32: cute.Tensor, + state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + snapshot_rows: cute.Tensor, + snapshot_chunks: cute.Tensor, + tm_workspace: cute.Tensor, + num_seqs: cutlass.Int32, + stream: cuda.CUstream, + ): + launch( + a_raw, + b_raw, + v_raw, + q_raw, + aqc_raw, + kg_raw, + o_raw, + gk_last_exp, + s_fp32, + state_indices, + cu_seqlens, + chunk_offsets, + snapshot_rows, + snapshot_chunks, + tm_workspace, + num_seqs, + stream, + ) + + return host_fn_snapshot + + @cute.jit + def host_fn( + a_raw: cute.Tensor, + b_raw: cute.Tensor, + v_raw: cute.Tensor, + q_raw: cute.Tensor, + aqc_raw: cute.Tensor, + kg_raw: cute.Tensor, + o_raw: cute.Tensor, + gk_last_exp: cute.Tensor, + s_fp32: cute.Tensor, + state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + tm_workspace: cute.Tensor, + num_seqs: cutlass.Int32, + # Launch stream — runtime argument; launching on the DSL default + # stream races with the executor's non-blocking execution stream. + # Same stream as akk_inv, so the PDL chain stays intact. + stream: cuda.CUstream, + ): + launch( + a_raw, + b_raw, + v_raw, + q_raw, + aqc_raw, + kg_raw, + o_raw, + gk_last_exp, + s_fp32, + state_indices, + cu_seqlens, + chunk_offsets, + None, + None, + tm_workspace, + num_seqs, + stream, + ) + + return host_fn diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode.py new file mode 100644 index 000000000000..b57a140d017f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode.py @@ -0,0 +1,2143 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL device kernel for KDA multi-token speculative verify (conv-MTP). + +Source-integrated from the ``KDA_decode_mtp`` kernel drop ("readable KDA +conv-MTP") and restructured for the latency-bound production regime (one CTA +per (head, request), or a cluster pair of CTAs while the batch is small; TEP8 +launches only ``H * N = 24`` tiles on a 212-SM Rubin). Specialized for the +benchmarked contract: conv enabled, no bias, optional BF16 gated RMSNorm +output, lower_bound gate, Q/K L2 norm, beta sigmoid, W=4, K == V == 128, +HV == H, ``TILE_V == V`` (one V pass), 512 threads. + +Per request ``n`` the kernel processes an accepted count plus ``1 + +NUM_SPEC`` steps: it first *replays* the accepted draft tokens from the +``k/g/v/beta`` caches (raw conv inputs from the extended tail columns of +``cs_q/cs_k/cs_v``), then processes the ``1 + NUM_SPEC`` new tokens. The +recurrent state and the base conv windows are committed in place after the +first new (golden) token; the new spec tokens are cached for the next +round's replay. The pool invariant is therefore "state after the last +golden token; accepted drafts pending in the replay caches". + +Structure (all stage boundaries are CTA barriers): + +1. **Load phase.** Every warp issues its global loads up front, so the whole + CTA waits for one DRAM round trip: the recurrent state ``h0`` through a TMA + bulk load into swizzled smem boxes (consumed by the tensor-core stage), the + raw conv inputs of every token (``x_q/x_k/x_v/g/beta`` for new tokens, the + replay caches and ``cs_*`` tail columns for accepted drafts) into smem + staging buffers, the conv weights, and the gated-RMSNorm operands for the + epilogue. +2. **Per-token-parallel precompute.** Warp ``j`` computes the Q/K conv + SiLU + + L2 norm of new token ``j`` and warp ``num_warps/2 + j`` its gate, beta and + V conv, reading the width-4 windows from the staged raw sequence (rows + ``0..W-2`` hold the committed window). Replay tokens were written to + ``sQ/sK/sG/sBeta/sVall`` directly by the load phase. Conv-window commits + and replay-cache writes are done by the warp that owns the token. +3. **Recurrence** in chunk (WY) form on the tcgen05 tensor cores, tf32 with + 3xTF32 (hi/lo splits) on every product that feeds the state: + + * GEMM1 ``[U0; O0] = S0 x [K~; Q~]`` -- S0 hi/lo staged into TMEM from the + smem boxes, ``K~ = k e^{G_t}``, ``Q~ = q e^{G_t}`` (``G_t`` = cumulative + gate of the chunk); + * GEMM3 ``L, B`` -- the T x T coupling matrices of the chunk (decays + factored around the chunk middle and clamped at ``e^87``); + * ``V' = (I + diag(beta) L)^-1 diag(beta) (V - U0)`` by forward + substitution (one warpgroup, one value row per thread); + * GEMM2 ``S_c = S0 e^{G_c} + V'^T Kbar`` -- accumulator seeded from the + registers holding S0, issued as two column halves so the first half can + start storing; the output combination ``O = O0 + B V'`` rides along as 16 + extra columns of the first half. +4. **Store / epilogue.** The committed state goes through swizzled smem boxes + and a TMA bulk tensor store; the gated RMSNorm and MXFP8 quantization read + their operands from smem and run on the warps that own no state rows while + the store drains. + +With ``SPLIT_V == 2`` two CTAs per (head, request) run as a cluster: both +recompute the k-side and the outputs, each seeds and stores half of the state +rows (the SM -> L2 drain of the 64 KB state is the tail of the kernel). The pair +must be co-scheduled because the state is updated in place, and a cluster +barrier separates every CTA's S0 load from the pair's stores. The TMA views of +the state follow the tensor's own strides (the runtime's per-layer state cache +is a strided slice of a larger buffer). The state may be fp32 or bf16; a bf16 +state is exact in tf32 (the ``S0_lo`` sweep is skipped) and the committed state +is rounded to nearest on the way out. + +Measured on Rubin (nsys, state L2-resident) against the previous CUDA-core +kernel at H=12/num_spec=7 with mixed acceptance: N=2 8.61 us -> 7.07 us, +N=16 11.78 us -> 9.06 us; H=96/N=16 91.1 us -> 74.2 us. The committed state +is within 2e-6 of the fp32 kernel (3xTF32), the FP8 output codes differ in +~0.2 % of the elements by one step; the conv pools and replay caches are +bit-exact. + +``BF16_MMA`` (``TRTLLM_KDA_MTP_BF16_MATH=1``) trades that precision for the +bf16-operand / fp32-accumulate scheme of the prefill kernel: S0 packed as bf16 +in TMEM (64 columns), bf16 K-major smem tiles and K = 16 per instruction (18 +MMAs per CTA instead of 76). Gates, cumulative gates, exponents, reductions +and the state accumulator stay fp32; the committed-state error vs fp32 grows to +~2e-2, the conv pools and replay caches stay bit-exact between the two paths. +""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import Int64 +from cutlass.cutlass_dsl import T, dsl_user_op + +# 16 warps: the full 128-row recurrent state is processed in one pass (eight +# rows per warp) and the per-token precompute runs one (token, role) per warp. +NUM_THREADS = 512 +# Two CTAs per (head, request) while HV * N <= SPLIT_MAX_TILES (up to 128 CTAs: one wave, the +# halved state store pays off); larger batches run one CTA per tile. +SPLIT_MAX_TILES = 64 +TILE_K = 128 +# Chunk length padding (MMA N granularity); NUM_SPEC <= 7 -> T_loop <= 15. +T_PAD = 16 +# TMEM columns: D4 (16) + D2 halves (2 x 64) + S0 hi (128) + S0 lo (128) + D1 (48) + D3 (16) +# = 464 of the 512 allocated. +TMEM_COLS = 512 +# One 114 KB scratch region hosts, at different times, the TMA boxes of S0 / the committed +# state, the swizzled MMA operand tiles and (bf16 output path) the per-warp 32x32 transposes. +# Offsets in floats, all multiples of 256 (1024 B) for the 128B-swizzle atoms. +SCRATCH_FLOATS = 29184 +OFF_B4HI = 10240 # 16 x 16 Bcoef hi, immediately followed by Kbar hi -> N=80 GEMM2 half-0 tile +OFF_B4LO = 12544 # 16 x 16 Bcoef lo, immediately followed by Kbar lo +OFF_B1K = 27136 # 16 x 128 K~hi (duplicate, N=16 tile for the S0_lo sweep) +OFF_B1 = 0 # 48 x 128 [K~hi; Q~; K~lo] +OFF_A2HI = 6144 # 128 x 16 V'^T hi +OFF_A2LO = 8192 # 128 x 16 V'^T lo +OFF_B2HI = 10496 # 128 x 16 Kbar hi +OFF_B2LO = 12800 # 128 x 16 Kbar lo +OFF_A3 = 14848 # 64 x 128 [Kf; Qf hi; Qf lo; -] +OFF_B3HI = 23040 # 16 x 128 Kb hi +OFF_B3LO = 25088 # 16 x 128 Kb lo +TP_STRIDE = 17 # padded row stride of the transposed (v|k, s) staging buffers +# bf16 operand tiles (BF16_MMA): offsets in bf16 elements inside the same scratch region, all +# multiples of 512 (1024 B). [Bcoef; Kbar] is one contiguous SW32 tile (Bcoef rows 0..15, then the +# 128 Kbar rows); the GEMM2 half-1 B view starts at Kbar row 64. +OFFH_B1 = 0 # 32 x 128 [K~; Q~] +OFFH_A3 = 4096 # 64 x 128 [Kf; Qf; 0; -] +OFFH_B3 = 12288 # 16 x 128 Kb +OFFH_A2 = 14336 # 128 x 16 V'^T +OFFH_B4 = 16384 # 16 x 16 Bcoef, immediately followed by +OFFH_B2 = 16640 # 128 x 16 Kbar +XP_STRIDE = 33 # padded row stride of the 32x32 transpose blocks + +# Compact layout (bf16 state, fused output norm, 3xTF32, one CTA per tile, NUM_SPEC <= 3): the +# per-CTA footprint is cut to <= 113 KB of smem and 256 TMEM columns so that TWO CTAs +# co-reside on an SM. At large batch (HV * N tiles >> SM count) the kernel is a +# sequence of waves of latency-bound CTAs; with one CTA per SM the load phase of a +# CTA never overlaps the tensor-core chain of another. The arithmetic (every +# instruction that produces a value) is unchanged; only buffer placement differs: +# * one 90 KB region is time-shared: phase A holds the S0 TMA box and the +# load/precompute staging (raw conv rows, raw gates, conv weights, dt_bias, sG); +# phase B the GEMM1/GEMM3 operands and the B4|B2 tiles; phase C (after GEMM1/3 +# completed) the A2 tiles, sLB and sVpT in the dead B1/A3 space; phase D the +# committed-state box in the dead A3 space; +# * Kbar goes straight into the B2 tiles (no sKbT staging); +# * S0 hi (the TMEM A operand of GEMM1) shares its 128 columns with the GEMM2 +# accumulator, which is therefore seeded only after GEMM1 completed. +# The host launches it with min_blocks_per_mp=2 (<= 64 registers per thread: two 512-thread CTAs +# fill the register file). +COMPACT_TMEM_COLS = 256 +# Named barrier ids (0 is the CTA-wide barrier). An SM reserves (max id + 1) named barriers per +# CTA; on SM107 ids up to 8 (9 barriers) allow only one CTA per SM, so the ids stay <= 4. +BAR_KBAR = 1 # Kbar / B2 tiles staged -> warp 0 issues GEMM3 +BAR_STORE_LO = 2 # committed-state boxes staged, threads 0..255 +BAR_STORE_HI = 3 # committed-state boxes staged, threads 256..511 +BAR_TMEM_FREE = 4 # last TMEM reads done -> warp 0 frees the columns +C_OFF_B1 = 0 # 48 x 128 +C_OFF_A3 = 6144 # 64 x 128 +C_OFF_B3HI = 14336 # 16 x 128 +C_OFF_B3LO = 16384 # 16 x 128 +C_OFF_B4HI = 18432 # 16 x 16, followed by B2 hi (N = 80 GEMM2 half-0 tile) +C_OFF_B2HI = 18688 # 128 x 16 +C_OFF_B4LO = 20736 +C_OFF_B2LO = 20992 +C_REGION_FLOATS = 23040 +C_OFF_A2HI = 0 # phase C (B1 dead) +C_OFF_A2LO = 2048 +C_OFF_LB = 6144 # phase C (A3 dead): 64 x 16 +C_OFF_VPT = 7168 # phase C: V x TP_STRIDE +C_OFF_STBOX = 6144 # phase D: committed-state bf16 boxes (A3 / dead staging space) +C_OFF_STAGE = 8192 # phase A staging, after the bf16 S0 box + + +def swz128(n, k, rows): + """Element offset of (row n, col k) in a K-major 128B-swizzled fp32 tile with + `rows` rows and 128 columns (canonical UMMA layout, see probe_swizzle.py).""" + off = n * 32 + (k % 32) + (k // 32) * (rows * 32) + return off ^ (((off >> 5) & 7) << 2) + + +def swz64_k16(n, k): + """Element offset in a K-major 64B-swizzled fp32 tile with 16 columns.""" + off = n * 16 + k + return off ^ (((off >> 5) & 3) << 2) + + +def sw128_bf16(n, k, rows): + """Element offset of (n, k) in a K-major 128B-swizzled bf16 tile with `rows` rows and K = 128: + two K blocks of 64 elements (rows x 128 B each), 16 B chunks XORed with n % 8.""" + kc = k % 64 + return (k // 64) * (rows * 64) + n * 64 + ((((kc >> 3) ^ (n & 7)) << 3) | (kc & 7)) + + +def sw32_k16_bf16(n, k): + """Element offset of (n, k) in a K-major 32B-swizzled bf16 tile with 16 columns.""" + off = n * 16 + k + return off ^ (((off >> 6) & 1) << 3) + + +@dsl_user_op +def cvt_bf16x2(hi: cutlass.Float32, lo: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Float32: + """Pack two fp32 values into one bf16x2 word (round to nearest even, `lo` in the low half) + with a single cvt; returned as the fp32 bit pattern for the TMEM word stores.""" + w = cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + return w.bitcast(cutlass.Float32) + + +@dsl_user_op +def tcgen05_fence_before_thread_sync(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "tcgen05.fence::before_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def tcgen05_fence_after_thread_sync(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "tcgen05.fence::after_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def round_tf32(x: cutlass.Float32) -> cutlass.Float32: + """Round to the tf32 grid (rna on mantissa bit 13) so hi/lo splits are exact.""" + b = x.bitcast(cutlass.Int32) + b = (b + cutlass.Int32(0x1000)) & cutlass.Int32(-8192) + return b.bitcast(cutlass.Float32) + + +def round_bf16(x: cutlass.Float32) -> cutlass.Float32: + """Round to the bf16 grid (rna on mantissa bit 16, two integer ops like round_tf32); the + value stays an fp32 operand that is exact in tf32.""" + b = x.bitcast(cutlass.Int32) + b = (b + cutlass.Int32(0x8000)) & cutlass.Int32(-65536) + return b.bitcast(cutlass.Float32) + + +def round_hi(x: cutlass.Float32, BF16_MMA: bool) -> cutlass.Float32: + """MMA operand rounding: tf32 grid for the 3xTF32 path (so the hi/lo split is exact), bf16 + grid for the bf16 operand path (the conversion into the bf16 tiles is then exact).""" + if BF16_MMA: + return round_bf16(x) + return round_tf32(x) + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None) -> Int64: + """Read the SM global timer for optional in-kernel stage profiling.""" + return Int64( + llvm.inline_asm( + T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.kernel +def kda_decode_mtp_kernel( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + output_scale: cute.Tensor, + ht: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + smem_qk_layout: cute.Layout, + ssm_state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + num_accepted_tokens: cute.Tensor, + precompute_control: cute.Tensor, + TILE_V: cutlass.Constexpr[int], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + scale_leading_dim: cutlass.Constexpr[int], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + USE_SETMAXREG: cutlass.Constexpr[bool], + USE_REGULAR_METADATA: cutlass.Constexpr[bool], + USE_PACKED_TOKEN_LAYOUT: cutlass.Constexpr[bool], + USE_REG_Q_WEIGHTS: cutlass.Constexpr[bool], + USE_ZERO_ACCEPTED: cutlass.Constexpr[bool], + FUSE_PRECOMPUTE: cutlass.Constexpr[bool], + RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], + FUSE_OUTPUT_NORM: cutlass.Constexpr[bool], + QUANTIZE_OUTPUT: cutlass.Constexpr[bool], + stage_timing: cute.Tensor, + PROFILE_STAGES: cutlass.Constexpr[bool], + tiled_mma1: cute.TiledMma, + tiled_mma3: cute.TiledMma, + tiled_mma1k: cute.TiledMma, + smem_layout_b1: cute.ComposedLayout, + smem_layout_a2: cute.ComposedLayout, + smem_layout_a3: cute.ComposedLayout, + smem_layout_b3: cute.ComposedLayout, + smem_layout_b1k: cute.ComposedLayout, + tiled_mma4: cute.TiledMma, + tma_atom_t: cute.CopyAtom, + mT_tma: cute.Tensor, + box_layout: cute.ComposedLayout, + tma_atom_l: cute.CopyAtom, + mS_tma: cute.Tensor, + tiled_mma2h: cute.TiledMma, + smem_layout_b2h: cute.ComposedLayout, + box_layout_t: cute.ComposedLayout, + tiled_mma2q: cute.TiledMma, + smem_layout_b2q: cute.ComposedLayout, + SPLIT_V: cutlass.Constexpr[int], + BF16_MMA: cutlass.Constexpr[bool], + tiled_mma_f32a: cute.TiledMma, + COMPACT: cutlass.Constexpr[bool] = False, +): + """KDA MTP verify: smem precompute + tensor-core chunk recurrence. + + With ``PROFILE_STAGES=True``, ``stage_timing`` must be an int64 tensor with at + least ``SPLIT_V * HV * N * 64`` elements, indexed as + ``((vh * HV + i_hv) * grid_n + i_n) * 64``. With profiling off it is never + accessed and the host may pass any placeholder tensor. + + ``COMPACT`` selects the two-CTAs-per-SM buffer placement (see ``COMPACT_TMEM_COLS``). + """ + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + i_hv, i_n, vh = cute.arch.block_idx() + # Two CTAs per (head, request): CTA vh owns the state rows [vh * V/2, (vh + 1) * V/2). + i_h = i_hv + if cutlass.const_expr(PROFILE_STAGES): + t_stage0 = read_globaltimer() + # Pre-declare so the dynamic `run_precompute` branch below only + # reassigns (first assignment inside a dynamic branch is untraceable; + # the name has to exist before the branch). + t_stage1 = Int64(0) + tp_a = Int64(0) + tp_b = Int64(0) + tp_c = Int64(0) + tp_d = Int64(0) + tp_e = Int64(0) + tp_f = Int64(0) + tp_g = Int64(0) + tp_h = Int64(0) + tp_i = Int64(0) + tp_j = Int64(0) + tp_k = Int64(0) + tp_l = Int64(0) + tp_m = Int64(0) + if cutlass.const_expr(USE_REGULAR_METADATA): + bos = i_n * (2 * NUM_SPEC + 1) + eos = bos + (2 * NUM_SPEC + 1) + slot = i_n + elif cutlass.const_expr(USE_PACKED_TOKEN_LAYOUT): + # Runtime verification packs exactly 1 + NUM_SPEC new rows per + # request. Accepted drafts live in slot-indexed replay pools and do + # not occupy rows in the projection outputs. + bos = cutlass.Int32(0) + eos = cutlass.Int32(0) + slot = ssm_state_indices[i_n] + else: + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + slot = ssm_state_indices[i_n] + # V2 state views may place many layer pages between adjacent slots. Widen + # the slot before indexing: slot * state stride can exceed INT32_MAX. + slot = Int64(slot) + h0_idx = slot * HV + i_hv + hk_off = i_h * K + hv_off = i_hv * V + # Recurrent-state element type (fp32 or bf16). The TMA boxes have 128 B swizzled rows, + # i.e. BOX_COLS = 32 fp32 or 64 bf16 columns; the state tile is NUM_SLABS such boxes. + STATE_DTYPE = h0.element_type + STATE_BF16 = STATE_DTYPE == cutlass.BFloat16 + # Recurrence precision. Default: 3xTF32 (hi/lo splits on every product that feeds the state). + # BF16_MMA: bf16 tensor-core operands -- S0 packed as bf16 in TMEM (64 columns), bf16 K-major + # smem tiles (half the bytes), K = 16 per MMA (a quarter of the instructions); the + # bf16-operand / fp32-accumulate scheme of the prefill kernel. Gates, cumulative gates, + # exponents, reductions and the state accumulator stay fp32 on both paths. + # S0_LO: the 3xTF32 path needs the lo part of a fp32 state (a bf16 state is exact in tf32). + S0_LO = (not STATE_BF16) and (not BF16_MMA) + N1 = 2 * T_PAD if BF16_MMA else 3 * T_PAD # GEMM1 N: [K~; Q~] (+ K~lo for the 3xTF32 sweep) + KB_MMA = 16 if BF16_MMA else 8 # K per tcgen05 instruction + # raw conv inputs of the committed window and of the pending drafts (bf16 in the runtime: the + # projections are bf16, so the copy is exact; fp32 caches are accepted as well) + CONV_DTYPE = cs_q.element_type + # replay caches: k / v / beta may be bf16 (one bf16 ulp on the replayed drafts), the gate + # cache stays fp32 (it feeds the cumulative-gate exponents) + K_CACHE_DTYPE = k_cache.element_type + V_CACHE_DTYPE = v_cache.element_type + BETA_CACHE_DTYPE = beta_cache.element_type + BOX_COLS = 128 // (STATE_DTYPE.width // 8) + NUM_SLABS = K // BOX_COLS + if cutlass.const_expr(USE_ZERO_ACCEPTED): + commit_len = cutlass.Int32(0) + else: + if cutlass.const_expr(USE_PACKED_TOKEN_LAYOUT): + commit_len = num_accepted_tokens[slot] + else: + commit_len = num_accepted_tokens[i_n] + # Only NUM_SPEC drafts can be pending from the previous round. Clamp + # so a malformed count cannot drive T_loop past the t_max-sized SMEM + # buffers or the num_spec extents of the replay caches. + if commit_len > NUM_SPEC: + commit_len = cutlass.Int32(NUM_SPEC) + if cutlass.const_expr(USE_PACKED_TOKEN_LAYOUT): + # Existing body addresses every new row as bos + i_t, where i_t is + # offset by commit_len after replay. Shift bos so those accesses map + # to request-major packed rows without materialized cu_seqlens. + bos = i_n * (1 + NUM_SPEC) - commit_len + eos = bos + commit_len + 1 + NUM_SPEC + if cutlass.const_expr(USE_ZERO_ACCEPTED): + T_loop = cutlass.Int32(1 + NUM_SPEC) + t_max = 1 + NUM_SPEC + else: + T_loop = commit_len + 1 + NUM_SPEC + t_max = 2 * NUM_SPEC + 1 + vec_size = TILE_K // 32 + num_v_tiles = V // TILE_V + NUM_V_ROWS = TILE_V // (NUM_THREADS // 32) + # The restructured kernel assumes one V pass with eight rows per warp, + # smem-resident conv weights and no setmaxreg register juggling. + assert num_v_tiles == 1 and NUM_V_ROWS == 8, ( + f"kda_decode_mtp_kernel expects TILE_V == V and 8 rows per warp, got TILE_V={TILE_V}, V={V}" + ) + assert not USE_REG_Q_WEIGHTS and not USE_SETMAXREG, ( + "kda_decode_mtp_kernel reads conv weights from smem and does not use setmaxreg" + ) + q_weight_elems = KERNEL_WIDTH * K + v_weight_elems = KERNEL_WIDTH * V + k_weight_base = q_weight_elems + v_weight_base = q_weight_elems + KERNEL_WIDTH * K + conv_weight_elems = q_weight_elems + KERNEL_WIDTH * K + v_weight_elems + smem = cutlass.utils.SmemAllocator() + if cutlass.const_expr(COMPACT): + assert ( + STATE_BF16 + and FUSE_OUTPUT_NORM + and SPLIT_V == 1 + and NUM_SPEC <= 3 + and K == 128 + and not BF16_MMA + ), "compact layout: bf16 state, fused output norm, SPLIT_V == 1, NUM_SPEC <= 3, 3xTF32" + # buffers live across phases + sQ = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sK = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) + sOnG = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((1 + NUM_SPEC) * V,)), 16) + sOnW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V,)), 16) + sE = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * K,)), 16) + sEc = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + assert NUM_SPEC <= 7 and t_max <= T_PAD, "chunk form pads the token axis to 16" + tmem_holder = smem.allocate_array(cutlass.Int32, 1) + mbar1 = smem.allocate_array(cutlass.Int64, 1) + mbar2 = smem.allocate_array(cutlass.Int64, 1) + mbar_ld = smem.allocate_array(cutlass.Int64, 1) + mbar4 = smem.allocate_array(cutlass.Int64, 1) + mbar2b = smem.allocate_array(cutlass.Int64, 1) + sScratch = smem.allocate_tensor(cutlass.Float32, cute.make_layout((C_REGION_FLOATS,)), 1024) + # phase A views (dead once the precompute and the cumulative gates are done) + c_off_xk = C_OFF_STAGE + (t_max + 3) * K + c_off_xv = c_off_xk + (t_max + 3) * K + c_off_gin = c_off_xv + (t_max + 3) * V + c_off_g = c_off_gin + t_max * K + c_off_convw = c_off_g + t_max * K + c_off_dtb = c_off_convw + conv_weight_elems + c_off_alog = c_off_dtb + K + c_off_betaraw = c_off_alog + 32 + assert c_off_betaraw + 16 <= C_REGION_FLOATS, "compact phase-A staging overflows" + sXQ = cute.make_tensor( + sScratch.iterator + C_OFF_STAGE, cute.make_layout(((t_max + 3) * K,)) + ) + sXK = cute.make_tensor(sScratch.iterator + c_off_xk, cute.make_layout(((t_max + 3) * K,))) + sXV = cute.make_tensor(sScratch.iterator + c_off_xv, cute.make_layout(((t_max + 3) * V,))) + sGin = cute.make_tensor(sScratch.iterator + c_off_gin, cute.make_layout((t_max * K,))) + sG = cute.make_tensor(sScratch.iterator + c_off_g, smem_qk_layout) + sConvW = cute.make_tensor( + sScratch.iterator + c_off_convw, cute.make_layout((conv_weight_elems,)) + ) + sDtBias = cute.make_tensor(sScratch.iterator + c_off_dtb, cute.make_layout((K,))) + sAlog = cute.make_tensor(sScratch.iterator + c_off_alog, cute.make_layout((32,))) + sBetaRaw = cute.make_tensor(sScratch.iterator + c_off_betaraw, cute.make_layout((t_max,))) + # phase B operand tiles + sB1 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + C_OFF_B1, smem_layout_b1.inner, dtype=cutlass.Float32 + ), + smem_layout_b1.outer, + ) + sA3 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + C_OFF_A3, smem_layout_a3.inner, dtype=cutlass.Float32 + ), + smem_layout_a3.outer, + ) + sB3hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + C_OFF_B3HI, smem_layout_b3.inner, dtype=cutlass.Float32 + ), + smem_layout_b3.outer, + ) + sB3lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + C_OFF_B3LO, smem_layout_b3.inner, dtype=cutlass.Float32 + ), + smem_layout_b3.outer, + ) + # the S0_lo sweep (fp32 state only) never runs: alias its tile onto B1 + sB1k = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + C_OFF_B1, smem_layout_b1k.inner, dtype=cutlass.Float32 + ), + smem_layout_b1k.outer, + ) + sB1kf = cute.make_tensor(sScratch.iterator + C_OFF_B1, cute.make_layout((T_PAD * 128,))) + sB1f = cute.make_tensor(sScratch.iterator + C_OFF_B1, cute.make_layout((48 * 128,))) + sB2hif = cute.make_tensor(sScratch.iterator + C_OFF_B2HI, cute.make_layout((128 * 16,))) + sB2lof = cute.make_tensor(sScratch.iterator + C_OFF_B2LO, cute.make_layout((128 * 16,))) + sB4hif = cute.make_tensor( + sScratch.iterator + C_OFF_B4HI, cute.make_layout((T_PAD * T_PAD,)) + ) + sB4lof = cute.make_tensor( + sScratch.iterator + C_OFF_B4LO, cute.make_layout((T_PAD * T_PAD,)) + ) + sA3f = cute.make_tensor(sScratch.iterator + C_OFF_A3, cute.make_layout((64 * 128,))) + sB3hif = cute.make_tensor(sScratch.iterator + C_OFF_B3HI, cute.make_layout((16 * 128,))) + sB3lof = cute.make_tensor(sScratch.iterator + C_OFF_B3LO, cute.make_layout((16 * 128,))) + # phase C views (GEMM1 / GEMM3 operands are dead once mbar1 completed) + sA2hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + C_OFF_A2HI, smem_layout_a2.inner, dtype=cutlass.Float32 + ), + smem_layout_a2.outer, + ) + sA2lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + C_OFF_A2LO, smem_layout_a2.inner, dtype=cutlass.Float32 + ), + smem_layout_a2.outer, + ) + sA2hif = cute.make_tensor(sScratch.iterator + C_OFF_A2HI, cute.make_layout((128 * 16,))) + sA2lof = cute.make_tensor(sScratch.iterator + C_OFF_A2LO, cute.make_layout((128 * 16,))) + sLB = cute.make_tensor( + sScratch.iterator + C_OFF_LB, + cute.make_layout((4 * T_PAD, T_PAD), stride=(T_PAD, 1)), + ) + sVpT = cute.make_tensor(sScratch.iterator + C_OFF_VPT, cute.make_layout((V * TP_STRIDE,))) + sKbT = sVpT # unused: Kbar is written straight into the B2 tiles + sXp = cute.make_tensor(sScratch.iterator, cute.make_layout((16 * 32 * XP_STRIDE,))) + R_OFF_B2HI = C_OFF_B2HI + R_OFF_B2LO = C_OFF_B2LO + R_OFF_B4HI = C_OFF_B4HI + R_OFF_B4LO = C_OFF_B4LO + else: + sQ = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sK = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sG = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + # Preserve the offsets and bank mapping of the shared-memory buffers from + # the source kernel. The fused norm below reduces within one warp and does + # not need this legacy CTA-reduction scratch. + smem.allocate_tensor(cutlass.Float32, cute.make_layout((8,)), 16) + sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) + sConvW = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((conv_weight_elems,)), + 16, + ) + if cutlass.const_expr(FUSE_OUTPUT_NORM): + # Epilogue operands are staged in smem during the load phase. + sOnG = smem.allocate_tensor( + cutlass.Float32, cute.make_layout(((1 + NUM_SPEC) * V,)), 16 + ) + sOnW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V,)), 16) + # Staged per-token inputs (raw conv sequence with the committed window + # prepended in rows 0..W-2, raw gates, raw betas). + sXQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXK = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXV = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * V,)), 16) + sGin = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * K,)), 16) + sBetaRaw = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + # Per-head gate constants, staged in the load phase together with the other inputs. + sDtBias = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + sAlog = smem.allocate_tensor(cutlass.Float32, cute.make_layout((32,)), 16) + # ---- tensor-core chunk recurrence buffers ---- + assert NUM_SPEC <= 7 and t_max <= T_PAD, "chunk form pads the token axis to 16" + tmem_holder = smem.allocate_array(cutlass.Int32, 1) + mbar1 = smem.allocate_array(cutlass.Int64, 1) + mbar2 = smem.allocate_array(cutlass.Int64, 1) + mbar_ld = smem.allocate_array(cutlass.Int64, 1) + mbar4 = smem.allocate_array(cutlass.Int64, 1) + mbar2b = smem.allocate_array(cutlass.Int64, 1) + sScratch = smem.allocate_tensor(cutlass.Float32, cute.make_layout((SCRATCH_FLOATS,)), 1024) + # Swizzled fp32 MMA operand tiles as views into the scratch region (tf32 modes; the layouts + # passed in are bf16 in mode 2, whose tiles are built below). + if cutlass.const_expr(not BF16_MMA): + sB1 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B1, smem_layout_b1.inner, dtype=cutlass.Float32 + ), + smem_layout_b1.outer, + ) + sA2hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_A2HI, smem_layout_a2.inner, dtype=cutlass.Float32 + ), + smem_layout_a2.outer, + ) + sA2lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_A2LO, smem_layout_a2.inner, dtype=cutlass.Float32 + ), + smem_layout_a2.outer, + ) + sA3 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_A3, smem_layout_a3.inner, dtype=cutlass.Float32 + ), + smem_layout_a3.outer, + ) + sB3hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B3HI, smem_layout_b3.inner, dtype=cutlass.Float32 + ), + smem_layout_b3.outer, + ) + sB3lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B3LO, smem_layout_b3.inner, dtype=cutlass.Float32 + ), + smem_layout_b3.outer, + ) + sB1k = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B1K, smem_layout_b1k.inner, dtype=cutlass.Float32 + ), + smem_layout_b1k.outer, + ) + sB1kf = cute.make_tensor(sScratch.iterator + OFF_B1K, cute.make_layout((T_PAD * 128,))) + # Flat (unswizzled-pointer) views for chunked element writes with swz128/swz64_k16 offsets. + sB1f = cute.make_tensor(sScratch.iterator + OFF_B1, cute.make_layout((48 * 128,))) + sA2hif = cute.make_tensor(sScratch.iterator + OFF_A2HI, cute.make_layout((128 * 16,))) + sA2lof = cute.make_tensor(sScratch.iterator + OFF_A2LO, cute.make_layout((128 * 16,))) + sB2hif = cute.make_tensor(sScratch.iterator + OFF_B2HI, cute.make_layout((128 * 16,))) + sB2lof = cute.make_tensor(sScratch.iterator + OFF_B2LO, cute.make_layout((128 * 16,))) + sB4hif = cute.make_tensor(sScratch.iterator + OFF_B4HI, cute.make_layout((T_PAD * T_PAD,))) + sB4lof = cute.make_tensor(sScratch.iterator + OFF_B4LO, cute.make_layout((T_PAD * T_PAD,))) + sA3f = cute.make_tensor(sScratch.iterator + OFF_A3, cute.make_layout((64 * 128,))) + sB3hif = cute.make_tensor(sScratch.iterator + OFF_B3HI, cute.make_layout((16 * 128,))) + sB3lof = cute.make_tensor(sScratch.iterator + OFF_B3LO, cute.make_layout((16 * 128,))) + # bf16 operand tiles (BF16_MMA): swizzled MMA views and flat views for the closed-form writes. + sScratchH = cute.make_tensor( + cute.recast_ptr(sScratch.iterator, dtype=cutlass.BFloat16), + cute.make_layout((2 * SCRATCH_FLOATS,)), + ) + if cutlass.const_expr(BF16_MMA): + sB1b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B1, smem_layout_b1.inner, dtype=cutlass.BFloat16 + ), + smem_layout_b1.outer, + ) + sA3b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_A3, smem_layout_a3.inner, dtype=cutlass.BFloat16 + ), + smem_layout_a3.outer, + ) + sB3b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B3, smem_layout_b3.inner, dtype=cutlass.BFloat16 + ), + smem_layout_b3.outer, + ) + sA2b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_A2, smem_layout_a2.inner, dtype=cutlass.BFloat16 + ), + smem_layout_a2.outer, + ) + sB1b_f = cute.make_tensor(sScratchH.iterator + OFFH_B1, cute.make_layout((32 * 128,))) + sA3b_f = cute.make_tensor(sScratchH.iterator + OFFH_A3, cute.make_layout((64 * 128,))) + sB3b_f = cute.make_tensor(sScratchH.iterator + OFFH_B3, cute.make_layout((16 * 128,))) + sA2b_f = cute.make_tensor(sScratchH.iterator + OFFH_A2, cute.make_layout((128 * 16,))) + sB4b_f = cute.make_tensor(sScratchH.iterator + OFFH_B4, cute.make_layout((16 * 16,))) + sB2b_f = cute.make_tensor(sScratchH.iterator + OFFH_B2, cute.make_layout((128 * 16,))) + # Per-warp 32 x 32 transpose blocks (bf16 output path: per-warp state store). + sXp = cute.make_tensor(sScratch.iterator, cute.make_layout((16 * 32 * XP_STRIDE,))) + # GEMM3 result [L; B hi; B lo; -] as a plain (64 x 16) row-major matrix. + sLB = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((4 * T_PAD, T_PAD), stride=(T_PAD, 1)), 16 + ) + sEc = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + # V'^T (v, s) and Kbar^T (k, s) staged with a padded row stride (conflict-free for both + # the row-per-thread producers and the 16-lanes-per-row tile writers). + sVpT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V * TP_STRIDE,)), 16) + sKbT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K * TP_STRIDE,)), 16) + # Cumulative gates G_t (per k) reuse the raw-gate staging buffer (no longer needed after the precompute). + sE = cute.make_tensor(sGin.iterator, cute.make_layout((t_max * K,))) + R_OFF_B2HI = OFF_B2HI + R_OFF_B2LO = OFF_B2LO + R_OFF_B4HI = OFF_B4HI + R_OFF_B4LO = OFF_B4LO + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mbar1, 1) + cute.arch.mbarrier_init(mbar2, 1) + cute.arch.mbarrier_init(mbar4, 1) + cute.arch.mbarrier_init(mbar2b, 1) + cute.arch.mbarrier_init(mbar_ld, 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + # ---- TMA bulk load of S0 into the boxes; completes under the token staging / + # precompute and is consumed (mbar_ld) at the start of the chunk stage. + # Issued by the whole of one warp (a TMA load issued from a single thread never + # completes on this DSL build; the copy elects its issuing lane internally). + # Load-phase roles besides the per-warp token rows (warp w stages token w): the last K + # threads stage the committed q/k windows, the q/k conv weights and dt_bias, the V threads + # below them the committed v window, and the S0 TMA is issued by the last v-window warp + # (three window loads per lane). With 512 threads and K = V = 128 the staging lands on + # warps 8-15, of which 11-15 own no token row at production trip counts (t_max <= 15). + QK_STAGE_TID0 = NUM_THREADS - K + V_STAGE_TID0 = QK_STAGE_TID0 - V + TMA_WARP = QK_STAGE_TID0 // 32 - 1 + assert V_STAGE_TID0 >= 0 and V_STAGE_TID0 % 32 == 0 and QK_STAGE_TID0 % 32 == 0, ( + f"load-phase staging roles need whole warps for the last K + V threads, got " + f"NUM_THREADS={NUM_THREADS}, K={K}, V={V}" + ) + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ld, V * K * (STATE_DTYPE.width // 8)) + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sTl_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gS_h_all = cute.local_tile(mS_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gS_h = gS_h_all[(None, None, 0, h0_idx)] + else: + gS_h = gS_h_all[(None, None, 0, (i_hv, slot))] + tSl_h, tGl_h = cpasync.tma_partition( + tma_atom_l, + 0, + cute.make_layout(1), + cute.group_modes(sTl_h, 0, 2), + cute.group_modes(gS_h, 0, 2), + ) + if warp_idx == TMA_WARP: + cute.copy(tma_atom_l, tGl_h, tSl_h, tma_bar_ptr=mbar_ld) + # TMEM allocation overlaps the in-flight loads; the pointer is retrieved after the + # precompute barrier (first TMEM use is the S0 staging in the chunk stage). + TMEM_ALLOC = COMPACT_TMEM_COLS if cutlass.const_expr(COMPACT) else TMEM_COLS + if warp_idx == 0: + cute.arch.alloc_tmem(TMEM_ALLOC, tmem_holder) + if cutlass.const_expr(COMPACT): + # An SM does not launch another TMEM-allocating CTA while a resident CTA holds + # the allocation permit: give it up right away so the second CTA can co-reside. + cute.arch.relinquish_tmem_alloc_permit() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + t_alloc = read_globaltimer() + _, grid_n4, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * grid_n4 + i_n) * 64 + 18] = t_alloc - t_stage0 + wg = warp_idx // 4 + wg_tidx = tidx % 128 + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_decay = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_bk = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_state = cute.make_rmem_tensor( + cute.make_layout((NUM_V_ROWS * vec_size,), stride=(1,)), cutlass.Float32 + ) + r_exp_A = cutlass.Float32(0.0) + # Issue the recurrent-state loads first so their DRAM latency overlaps the + # whole precompute stage (one V pass: TILE_V == V). + # S0 arrives through the TMA bulk load issued at kernel start (mbar_ld); it is read + # from the swizzled smem boxes at the start of the chunk stage. + # Warp-uniform request bounds: the TMA store inside this region needs a provably + # uniform branch condition (legacy shifted layout only; the other layouts are unconditional). + bos = cute.arch.make_warp_uniform(bos) + eos = cute.arch.make_warp_uniform(eos) + if cutlass.const_expr(USE_REGULAR_METADATA or USE_PACKED_TOKEN_LAYOUT) or eos > bos: + if cutlass.const_expr(FUSE_PRECOMPUTE or RUNTIME_PRECOMPUTE_FLAG): + if cutlass.const_expr(RUNTIME_PRECOMPUTE_FLAG): + run_precompute = precompute_control[0] != 0 + else: + run_precompute = True + if run_precompute: + # The last K threads stage the committed q/k windows, the q/k conv weights and + # dt_bias (see the load-phase roles above). + ct = tidx - QK_STAGE_TID0 + if ct >= 0: + for w in cutlass.range(KERNEL_WIDTH - 1, unroll_full=True): + sXQ[w * K + ct] = cutlass.Float32(cs_q[slot, hk_off + ct, w]) + sXK[w * K + ct] = cutlass.Float32(cs_k[slot, hk_off + ct, w]) + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + if ct >= 0: + sConvW[w * K + ct] = cutlass.Float32(w_q[hk_off + ct, w]) + sConvW[k_weight_base + w * K + ct] = cutlass.Float32(w_k[hk_off + ct, w]) + for ld in cutlass.range(max(1, V * KERNEL_WIDTH // NUM_THREADS), unroll_full=True): + flat = ld * NUM_THREADS + tidx + if flat < V * KERNEL_WIDTH: + sConvW[v_weight_base + flat] = cutlass.Float32( + w_v[hv_off + flat % V, flat // V] + ) + # Load phase: one round trip for every per-token input. Warp w + # stages tokens w, w + num_warps, ...; replay tokens go straight + # into sQ/sK/sG/sBeta/sVall, new tokens stage their raw inputs. + if ct >= 0: + sDtBias[ct] = cutlass.Float32(dt_bias[i_h * K + ct]) + if tidx == 0: + sAlog[0] = cutlass.Float32(A_log[i_h]) + cv = tidx - V_STAGE_TID0 + if cv >= 0 and cv < V: + for _w in cutlass.range(KERNEL_WIDTH - 1, unroll_full=True): + sXV[_w * V + cv] = cutlass.Float32(cs_v[slot, hv_off + cv, _w]) + for _rep in cutlass.range( + (t_max + NUM_THREADS // 32 - 1) // (NUM_THREADS // 32), unroll_full=True + ): + _t = _rep * (NUM_THREADS // 32) + warp_idx + if _t < T_loop: + if _t < commit_len: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + # replayed drafts produce no output: their q rows are unused + sQ[_t, k_idx] = cutlass.Float32(0.0) + sK[_t, k_idx] = cutlass.Float32(k_cache[slot, _t, hk_off + k_idx]) + sG[_t, k_idx] = cutlass.Float32(g_cache[slot, _t, hk_off + k_idx]) + sXQ[(_t + 3) * K + k_idx] = cutlass.Float32( + cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + _t] + ) + sXK[(_t + 3) * K + k_idx] = cutlass.Float32( + cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + _t] + ) + sXV[(_t + 3) * V + k_idx] = cutlass.Float32( + cs_v[slot, hv_off + k_idx, KERNEL_WIDTH - 1 + _t] + ) + sVall[_t * V + k_idx] = cutlass.Float32( + v_cache[slot, _t, hv_off + k_idx] + ) + if in_warp_tid == 0: + sBeta[_t] = cutlass.Float32(beta_cache[slot, _t, i_hv]) + else: + _tok = bos + _t + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + if cutlass.const_expr(USE_FLAT_LAYOUT): + sXQ[(_t + 3) * K + k_idx] = cutlass.Float32( + x_q[0, _tok, hk_off + k_idx] + ) + sXK[(_t + 3) * K + k_idx] = cutlass.Float32( + x_k[0, _tok, hk_off + k_idx] + ) + sXV[(_t + 3) * V + k_idx] = cutlass.Float32( + x_v[0, _tok, hv_off + k_idx] + ) + else: + sXQ[(_t + 3) * K + k_idx] = cutlass.Float32( + x_q[0, _tok, i_h, k_idx] + ) + sXK[(_t + 3) * K + k_idx] = cutlass.Float32( + x_k[0, _tok, i_h, k_idx] + ) + sXV[(_t + 3) * V + k_idx] = cutlass.Float32( + x_v[0, _tok, i_hv, k_idx] + ) + sGin[_t * K + k_idx] = cutlass.Float32(g[0, _tok, i_hv, k_idx]) + if in_warp_tid == 0: + sBetaRaw[_t] = cutlass.Float32(beta[0, _tok, i_hv]) + # Padding rows [T_loop, t_max) of the gates, values and betas are zeroed (under the + # load latency) so the chunk stage's r_Y / cumulative-gate / V' loops are branch-free. + if tidx < K: + for _z in cutlass.range(T_loop, t_max, 1): + sG[_z, tidx] = cutlass.Float32(0.0) + sVall[_z * V + tidx] = cutlass.Float32(0.0) + if tidx == 0: + sBeta[_z] = cutlass.Float32(0.0) + if cutlass.const_expr(FUSE_OUTPUT_NORM): + # Threads 224+ stage the gated-RMSNorm operands for the epilogue. + if tidx >= 224: + _pf_tid = tidx - 224 + for _e in cutlass.range( + ((1 + NUM_SPEC) * V + (NUM_THREADS - 224) - 1) // (NUM_THREADS - 224), + unroll_full=True, + ): + _flat = _e * (NUM_THREADS - 224) + _pf_tid + if _flat < (1 + NUM_SPEC) * V: + _tok = _flat // V + _vv = _flat % V + sOnG[_flat] = cutlass.Float32( + onorm_g[0, i_n * (1 + NUM_SPEC) + _tok, i_hv, _vv] + ) + for _e in cutlass.range( + (V + (NUM_THREADS - 224) - 1) // (NUM_THREADS - 224), unroll_full=True + ): + _flat = _e * (NUM_THREADS - 224) + _pf_tid + if _flat < V: + sOnW[_flat] = cutlass.Float32(onorm_weight[_flat]) + cute.arch.barrier() + # ---- per-token-parallel precompute ---- + num_warps = NUM_THREADS // 32 + half = num_warps // 2 + role = warp_idx // half + r_exp_A = cute.math.exp(sAlog[0], fastmath=True) + for _rep in cutlass.range((1 + NUM_SPEC + half - 1) // half, unroll_full=True): + j = _rep * half + warp_idx % half + if j < 1 + NUM_SPEC: + t = commit_len + j + if role == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_conv_0 = cutlass.Float32(0.0) + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + r_conv_0 += sXQ[(t + w) * K + k_idx] * sConvW[w * K + k_idx] + e0 = cute.math.exp(-r_conv_0, fastmath=True) + r_q[i] = r_conv_0 * cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) + sum_q = 0.0 + for i in cutlass.range(vec_size, unroll_full=True): + sum_q += r_q[i] * r_q[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_q_scaled = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + for i in cutlass.range(vec_size, unroll_full=True): + r_q[i] = r_q[i] * rnorm_q_scaled + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + sQ[t, k_idx] = r_q[i] + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_conv = sXK[t * K + k_idx] * sConvW[k_weight_base + 0 * K + k_idx] + for w in cutlass.range(1, KERNEL_WIDTH, unroll_full=True): + r_conv += ( + sXK[(t + w) * K + k_idx] + * sConvW[k_weight_base + w * K + k_idx] + ) + r_conv = r_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) + ) + r_k[i] = r_conv + sum_k = 0.0 + for i in cutlass.range(vec_size, unroll_full=True): + sum_k += r_k[i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in cutlass.range(vec_size, unroll_full=True): + r_k[i] = r_k[i] * rnorm_k + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + sK[t, k_idx] = r_k[i] + if j == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + for w in cutlass.range(KERNEL_WIDTH - 1, unroll_full=True): + cs_q[slot, hk_off + k_idx, w] = sXQ[ + (t + 1 + w) * K + k_idx + ].to(CONV_DTYPE) + cs_k[slot, hk_off + k_idx, w] = sXK[ + (t + 1 + w) * K + k_idx + ].to(CONV_DTYPE) + else: + cache_pos = j - 1 + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + k_cache[slot, cache_pos, hk_off + k_idx] = r_k[i].to( + K_CACHE_DTYPE + ) + cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = sXQ[ + (t + 3) * K + k_idx + ].to(CONV_DTYPE) + cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = sXK[ + (t + 3) * K + k_idx + ].to(CONV_DTYPE) + else: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_g_raw = sGin[t * K + k_idx] + r_g_raw = r_g_raw + sDtBias[k_idx] + exp_A_x = r_exp_A * r_g_raw + sigmoid_val = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) + ) + r_gk = lower_bound * sigmoid_val + sG[t, k_idx] = r_gk + r_decay[i] = r_gk + r_beta_t = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-sBetaRaw[t], fastmath=True) + ) + if in_warp_tid == 0: + sBeta[t] = r_beta_t + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid + _v_conv = 0.0 + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + _v_conv += ( + sXV[(t + w) * V + v_idx] + * sConvW[v_weight_base + w * V + v_idx] + ) + _v_conv = _v_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) + ) + sVall[t * V + v_idx] = _v_conv + r_bk[i] = _v_conv + if j == 0: + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid + for w in cutlass.range(KERNEL_WIDTH - 1, unroll_full=True): + cs_v[slot, hv_off + v_idx, w] = sXV[ + (t + 1 + w) * V + v_idx + ].to(CONV_DTYPE) + else: + cache_pos = j - 1 + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + g_cache[slot, cache_pos, hk_off + k_idx] = r_decay[i] + v_cache[slot, cache_pos, hv_off + k_idx] = r_bk[i].to( + V_CACHE_DTYPE + ) + cs_v[slot, hv_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = sXV[ + (t + 3) * V + k_idx + ].to(CONV_DTYPE) + if in_warp_tid == 0: + beta_cache[slot, cache_pos, i_hv] = r_beta_t.to( + BETA_CACHE_DTYPE + ) + if cutlass.const_expr(PROFILE_STAGES): + cute.arch.barrier() + t_stage1 = read_globaltimer() + else: + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + cute.arch.barrier() + t_stage1 = read_globaltimer() + else: + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + cute.arch.barrier() + t_stage1 = read_globaltimer() + cute.arch.barrier() + # ================================================================= + # Chunk (WY-form) recurrence on tensor cores. + # ================================================================= + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + pool = cutlass.utils.TmemBufferPool(tmem_ptr, TMEM_ALLOC) + tD4_frag = tiled_mma4.make_fragment_C(tiled_mma4.partition_shape_C((V, T_PAD))) + tD4 = pool.allocate_tensor(tD4_frag.layout, cutlass.Float32) # outputs (columns 0..15) + tD2h_frag = tiled_mma2h.make_fragment_C(tiled_mma2h.partition_shape_C((V, K // 2))) + tD2h0 = pool.allocate_tensor( + tD2h_frag.layout, cutlass.Float32 + ) # GEMM2 half 0: k columns 0..63 + tD2h1 = pool.allocate_tensor( + tD2h_frag.layout, cutlass.Float32 + ) # GEMM2 half 1: k columns 64..127 + tD2q_frag = tiled_mma2q.make_fragment_C(tiled_mma2q.partition_shape_C((V, T_PAD + K // 2))) + tD2q = cute.make_tensor(tD4.iterator, tD2q_frag.layout) # [D4 | ΔS half 0] accumulator view + tA_frag = tiled_mma_f32a.make_fragment_A(tiled_mma_f32a.partition_shape_A((V, K))) + if cutlass.const_expr(BF16_MMA): + # S0 as packed bf16 (even k in the low half of each 32-bit column): 64 columns, staged + # through an fp32 word view and consumed by the MMA through the bf16 A-fragment view. + tAw = pool.allocate_tensor( + cute.make_layout(((V, 8), 1, K // 16), stride=((65536, 1), 0, 8)), cutlass.Float32 + ) + tAb = cute.make_tensor( + cute.recast_ptr(tAw.iterator, dtype=cutlass.BFloat16), + tiled_mma1.make_fragment_A(tiled_mma1.partition_shape_A((V, K))).layout, + ) + tAhi = tAw + tAlo = tAw + elif cutlass.const_expr(COMPACT): + # S0 hi shares the GEMM2 accumulator columns (seeded after GEMM1 completed); + # there is no S0 lo for a bf16 state. + tAhi = cute.make_tensor(tD2h0.iterator, tA_frag.layout) + tAlo = tAhi + else: + tAhi = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tAlo = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tD1_frag = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C((V, N1))) + tD1 = pool.allocate_tensor(tD1_frag.layout, cutlass.Float32) + tD3_frag = tiled_mma3.make_fragment_C(tiled_mma3.partition_shape_C((4 * T_PAD, T_PAD))) + tD3 = pool.allocate_tensor(tD3_frag.layout, cutlass.Float32) + tD2 = cute.make_tensor( + tD2h0.iterator, tA_frag.layout + ) # A-view layout for st/ld (all 128 columns) + + st32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + ld16 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + ld32 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + tiled_stA = tcgen05.make_tmem_copy(st32, tD2) + thr_stA = tiled_stA.get_slice(wg_tidx) + tD2_st = thr_stA.partition_D(tD2) + rA_shape = thr_stA.partition_S(tD2).shape + if cutlass.const_expr(BF16_MMA): + st16w = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_stAw = tcgen05.make_tmem_copy(st16w, tAw).get_slice(wg_tidx) + tAw_st = thr_stAw.partition_D(tAw) + r_pack = cute.make_rmem_tensor(thr_stAw.partition_S(tAw).shape[:-1], cutlass.Float32) + r_pack_flat = cute.make_tensor(r_pack.iterator, cute.make_layout((16,))) + tAhi_st = tD2_st + tAlo_st = tD2_st + else: + tAhi_st = thr_stA.partition_D(tAhi) + tAlo_st = thr_stA.partition_D(tAlo) + tiled_ldD2 = tcgen05.make_tmem_copy(ld32, tD2) + thr_ldD2 = tiled_ldD2.get_slice(wg_tidx) + tD2_ld = thr_ldD2.partition_S(tD2) + rS = cute.make_rmem_tensor(thr_ldD2.partition_D(tD2).shape[:-1], cutlass.Float32) + rS_flat = cute.make_tensor(rS.iterator, cute.make_layout((32,))) + # TMA store partitions for the NUM_SLABS boxes of this CTA's committed-state tile. + # Created here, outside the dynamic request-bounds region (the exec-TMA lowering is + # only legal at this level); the copies are issued inside. + sBoxf = cute.make_tensor(sScratch.iterator, cute.make_layout((V * K,))) + sBoxh = cute.make_tensor( + cute.recast_ptr(sScratch.iterator, dtype=cutlass.BFloat16), cute.make_layout((V * K,)) + ) + # committed-state boxes (phase D): the compact layout stores them away from the + # A2 / B2-half-1 tiles still read by GEMM2 half 1 + ST_BOX_OFF = C_OFF_STBOX if cutlass.const_expr(COMPACT) else 0 + sBoxhS = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + ST_BOX_OFF, dtype=cutlass.BFloat16), + cute.make_layout((V * K,)), + ) + tS_list = [] + tG_list = [] + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sT_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + ST_BOX_OFF + sl_h * ((V // SPLIT_V) * 32), + box_layout_t.inner, + dtype=STATE_DTYPE, + ), + box_layout_t.outer, + ) + gT_h_all = cute.local_tile(mT_tma, (V // SPLIT_V, BOX_COLS, 1), (vh, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gT_h = gT_h_all[(None, None, 0, h0_idx)] + else: + gT_h = gT_h_all[(None, None, 0, (i_hv, slot))] + tS_h, tG_h = cpasync.tma_partition( + tma_atom_t, + 0, + cute.make_layout(1), + cute.group_modes(sT_h, 0, 2), + cute.group_modes(gT_h, 0, 2), + ) + tS_list.append(tS_h) + tG_list.append(tG_h) + r_hi = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_lo = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_hi_flat = cute.make_tensor(r_hi.iterator, cute.make_layout((32,))) + r_lo_flat = cute.make_tensor(r_lo.iterator, cute.make_layout((32,))) + + # ---- S0 row v = wg_tidx, columns [32 wg, +32), from the TMA-loaded swizzled box + # (16 B chunks: conflict-free LDS.128 under the 128B swizzle) + cute.arch.mbarrier_wait(mbar_ld, 0) + # in-place state update: the peer CTA must not store its half before our S0 load + # completed -> arrive here, wait right before the state store. + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_arrive_relaxed() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq28 = read_globaltimer() + _, _gq28, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq28 + i_n) * 64 + 28] = _tq28 - t_stage0 + vx = 4 * (wg_tidx % 8) + if cutlass.const_expr(STATE_BF16): + # bf16 box: 64 columns per 128 B row, so this warpgroup's 32 columns are the 16 B + # chunks [4 (wg % 2), +4) of box wg // 2. A bf16 value is exact in tf32: no lo part, + # and the S0_lo sweep of GEMM1 is skipped. + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (wg % 2) + c) ^ (wg_tidx % 8) + off_l = (wg // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + for i in cutlass.range(8, unroll_full=True): + r_state[8 * c + i] = cutlass.Float32(sBoxh[off_l + i]) + if cutlass.const_expr(not BF16_MMA): + for i in cutlass.range(32, unroll_full=True): + r_hi_flat[i] = r_state[i] + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, j)]) + else: + for c in cutlass.range(8, unroll_full=True): + off_l = wg * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range(4, unroll_full=True): + r_state[4 * c + i] = sBoxf[off_l + i] + if cutlass.const_expr(not BF16_MMA): + for i in cutlass.range(32, unroll_full=True): + s_hi = round_hi(r_state[i], BF16_MMA) + r_hi_flat[i] = s_hi + if cutlass.const_expr(S0_LO): + r_lo_flat[i] = r_state[i] - s_hi + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, j)]) + if cutlass.const_expr(S0_LO): + cute.copy(st32, r_lo, tAlo_st[(None, None, None, j)]) + if cutlass.const_expr(BF16_MMA): + # this warpgroup's 32 state columns -> 16 packed bf16x2 words -> TMEM word slab wg + for i in cutlass.range(16, unroll_full=True): + r_pack_flat[i] = cvt_bf16x2(r_state[2 * i + 1], r_state[2 * i]) + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st16w, r_pack, tAw_st[(None, None, None, j)]) + cute.arch.fence_view_async_tmem_store() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq29 = read_globaltimer() + _, _gq29, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq29 + i_n) * 64 + 29] = _tq29 - t_stage0 + + # ---- cumulative gates per k column (thread kq, quarter q owns tokens t = q mod 4): + # G_t -> sE, the T x T coupling factors, and the GEMM1 / GEMM3 operand rows written + # with closed-form 128B-swizzle offsets (lanes = consecutive k: conflict-free). + kq = tidx % K + quarter = tidx // K + t_mid = T_loop // 2 + r_G = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + g_run = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t < t_max: + g_run = g_run + sG[t, kq] + r_G[t] = g_run + g_mid = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t == t_mid: + g_mid = r_G[t] + # Factors of the T x T coupling: e^{G_t - G_mid} (t side) and e^{G_mid - G_s} (s side), + # clamped at e^87 so a chunk whose decay exceeds the fp32 range yields ~0 instead of + # inf/NaN (the affected products are < e^{-|g|} and negligible either way). + exp_clamp = cutlass.Float32(87.0) + # the scratch region (S0 boxes) is reused for the MMA operand tiles below + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq30 = read_globaltimer() + _, _gq30, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq30 + i_n) * 64 + 30] = _tq30 - t_stage0 + ka = kq // 32 + kc = kq % 32 + b1_base = ka * (3 * T_PAD * 32) + a3_base = ka * (4 * T_PAD * 32) + b3_base = ka * (T_PAD * 32) + b1k_base = ka * (T_PAD * 32) + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + kt_hi = cutlass.Float32(0.0) + kt_lo = cutlass.Float32(0.0) + qt_r = cutlass.Float32(0.0) + kf = cutlass.Float32(0.0) + qf_hi = cutlass.Float32(0.0) + qf_lo = cutlass.Float32(0.0) + kb_hi = cutlass.Float32(0.0) + kb_lo = cutlass.Float32(0.0) + if t < T_loop: + d_t = r_G[t] - g_mid + e_t = cute.math.exp(-cute.arch.fmax(-d_t, -exp_clamp), fastmath=True) + e_bt = cute.math.exp(-cute.arch.fmax(d_t, -exp_clamp), fastmath=True) + eg_t = cute.math.exp(r_G[t], fastmath=True) + sE[t * K + kq] = r_G[t] + if t == commit_len: + sEc[kq] = eg_t + k_t = sK[t, kq] + q_t = sQ[t, kq] + kt = k_t * eg_t + kt_hi = round_hi(kt, BF16_MMA) + qt_r = round_hi(q_t * eg_t, BF16_MMA) + kf = round_hi(k_t * e_t, BF16_MMA) + qf = q_t * e_t + qf_hi = round_hi(qf, BF16_MMA) + kb = k_t * e_bt + kb_hi = round_hi(kb, BF16_MMA) + if cutlass.const_expr(not BF16_MMA): + kt_lo = kt - kt_hi + qf_lo = qf - qf_hi + kb_lo = kb - kb_hi + if cutlass.const_expr(BF16_MMA): + sB1b_f[sw128_bf16(t, kq, 2 * T_PAD)] = cutlass.BFloat16(kt_hi) + sB1b_f[sw128_bf16(T_PAD + t, kq, 2 * T_PAD)] = cutlass.BFloat16(qt_r) + sA3b_f[sw128_bf16(t, kq, 4 * T_PAD)] = cutlass.BFloat16(kf) + sA3b_f[sw128_bf16(T_PAD + t, kq, 4 * T_PAD)] = cutlass.BFloat16(qf_hi) + # rows 32..47 feed the (unused, must-be-zero) B_lo block of D3 + sA3b_f[sw128_bf16(2 * T_PAD + t, kq, 4 * T_PAD)] = cutlass.BFloat16(0.0) + sB3b_f[sw128_bf16(t, kq, T_PAD)] = cutlass.BFloat16(kb_hi) + else: + # rows t, 16+t, 32+t share t % 8 -> the same 16-byte-chunk XOR + xc = kc ^ (4 * (t % 8)) + sB1f[b1_base + 32 * t + xc] = kt_hi + if cutlass.const_expr(S0_LO): + sB1kf[b1k_base + 32 * t + xc] = kt_hi + sB1f[b1_base + 32 * (T_PAD + t) + xc] = qt_r + sB1f[b1_base + 32 * (2 * T_PAD + t) + xc] = kt_lo + sA3f[a3_base + 32 * t + xc] = kf + sA3f[a3_base + 32 * (T_PAD + t) + xc] = qf_hi + sA3f[a3_base + 32 * (2 * T_PAD + t) + xc] = qf_lo + sB3hif[b3_base + 32 * t + xc] = kb_hi + sB3lof[b3_base + 32 * t + xc] = kb_lo + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + tp_a = read_globaltimer() + + # ---- operand fragments of GEMM1 (D1 = S0_hi B1^T + S0_lo K~hi^T) and GEMM3 + # (D3 = A3 Kb_hi^T + A3 Kb_lo^T); the MMAs are issued by warp 0 after the Kbar staging + if cutlass.const_expr(BF16_MMA): + tCrB1 = tiled_mma1.make_fragment_B(sB1b) + tCrA3 = tiled_mma3.make_fragment_A(sA3b) + tCrB3hi = tiled_mma3.make_fragment_B(sB3b) + tA1 = tAb + else: + sB1_tf = cute.make_tensor( + cute.recast_ptr(sB1.iterator, smem_layout_b1.inner, dtype=cutlass.TFloat32), + sB1.layout, + ) + tCrB1 = tiled_mma1.make_fragment_B(sB1_tf) + sB1k_tf = cute.make_tensor( + cute.recast_ptr(sB1k.iterator, smem_layout_b1k.inner, dtype=cutlass.TFloat32), + sB1k.layout, + ) + tCrB1k = tiled_mma1k.make_fragment_B(sB1k_tf) + tD1k = cute.make_tensor( + tD1.iterator, + tiled_mma1k.make_fragment_C(tiled_mma1k.partition_shape_C((V, T_PAD))).layout, + ) + tAhi_tf = cute.make_tensor( + cute.recast_ptr(tAhi.iterator, dtype=cutlass.TFloat32), tAhi.layout + ) + tAlo_tf = cute.make_tensor( + cute.recast_ptr(tAlo.iterator, dtype=cutlass.TFloat32), tAlo.layout + ) + sA3_tf = cute.make_tensor( + cute.recast_ptr(sA3.iterator, smem_layout_a3.inner, dtype=cutlass.TFloat32), + sA3.layout, + ) + sB3hi_tf = cute.make_tensor( + cute.recast_ptr(sB3hi.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3hi.layout, + ) + sB3lo_tf = cute.make_tensor( + cute.recast_ptr(sB3lo.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3lo.layout, + ) + tCrA3 = tiled_mma3.make_fragment_A(sA3_tf) + tCrB3hi = tiled_mma3.make_fragment_B(sB3hi_tf) + tCrB3lo = tiled_mma3.make_fragment_B(sB3lo_tf) + tA1 = tAhi_tf + # ---- GEMM2 accumulator seed (D2 = S0 e^{G_c}) and the Kbar tiles first; warp 0 then + # issues GEMM1/GEMM3 while the other warps build the B2 tiles. + # Kbar^T staging: thread (kq, quarter) -> sKbT[kq, s] for its 4 tokens (lanes = k: conflict-free) + if cutlass.const_expr(not COMPACT): + for i in cutlass.range(4, unroll_full=True): + s_idx = quarter * 4 + i + b_val = cutlass.Float32(0.0) + if s_idx <= commit_len: + b_val = sK[s_idx, kq] * cute.math.exp( + sE[commit_len * K + kq] - sE[s_idx * K + kq], fastmath=True + ) + sKbT[kq * TP_STRIDE + s_idx] = b_val + if warp_idx == 0: + cute.arch.barrier_arrive(barrier_id=BAR_KBAR, number_of_threads=NUM_THREADS) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // KB_MMA): + cute.gemm( + tiled_mma3, + tD3, + tCrA3[(None, None, kb_i, 0)], + tCrB3hi[(None, None, kb_i, 0)], + tD3, + ) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(not BF16_MMA): + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, + tD3, + tCrA3[(None, None, kb_i, 0)], + tCrB3lo[(None, None, kb_i, 0)], + tD3, + ) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // KB_MMA): + cute.gemm( + tiled_mma1, tD1, tA1[(None, None, kb_i)], tCrB1[(None, None, kb_i, 0)], tD1 + ) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(S0_LO): + tiled_mma1k.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1k, + tD1k, + tAlo_tf[(None, None, kb_i)], + tCrB1k[(None, None, kb_i, 0)], + tD1k, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar1) + else: + cute.arch.barrier(barrier_id=BAR_KBAR, number_of_threads=NUM_THREADS) + if cutlass.const_expr(PROFILE_STAGES): + tp_j = read_globaltimer() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 32: + _tq31 = read_globaltimer() + _, _gq31, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq31 + i_n) * 64 + 31] = _tq31 - t_stage0 + # ---- GEMM2 accumulator seed (overlaps the GEMM1/GEMM3 execution) + # D4 (output-combination) columns start from zero: GEMM2 half 0 accumulates into [D4 | ΔS] + st16 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + tiled_stD4 = tcgen05.make_tmem_copy(st16, tD4) + tD4_st = tiled_stD4.get_slice(wg_tidx).partition_D(tD4) + r_zero = cute.make_rmem_tensor( + tiled_stD4.get_slice(wg_tidx).partition_S(tD4).shape, cutlass.Float32 + ) + for i in cutlass.range(cute.size(r_zero), unroll_full=True): + r_zero[i] = cutlass.Float32(0.0) + if wg == 0: + cute.copy(st16, r_zero, tD4_st) + # only this CTA's state rows (TMEM lanes [vh * V/2, +V/2)) are seeded and stored + if cutlass.const_expr(not COMPACT): + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + for i in cutlass.range(32, unroll_full=True): + kcol = wg * 32 + i + r_hi_flat[i] = r_state[i] * sEc[kcol] + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st32, r_hi, tD2_st[(None, None, None, j)]) + cute.arch.fence_view_async_tmem_store() + s_col = in_warp_tid % T_PAD + if cutlass.const_expr(COMPACT): + # ---- B2 tiles straight from Kbar (same values as the sKbT path), thread + # (kq, quarter) -> rows kq, columns 4 quarter .. 4 quarter + 3 + for i in cutlass.range(4, unroll_full=True): + s_idx = quarter * 4 + i + b_val = cutlass.Float32(0.0) + if s_idx <= commit_len: + b_val = sK[s_idx, kq] * cute.math.exp( + sE[commit_len * K + kq] - sE[s_idx * K + kq], fastmath=True + ) + b_hi = round_tf32(b_val) + o_b2 = swz64_k16(kq, s_idx) + sB2hif[o_b2] = b_hi + sB2lof[o_b2] = b_val - b_hi + # ---- B2 tiles (k rows x s cols, K-major SW64) from sKbT while the MMAs run: + # lanes 0..15 -> row 2w, lanes 16..31 -> row 2w+1, four 32-row rounds + if cutlass.const_expr(not COMPACT): + if warp_idx > 0: + for r in cutlass.range(4, unroll_full=True): + k_row = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + b_val2 = sKbT[k_row * TP_STRIDE + s_col] + b_hi = round_hi(b_val2, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sB2b_f[sw32_k16_bf16(k_row, s_col)] = cutlass.BFloat16(b_hi) + else: + o_b2 = swz64_k16(k_row, s_col) + sB2hif[o_b2] = b_hi + sB2lof[o_b2] = b_val2 - b_hi + # warp 0 is issuing the MMAs: warps 1..4 cover its rows 32 r + {0, 1} + if warp_idx <= 4: + k_row0 = 32 * (warp_idx - 1) + in_warp_tid // T_PAD + b_val0 = sKbT[k_row0 * TP_STRIDE + s_col] + b_hi0 = round_hi(b_val0, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sB2b_f[sw32_k16_bf16(k_row0, s_col)] = cutlass.BFloat16(b_hi0) + else: + o_b20 = swz64_k16(k_row0, s_col) + sB2hif[o_b20] = b_hi0 + sB2lof[o_b20] = b_val0 - b_hi0 + + if cutlass.const_expr(PROFILE_STAGES): + tp_k = read_globaltimer() + # ---- wait for GEMM1/GEMM3 + cute.arch.mbarrier_wait(mbar1, 0) + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(COMPACT): + # GEMM1 has consumed S0 hi: seed the (aliased) GEMM2 accumulator D2 = S0 e^{G_c} + for i in cutlass.range(32, unroll_full=True): + kcol = wg * 32 + i + r_hi_flat[i] = r_state[i] * sEc[kcol] + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st32, r_hi, tD2_st[(None, None, None, j)]) + cute.arch.fence_view_async_tmem_store() + if cutlass.const_expr(PROFILE_STAGES): + tp_b = read_globaltimer() + tiled_ld1 = tcgen05.make_tmem_copy(ld16, tD1) + thr_ld1 = tiled_ld1.get_slice(wg_tidx) + tD1_ld = thr_ld1.partition_S(tD1) + rD1 = cute.make_rmem_tensor(thr_ld1.partition_D(tD1).shape, cutlass.Float32) + cute.copy(ld16, tD1_ld, rD1) + cute.arch.fence_view_async_tmem_load() + rD1_flat = cute.make_tensor(rD1.iterator, cute.make_layout((N1,))) + v_row = wg_tidx + # warpgroup 0: D3 (64 x 16) -> smem [L; B_hi; B_lo; -] + ld3 = cute.make_copy_atom( + tcgen05.Ld16x128bOp(tcgen05.Repetition(T_PAD * 32 // 128)), cutlass.Float32 + ) + tiled_ld3 = tcgen05.make_tmem_copy(ld3, tD3) + thr_ld3 = tiled_ld3.get_slice(wg_tidx) + tD3_ld = thr_ld3.partition_S(tD3) + sLB3 = cute.make_tensor( + sLB.iterator, + cute.make_layout(((4 * T_PAD, T_PAD), 1, 1), stride=((T_PAD, 1), 0, 0)), + ) + tD3_sm = thr_ld3.partition_D(sLB3) + rD3 = cute.make_rmem_tensor(tD3_sm.shape, cutlass.Float32) + if wg == 0: + cute.copy(ld3, tD3_ld, rD3) + cute.arch.fence_view_async_tmem_load() + cute.autovec_copy(rD3, tD3_sm) + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + tp_c = read_globaltimer() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _, _gq32, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq32 + i_n) * 64 + 32] = tp_c - t_stage0 + # ---- warpgroup 0: V' by forward substitution (row v per thread, four accumulators); + # warpgroups 1-2: output-combination operand Bcoef[t, s] = B_hi + B_lo (rows t < T_loop, + # s <= t; zero elsewhere), one element per thread, K-major SW64 16 x 16 tiles (hi / lo) + # that head the GEMM2 half-0 B tile + if wg == 0: + r_Y = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for s in cutlass.range(T_PAD, unroll_full=True): + y = cutlass.Float32(0.0) + if s < t_max: + if cutlass.const_expr(BF16_MMA): + u0 = rD1_flat[s] + else: + u0 = rD1_flat[s] + rD1_flat[2 * T_PAD + s] + y = sBeta[s] * (sVall[s * V + v_row] - u0) + r_Y[s] = y + r_V = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + vp = cutlass.Float32(0.0) + if t < t_max: + acc0 = cutlass.Float32(0.0) + acc1 = cutlass.Float32(0.0) + acc2 = cutlass.Float32(0.0) + acc3 = cutlass.Float32(0.0) + for s in range(t): + if s % 4 == 0: + acc0 = acc0 + sLB[t, s] * r_V[s] + elif s % 4 == 1: + acc1 = acc1 + sLB[t, s] * r_V[s] + elif s % 4 == 2: + acc2 = acc2 + sLB[t, s] * r_V[s] + else: + acc3 = acc3 + sLB[t, s] * r_V[s] + vp = r_Y[t] - sBeta[t] * ((acc0 + acc1) + (acc2 + acc3)) + r_V[t] = vp + for t in cutlass.range(T_PAD, unroll_full=True): + sVpT[v_row * TP_STRIDE + t] = r_V[t] + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq36 = read_globaltimer() + _, _gq36, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq36 + i_n) * 64 + 36] = _tq36 - t_stage0 + else: + if tidx < 128 + T_PAD * T_PAD: + t4 = (tidx - 128) // T_PAD + s4 = (tidx - 128) % T_PAD + bc = cutlass.Float32(0.0) + if (t4 < T_loop) & (s4 <= t4): + bc = sLB[T_PAD + t4, s4] + sLB[2 * T_PAD + t4, s4] + bc_hi = round_hi(bc, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sB4b_f[sw32_k16_bf16(t4, s4)] = cutlass.BFloat16(bc_hi) + else: + o_b4 = swz64_k16(t4, s4) + sB4hif[o_b4] = bc_hi + sB4lof[o_b4] = bc - bc_hi + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq37 = read_globaltimer() + _, _gq37, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq37 + i_n) * 64 + 37] = _tq37 - t_stage0 + # ---- A2 tiles (v rows x s cols) from sVpT, same 16-lanes-per-row mapping; s > c -> 0 + for r in cutlass.range(4, unroll_full=True): + v_row2 = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + a_val = sVpT[v_row2 * TP_STRIDE + s_col] + a_hi = round_hi(a_val, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sA2b_f[sw32_k16_bf16(v_row2, s_col)] = cutlass.BFloat16(a_hi) + else: + o_a2 = swz64_k16(v_row2, s_col) + sA2hif[o_a2] = a_hi + sA2lof[o_a2] = a_val - a_hi + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + tp_d = read_globaltimer() + if cutlass.const_expr(BF16_MMA): + tCrA2hi = tiled_mma2h.make_fragment_A(sA2b) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2b) + sB2b_h1 = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B2 + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.BFloat16, + ), + smem_layout_b2h.outer, + ) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B(sB2b_h1) + sB2qb = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B4, smem_layout_b2q.inner, dtype=cutlass.BFloat16 + ), + smem_layout_b2q.outer, + ) + tCrB2q_hi = tiled_mma2q.make_fragment_B(sB2qb) + else: + sA2hi_tf = cute.make_tensor( + cute.recast_ptr(sA2hi.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2hi.layout, + ) + sA2lo_tf = cute.make_tensor( + cute.recast_ptr(sA2lo.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2lo.layout, + ) + tCrA2hi = tiled_mma2h.make_fragment_A(sA2hi_tf) + tCrA2lo = tiled_mma2h.make_fragment_A(sA2lo_tf) + # B2 halves: k rows 0..63 / 64..127 of the SW64 K-major tiles (a 4 KB shift keeps the swizzle) + sB2hi_h1 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + R_OFF_B2HI + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.Float32, + ), + smem_layout_b2h.outer, + ) + sB2lo_h1 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + R_OFF_B2LO + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.Float32, + ), + smem_layout_b2h.outer, + ) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2hi_h1.iterator, smem_layout_b2h.inner, dtype=cutlass.TFloat32 + ), + sB2hi_h1.layout, + ) + ) + tCrB2lo_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2lo_h1.iterator, smem_layout_b2h.inner, dtype=cutlass.TFloat32 + ), + sB2lo_h1.layout, + ) + ) + sB2q_hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + R_OFF_B4HI, smem_layout_b2q.inner, dtype=cutlass.Float32 + ), + smem_layout_b2q.outer, + ) + sB2q_lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + R_OFF_B4LO, smem_layout_b2q.inner, dtype=cutlass.Float32 + ), + smem_layout_b2q.outer, + ) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2hi_tf) + tCrA2lo_q = tiled_mma2q.make_fragment_A(sA2lo_tf) + tCrB2q_hi = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2q_hi.iterator, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + sB2q_hi.layout, + ) + ) + tCrB2q_lo = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2q_lo.iterator, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + sB2q_lo.layout, + ) + ) + if warp_idx == 0: + # GEMM2 half 0 fused with the output combination: [D4 | ΔS_0..63] += A2 x [Bcoef; Kbar_0..63] + tiled_mma2q.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // KB_MMA): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + if cutlass.const_expr(not BF16_MMA): + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_lo[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2lo_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2) + tiled_mma2h.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // KB_MMA): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + if cutlass.const_expr(not BF16_MMA): + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2lo_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2lo[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2b) + if cutlass.const_expr(PROFILE_STAGES): + tp_e = read_globaltimer() + + # ---- outputs of the new tokens: o_t = O0_t + D4[v, t] (columns 0..15 of GEMM2 half 0), + # warpgroup w writes + # new tokens j = w mod 4 + cute.arch.mbarrier_wait(mbar2, 0) + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq38 = read_globaltimer() + _, _gq38, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq38 + i_n) * 64 + 38] = _tq38 - t_stage0 + tiled_ld4 = tcgen05.make_tmem_copy(ld16, tD4) + thr_ld4 = tiled_ld4.get_slice(wg_tidx) + rD4 = cute.make_rmem_tensor(thr_ld4.partition_D(tD4).shape, cutlass.Float32) + cute.copy(ld16, thr_ld4.partition_S(tD4), rD4) + cute.arch.fence_view_async_tmem_load() + rD4_flat = cute.make_tensor(rD4.iterator, cute.make_layout((T_PAD,))) + for j_new in cutlass.range(1 + NUM_SPEC, unroll_full=True): + if wg == j_new % 4: + t_new = commit_len + j_new + o_val = cutlass.Float32(0.0) + for tt in cutlass.range(T_PAD, unroll_full=True): + if tt == t_new: + o_val = rD1_flat[T_PAD + tt] + rD4_flat[tt] + if cutlass.const_expr(FUSE_OUTPUT_NORM): + sVall[j_new * V + v_row] = cutlass.Float32(cutlass.BFloat16(o_val)) + else: + o[0, bos + t_new, i_hv, v_row] = cutlass.BFloat16(o_val) + if cutlass.const_expr(PROFILE_STAGES): + tp_f = read_globaltimer() + # ---- committed-state store through TMA: D2 slab -> swizzled smem box (row per thread, + # 16 B chunks land in distinct banks under the 128B swizzle) -> one bulk tensor store + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_wait() + if cutlass.const_expr(not FUSE_OUTPUT_NORM): + # bf16 path has no epilogue to overlap with: every warp stores its own slab + # (per-warp 32x32 smem transpose; the TMA store is used on the fused-norm path). + xp_base = warp_idx * 32 * XP_STRIDE + cute.arch.mbarrier_wait(mbar2, 0) + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(ld32, tD2_ld[(None, None, None, j)], rS) + cute.arch.fence_view_async_tmem_load() + cute.arch.barrier() + for i in cutlass.range(32, unroll_full=True): + sXp[xp_base + in_warp_tid * XP_STRIDE + i] = rS_flat[i] + cute.arch.sync_warp() + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + for j in cutlass.range(32, unroll_full=True): + val = sXp[xp_base + j * XP_STRIDE + in_warp_tid] + if cutlass.const_expr(USE_FLAT_LAYOUT): + ht[h0_idx, (warp_idx % 4) * 32 + j, wg * 32 + in_warp_tid] = val.to( + STATE_DTYPE + ) + else: + ht[slot, i_hv, (warp_idx % 4) * 32 + j, wg * 32 + in_warp_tid] = val.to( + STATE_DTYPE + ) + if cutlass.const_expr(FUSE_OUTPUT_NORM): + # All 16 warps produced disjoint value rows above. One CTA + # barrier makes the BF16-rounded rows visible before assigning + # one complete token to each warp. + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq39 = read_globaltimer() + _, _gq39, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq39 + i_n) * 64 + 39] = _tq39 - t_stage0 + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + tp_g = read_globaltimer() + # ---- committed state: the state-owning warps of warpgroups 0/1 stage k columns 0..63 + # (GEMM2 half 0, mbar2), those of warpgroups 2/3 columns 64..127 (half 1, mbar2b) + # into the swizzled smem boxes; the first staging warp of each pair issues the TMA + # stores once its half is visible (named barriers 6 / 7) and waits for the bulk + # groups at the very end. The warps without state rows run the norm epilogue. + if wg < 2: + cute.arch.mbarrier_wait(mbar2, 0) + else: + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + tp_l = read_globaltimer() + for sl in cutlass.range_constexpr(4): + if wg == sl: + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + cute.copy(ld32, tD2_ld[(None, None, None, sl)], rS) + cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr(STATE_BF16): + # bf16 box: chunks [4 (sl % 2), +4) of box sl // 2, round to nearest + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (sl % 2) + c) ^ (wg_tidx % 8) + off_b = ( + (sl // 2) * ((V // SPLIT_V) * 64) + + 64 * (wg_tidx - (V // SPLIT_V) * vh) + + 8 * chunk + ) + for i in cutlass.range(8, unroll_full=True): + sBoxhS[off_b + i] = cutlass.BFloat16(rS_flat[8 * c + i]) + else: + for c in cutlass.range(8, unroll_full=True): + off_b = ( + sl * ((V // SPLIT_V) * 32) + + 32 * (wg_tidx - (V // SPLIT_V) * vh) + + ((4 * c) ^ vx) + ) + for i in cutlass.range(4, unroll_full=True): + sBoxf[off_b + i] = rS_flat[4 * c + i] + # last TMEM reads of the kernel: warps 1..15 signal BAR_TMEM_FREE, warp 0 frees TMEM + tcgen05_fence_before_thread_sync() + if warp_idx > 0: + cute.arch.barrier_arrive(barrier_id=BAR_TMEM_FREE, number_of_threads=NUM_THREADS) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq23 = read_globaltimer() + _, _gq23, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq23 + i_n) * 64 + 23] = _tq23 - t_stage0 + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 256: + _tq33 = read_globaltimer() + _, _gq33, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq33 + i_n) * 64 + 33] = _tq33 - t_stage0 + cute.arch.fence_view_async_shared() + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + if wg < 2: + cute.arch.barrier(barrier_id=BAR_STORE_LO, number_of_threads=256 // SPLIT_V) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq24 = read_globaltimer() + _, _gq24, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq24 + i_n) * 64 + 24] = ( + _tq24 - t_stage0 + ) + if tidx == (V // 2) * vh: + for s in cutlass.range_constexpr(NUM_SLABS // 2): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq25 = read_globaltimer() + _, _gq25, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq25 + i_n) * 64 + 25] = ( + _tq25 - t_stage0 + ) + else: + cute.arch.barrier(barrier_id=BAR_STORE_HI, number_of_threads=256 // SPLIT_V) + if tidx == 256 + (V // 2) * vh: + for s in cutlass.range_constexpr(NUM_SLABS // 2, NUM_SLABS): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 256: + _tq34 = read_globaltimer() + _, _gq34, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq34 + i_n) * 64 + 34] = ( + _tq34 - t_stage0 + ) + if warp_idx == 0: + # all TMEM reads are done (BAR_TMEM_FREE): free the columns while the store drains + cute.arch.barrier(barrier_id=BAR_TMEM_FREE, number_of_threads=NUM_THREADS) + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(not COMPACT): + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_ALLOC) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 64: + _tq35 = read_globaltimer() + _, _gq35, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq35 + i_n) * 64 + 35] = _tq35 - t_stage0 + # norm epilogue (CTA 0) on the warps that own no state rows: (warp % 4) in {2, 3} + token_idx = (warp_idx // 4) * 2 + (warp_idx % 4 - 2) + if (vh == 0) & ((warp_idx % 4) // 2 == 1): + while token_idx < 1 + NUM_SPEC: + # gated values p_v = o_v * w_v * sigmoid(g_v); the two warp reductions (sum of + # squares, max |p_v|) are interleaved. bf16 rounding is monotonic, so + # amax(bf16(p_v * inv_rms)) = bf16(max|p_v| * inv_rms). + pre_values = cute.make_rmem_tensor( + cute.make_layout((V // 32,), stride=(1,)), cutlass.Float32 + ) + rms_partial = cutlass.Float32(0.0) + pre_max = cutlass.Float32(0.0) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + value = sVall[token_idx * V + v_idx] + gate = sOnG[token_idx * V + v_idx] + sigmoid_gate = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) + ) + p_v = value * sOnW[v_idx] * sigmoid_gate + pre_values[i] = p_v + rms_partial += value * value + pre_max = cute.arch.fmax(pre_max, cute.arch.fmax(p_v, -p_v)) + for offset in [16, 8, 4, 2, 1]: + rms_partial += cute.arch.shuffle_sync_bfly( + rms_partial, offset=offset, mask=-1, mask_and_clamp=31 + ) + pre_max = cute.arch.fmax( + pre_max, + cute.arch.shuffle_sync_bfly( + pre_max, offset=offset, mask=-1, mask_and_clamp=31 + ), + ) + inv_rms = cute.math.rsqrt( + rms_partial / cutlass.Float32(V) + cutlass.Float32(onorm_eps), + fastmath=True, + ) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 64: + _tq19 = read_globaltimer() + _, _gq19, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq19 + i_n) * 64 + 19] = ( + _tq19 - t_stage0 + ) + output_row = i_n * (1 + NUM_SPEC) + token_idx + normalized_values = cute.make_rmem_tensor( + cute.make_layout((V // 32,), stride=(1,)), cutlass.Float32 + ) + for i in cutlass.range(V // 32, unroll_full=True): + normalized = pre_values[i] * inv_rms + if cutlass.const_expr(QUANTIZE_OUTPUT): + # Preserve the old standalone RMSNorm's BF16 output + # boundary in registers. The value goes straight to + # MXFP8 below; it is never written to or reloaded from + # HBM as BF16. + normalized = cutlass.Float32(cutlass.BFloat16(normalized)) + normalized_values[i] = normalized + if cutlass.const_expr(QUANTIZE_OUTPUT): + output_amax = cutlass.Float32(cutlass.BFloat16(pre_max * inv_rms)) + output_amax = cute.arch.fmax(output_amax, cutlass.Float32(1.0e-10)) + output_sf = (output_amax * cutlass.Float32(1.0 / 448.0)).to( + cutlass.Float8E8M0FNU + ) + quant_scale = cute.arch.rcp_approx(cutlass.Float32(output_sf)) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[output_row, i_hv * V + v_idx] = cutlass.Float8E4M3FN( + normalized_values[i] * quant_scale + ) + if in_warp_tid == 0: + scale_bytes = cute.make_tensor( + cute.recast_ptr( + output_scale.iterator, + dtype=cutlass.Float8E8M0FNU, + ), + cute.make_layout(output_scale.shape), + ) + scale_offset = ( + (output_row // 128) * HV * 512 + + i_hv * 512 + + (output_row % 32) * 16 + + ((output_row % 128) // 32) * 4 + ) + for replica in cutlass.range_constexpr(4): + scale_bytes[scale_offset + replica] = output_sf + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 64: + _tq20 = read_globaltimer() + _, _gq20, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq20 + i_n) * 64 + 20] = ( + _tq20 - t_stage0 + ) + else: + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[0, output_row, i_hv, v_idx] = cutlass.BFloat16(normalized_values[i]) + token_idx = token_idx + NUM_THREADS // 32 + if (tidx == (V // 2) * vh) | (tidx == 256 + (V // 2) * vh): + cute.arch.cp_async_bulk_wait_group(0, read=True) + if cutlass.const_expr(PROFILE_STAGES): + tp_m = read_globaltimer() + if tidx == 0: + _, grid_n2, _ = cute.arch.grid_dim() + tb2 = ((vh * HV + i_hv) * grid_n2 + i_n) * 64 + stage_timing[tb2 + 15] = tp_l - t_stage0 + stage_timing[tb2 + 16] = tp_m - t_stage0 + if tidx == 256: + _, _gq26, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq26 + i_n) * 64 + 26] = tp_l - t_stage0 + if cutlass.const_expr(PROFILE_STAGES): + tp_h = read_globaltimer() + tp_i = tp_h + if tidx == 32: + _, _gq21, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq21 + i_n) * 64 + 21] = tp_h - t_stage0 + if tidx == 224: + _, _gq22, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq22 + i_n) * 64 + 22] = tp_h - t_stage0 + if cutlass.const_expr(PROFILE_STAGES): + cute.arch.barrier() + t_stage2 = read_globaltimer() + if tidx == 256: + _, _gq27, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq27 + i_n) * 64 + 27] = t_stage2 - t_stage0 + if tidx == 0: + _, grid_n, _ = cute.arch.grid_dim() + timing_base = ((vh * HV + i_hv) * grid_n + i_n) * 64 + stage_timing[timing_base + 0] = t_stage1 - t_stage0 + stage_timing[timing_base + 1] = t_stage2 - t_stage1 + stage_timing[timing_base + 2] = t_stage2 - t_stage0 + stage_timing[timing_base + 3] = t_stage0 + stage_timing[timing_base + 4] = tp_a - t_stage0 + stage_timing[timing_base + 5] = tp_b - t_stage0 + stage_timing[timing_base + 6] = tp_c - t_stage0 + stage_timing[timing_base + 7] = tp_d - t_stage0 + stage_timing[timing_base + 8] = tp_e - t_stage0 + stage_timing[timing_base + 9] = tp_f - t_stage0 + stage_timing[timing_base + 10] = tp_g - t_stage0 + stage_timing[timing_base + 11] = tp_h - t_stage0 + stage_timing[timing_base + 12] = tp_i - t_stage0 + stage_timing[timing_base + 13] = tp_j - t_stage0 + stage_timing[timing_base + 14] = tp_k - t_stage0 + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + if cutlass.const_expr(not FUSE_OUTPUT_NORM): + # (the fused-norm path released TMEM right after the state staging) + if warp_idx == 0: + cute.arch.relinquish_tmem_alloc_permit() + tmem_ptr_d = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + cute.arch.dealloc_tmem(tmem_ptr_d, TMEM_ALLOC) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + t_exit = read_globaltimer() + _, grid_n3, _ = cute.arch.grid_dim() + tb3 = ((vh * HV + i_hv) * grid_n3 + i_n) * 64 + stage_timing[tb3 + 17] = t_exit - t_stage0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode_2cta.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode_2cta.py new file mode 100644 index 000000000000..56d772115261 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode_2cta.py @@ -0,0 +1,1256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""KDA multi-token verify with the MXFP8 output, sized for two CTAs per SM (large batches). + +Computes exactly what ``kda_mtp_decode_v2.kda_decode_mtp_v2_kernel`` computes on its runtime +contract (packed rows, fused gated RMSNorm, 3xTF32 chunk recurrence, fp32 or bf16 state, bf16 +projections and conv caches, fp32 replay gates, one CTA per (head, request) tile, forget gate +from the projection) with the same arithmetic in the same order, so the normalized outputs, +committed state, conv windows and replay caches are bitwise identical. A tile still takes one +512-thread CTA; the footprint shrinks so that two tiles share an SM when the grid has more tiles +than SMs: + +* **Tensor memory: 256 columns.** S0 enters GEMM1 as two sequential A operands in the same 128 + columns (S0 hi for GEMM1, then S0 lo for the K~ sweep, issued in v2's order), and the GEMM2 + accumulator (S0 * decay) is seeded into those columns once GEMM1 is done. D4 | D2 | D1 | D3 + = 208 columns. A bf16 state is exact in tf32, so (as in v2) it has no lo part and GEMM1 is + the hi pass alone. +* **Shared memory: <= 114 KB** (113.7 KB at seven drafts). One 64 KB region holds, in turn, the + S0 TMA boxes (S0 is moved into registers and its hi part into tensor memory right after the + precompute), the GEMM3 tiles with the cumulative gates and e^{G_c}, the GEMM1 / Kbar / Bcoef + tiles once GEMM3 is done, V' over the GEMM1 tiles once GEMM1 is done, and finally the + committed-state TMA boxes. The raw conv inputs stay bf16 (the committed window and replay + rows land through ``cp.async``; the new tokens are indexed by token, not position, so their + placement does not wait for the accepted count), q is staged for the new tokens only, the + gate warps keep their token's raw gate, beta, dt_bias and A_log in registers (gate of four + consecutive channels per lane), the epilogue keeps its output-norm operands in registers, and + the GEMM3 result reuses the conv-weight words. Kbar and V' go straight into their swizzled + tiles from the row-per-thread producers. +* **Registers: 64 per thread** (``min_blocks_per_mp=2``); D1 is read from tensor memory in the + 16-column pieces each consumer needs, and a bf16 S0 row stays in registers as packed bf16x2 + words. The tensor-memory allocation permit is released right after the allocation so that the + SM's other CTA can allocate. +* **Named barriers: ids <= 4, all compile-time constants** (see ``BAR_WG0``), which SM107 needs + to co-schedule two CTAs. + +The epilogue quantizes the normalized output with the scale layout and the arithmetic of +``kda_decode_mtp_kernel``'s MXFP8 epilogue (the operand of the Rubin o_proj GEMM). + +Alone on an SM a tile is not faster than in the v2 kernel (with an fp32 state it takes ~0.8 us +longer: the precompute reads bf16 and GEMM1's two A operands are serialized); two tiles per SM +finish ~20% (fp32 state) to ~27% (bf16 state) sooner than two waves of v2 tiles, so the kernel pays +off for grids with more tiles than SMs. On SM107 (bf16 state, MXFP8 output) a tile alone on an SM +also finishes 0.1-0.8 us sooner than in ``kda_decode_mtp_kernel`` (one CTA or a cluster pair). +""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import Int64 +from cutlass.cutlass_dsl import T, dsl_user_op + +from .kda_mtp_decode_v2 import ( + NUM_THREADS, + T_PAD, + TILE_K, + _st_asm, + cvt_bf16x2, + gaddr, + griddepcontrol_launch_dependents, + griddepcontrol_wait, + ld_row4, + round_tf32, + saddr, + st_global_b32x2, + st_row4, + st_shared_f32x4, + swz64_k16, + tcgen05_fence_after_thread_sync, + tcgen05_fence_before_thread_sync, + unpack4, +) +from .kda_mtp_decode_v2 import ld_shared_f32x4 as _ld_shared_f32x4 + +TMEM_COLS = 256 +# The 64 KB operand region (floats), used in turn as +# P0: the S0 TMA boxes (128 B swizzled rows: four 128 x 32 fp32 or two 128 x 64 bf16 boxes), +# until S0 is in registers +# P1: A3 = [Kf; Qf hi; Qf lo; -] 64 x 128, B3 hi / lo 16 x 128 (GEMM3), the cumulative gates +# and e^{G_c} (both kept through P2 / P3) +# P2 (once GEMM3 is done): B1 = [K~hi; Q~; K~lo] 48 x 128, B1k = K~hi 16 x 128, then +# [Bcoef hi; Kbar hi] and [Bcoef lo; Kbar lo] (16 + 128) x 16 each (GEMM2 B tiles) +# P3 (once GEMM1 is done): A2 hi / lo = V'^T 128 x 16 over B1 +# P4 (once GEMM2 is done): the committed-state TMA boxes +OFF_A3 = 0 +OFF_B3HI = 8192 +OFF_B3LO = 10240 +OFF_B1 = 0 +OFF_B1K = 6144 +OFF_B4HI = 8192 +OFF_B2HI = 8448 +OFF_B4LO = 10496 +OFF_B2LO = 10752 +OFF_E = 12800 +OFF_EC = 14720 +OFF_A2HI = 0 +OFF_A2LO = 2048 +REGION_FLOATS = 16384 +# Named barrier ids (0 is the CTA-wide barrier), all compile-time constants. An SM reserves +# (max id + 1) named barriers per CTA and SM107 co-schedules two CTAs only up to id 4, so the ids +# are reused across phases that CTA-wide barriers separate: BAR_WG0 + wg (1..4) are the +# per-warpgroup GEMM3 barriers, the store / TMEM-release barriers of the epilogue take ids 1..3 +# again. +BAR_WG0 = 1 +BAR_STORE_LO = 1 # committed-state boxes staged, threads 0..255 +BAR_STORE_HI = 2 # committed-state boxes staged, threads 256..511 +BAR_TMEM_FREE = 3 # last TMEM reads done -> warp 0 frees the columns + + +@dsl_user_op +def _cp_async_8(smem_addr, gmem_addr, *, loc=None, ip=None): + _st_asm( + "cp.async.ca.shared.global [$0], [$1], 8;", + "r,l", + [ + cutlass.Int32(smem_addr).ir_value(loc=loc, ip=ip), + Int64(gmem_addr).ir_value(loc=loc, ip=ip), + ], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cp_async_16(smem_addr, gmem_addr, *, loc=None, ip=None): + _st_asm( + "cp.async.cg.shared.global [$0], [$1], 16;", + "r,l", + [ + cutlass.Int32(smem_addr).ir_value(loc=loc, ip=ip), + Int64(gmem_addr).ir_value(loc=loc, ip=ip), + ], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cp_async_commit(*, loc=None, ip=None): + _st_asm("cp.async.commit_group;", "", [], loc=loc, ip=ip) + + +@dsl_user_op +def _cp_async_wait_all(*, loc=None, ip=None): + _st_asm("cp.async.wait_group 0;", "", [], loc=loc, ip=ip) + + +@dsl_user_op +def _tmem_ld_x1(taddr, *, loc=None, ip=None) -> cutlass.Float32: + """One 32-bit tensor-memory word per lane (``tcgen05.ld 32x32b.x1``) at ``taddr``.""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Int32(taddr).ir_value(loc=loc, ip=ip)], + "tcgen05.ld.sync.aligned.32x32b.x1.b32 {$0}, [$1];", + "=f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +def _ld_shared_b32x2(addr): + from .kda_mtp_decode_v2 import _asm_words + + return _asm_words( + 2, + "ld.shared.v2.b32 {$0, $1}, [$2];", + "=r,=r,r", + [cutlass.Int32(addr).ir_value()], + loc=None, + ip=None, + ) + + +def _st_shared_b32x2(addr, a, b): + _st_asm( + "st.shared.v2.b32 [$0], {$1, $2};", + "r,r,r", + [cutlass.Int32(addr).ir_value(), cutlass.Int32(a).ir_value(), cutlass.Int32(b).ir_value()], + ) + + +def _s0_value(r_state, r_state_w, idx, packed): + """fp32 value of S0 column ``idx`` of this thread's row: ``r_state[idx]``, or for a packed bf16 + state the bf16 value ``idx`` of the bf16x2 words ``r_state_w`` (even values in the low half).""" + if cutlass.const_expr(packed): + w = r_state_w[idx // 2] + if cutlass.const_expr(idx % 2 == 0): + w = w << 16 + else: + w = w & cutlass.Int32(-65536) + return w.bitcast(cutlass.Float32) + return r_state[idx] + + +def _load_new_token_qkv( + x_q, x_k, x_v, tok_row, hk_off, hv_off, k4, r_new, use_flat_layout, i_h, i_hv +): + """v2's new-token q / k / v loads (raw bf16 words, four channels per lane).""" + if cutlass.const_expr(use_flat_layout): + ld_row4(x_q, (0, tok_row, hk_off + k4), r_new, 0) + ld_row4(x_k, (0, tok_row, hk_off + k4), r_new, 4) + ld_row4(x_v, (0, tok_row, hv_off + k4), r_new, 8) + else: + ld_row4(x_q, (0, tok_row, i_h, k4), r_new, 0) + ld_row4(x_k, (0, tok_row, i_h, k4), r_new, 4) + ld_row4(x_v, (0, tok_row, i_hv, k4), r_new, 8) + + +@cute.kernel +def kda_decode_mtp_2cta_kernel( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + o8: cute.Tensor, + osf: cute.Tensor, + tiled_mma1: cute.TiledMma, + tiled_mma3: cute.TiledMma, + tiled_mma1k: cute.TiledMma, + smem_layout_b1: cute.ComposedLayout, + smem_layout_a2: cute.ComposedLayout, + smem_layout_a3: cute.ComposedLayout, + smem_layout_b3: cute.ComposedLayout, + smem_layout_b1k: cute.ComposedLayout, + tiled_mma4: cute.TiledMma, + tiled_mma2h: cute.TiledMma, + smem_layout_b2h: cute.ComposedLayout, + tiled_mma2q: cute.TiledMma, + smem_layout_b2q: cute.ComposedLayout, + tiled_mma_f32a: cute.TiledMma, + tma_atom_l: cute.CopyAtom, + mS_tma: cute.Tensor, + box_layout: cute.ComposedLayout, + tma_atom_t: cute.CopyAtom, + mT_tma: cute.Tensor, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + PDL: cutlass.Constexpr[bool], + EARLY_TRIGGER: cutlass.Constexpr[bool] = False, +): + """See the module docstring. Grid ``(HV, N, 1)``, one (head, request) tile per CTA. + + ``EARLY_TRIGGER`` (c240, with ``PDL``): release the dependent grid right after + ``griddepcontrol.wait`` (as ``kda_mtp_decode_v2`` does with ``PDL``), for a dependent that only + streams weights before its own wait (the mid-M o_proj GEMM with the all-reduce push). + + The epilogue writes the MXFP8 operand of the Rubin CuTe o_proj GEMM with the arithmetic of + ``kda_decode_mtp_kernel``'s ``QUANTIZE_OUTPUT`` epilogue: ``o8`` (E4M3 ``[T, HV * V]``) and + ``osf`` (uint8, UE8M0 1x128 scales in the R128c4 layout: four copies per (row, head)).""" + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + i_hv, i_n, _ = cute.arch.block_idx() + i_h = i_hv + slot32 = ssm_state_indices[i_n] + hk_off = i_h * K + hv_off = i_hv * V + N1 = 3 * T_PAD + K_CACHE_DTYPE = k_cache.element_type + V_CACHE_DTYPE = v_cache.element_type + BETA_CACHE_DTYPE = beta_cache.element_type + t_max = 2 * NUM_SPEC + 1 + NEW_TOK = 1 + NUM_SPEC + NUM_OLD = KERNEL_WIDTH - 1 + NUM_SPEC # committed window + replay rows of the raw staging + vec_size = TILE_K // 32 + REPLAY_WARP0 = NUM_THREADS // 64 + TMA_WARP = NUM_THREADS // 32 - 1 + assert NUM_SPEC <= 7 and t_max <= T_PAD + assert NEW_TOK <= REPLAY_WARP0 and REPLAY_WARP0 + NUM_SPEC <= TMA_WARP + assert K == 128 and V == 128 and KERNEL_WIDTH == 4 and NUM_THREADS == 512 + STATE_DTYPE = h0.element_type + STATE_BF16 = STATE_DTYPE == cutlass.BFloat16 + assert STATE_BF16 or STATE_DTYPE == cutlass.Float32, "fp32 or bf16 recurrent state" + # a bf16 state is exact in tf32: no S0 lo part (v2's S0_LO) + S0_LO = not STATE_BF16 + BOX_COLS = 128 // (STATE_DTYPE.width // 8) # state columns per 128 B swizzled box row + NUM_SLABS = K // BOX_COLS + assert g_cache.element_type == cutlass.Float32, "fp32 replay gates (cp.async into sG)" + assert cs_q.element_type == cutlass.BFloat16 and x_q.element_type == cutlass.BFloat16 + k_weight_base = KERNEL_WIDTH * K + v_weight_base = 2 * KERNEL_WIDTH * K + conv_weight_elems = 2 * KERNEL_WIDTH * K + KERNEL_WIDTH * V + + smem = cutlass.utils.SmemAllocator() + sQn = smem.allocate_tensor(cutlass.Float32, cute.make_layout((NEW_TOK * K,)), 16) + sK = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max, K), stride=(K, 1)), 16) + sG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max, K), stride=(K, 1)), 16) + sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) + # Conv weights during the precompute; the GEMM3 result [L; B hi; B lo; -] (64 x 16) after. + sConvW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((conv_weight_elems,)), 16) + sLB = cute.make_tensor(sConvW.iterator, cute.make_layout((4 * T_PAD, T_PAD), stride=(T_PAD, 1))) + # Raw conv inputs (bf16): [window; replayed drafts] and the new tokens, per q / k / v. + # rows [0, NUM_OLD): window + replay rows, then the new tokens' rows + sXq = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(((NUM_OLD + NEW_TOK) * K,)), 16) + sXk = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(((NUM_OLD + NEW_TOK) * K,)), 16) + sXv = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(((NUM_OLD + NEW_TOK) * V,)), 16) + sXrq, sXrk, sXrv = sXq, sXk, sXv + sXnq = cute.make_tensor(sXq.iterator + NUM_OLD * K, cute.make_layout((NEW_TOK * K,))) + sXnk = cute.make_tensor(sXk.iterator + NUM_OLD * K, cute.make_layout((NEW_TOK * K,))) + sXnv = cute.make_tensor(sXv.iterator + NUM_OLD * V, cute.make_layout((NEW_TOK * V,))) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + tmem_holder = smem.allocate_array(cutlass.Int32, 1) + mbar1a = smem.allocate_array(cutlass.Int64, 1) + mbar1 = smem.allocate_array(cutlass.Int64, 1) + mbar3 = smem.allocate_array(cutlass.Int64, 1) + mbar2 = smem.allocate_array(cutlass.Int64, 1) + mbar2b = smem.allocate_array(cutlass.Int64, 1) + mbar_ld = smem.allocate_array(cutlass.Int64, 1) + sRegion = smem.allocate_tensor(cutlass.Float32, cute.make_layout((REGION_FLOATS,)), 1024) + + def region_tile(off, layout): + return cute.make_tensor( + cute.recast_ptr(sRegion.iterator + off, layout.inner, dtype=cutlass.Float32), + layout.outer, + ) + + sA3 = region_tile(OFF_A3, smem_layout_a3) + sB3hi = region_tile(OFF_B3HI, smem_layout_b3) + sB3lo = region_tile(OFF_B3LO, smem_layout_b3) + sB1 = region_tile(OFF_B1, smem_layout_b1) + sB1k = region_tile(OFF_B1K, smem_layout_b1k) + sA2hi = region_tile(OFF_A2HI, smem_layout_a2) + sA2lo = region_tile(OFF_A2LO, smem_layout_a2) + # Flat views for closed-form swizzled element writes. + sA3f = cute.make_tensor(sRegion.iterator + OFF_A3, cute.make_layout((64 * 128,))) + sB3hif = cute.make_tensor(sRegion.iterator + OFF_B3HI, cute.make_layout((16 * 128,))) + sB3lof = cute.make_tensor(sRegion.iterator + OFF_B3LO, cute.make_layout((16 * 128,))) + sB1f = cute.make_tensor(sRegion.iterator + OFF_B1, cute.make_layout((48 * 128,))) + sB1kf = cute.make_tensor(sRegion.iterator + OFF_B1K, cute.make_layout((T_PAD * 128,))) + sB4hif = cute.make_tensor(sRegion.iterator + OFF_B4HI, cute.make_layout((T_PAD * T_PAD,))) + sB4lof = cute.make_tensor(sRegion.iterator + OFF_B4LO, cute.make_layout((T_PAD * T_PAD,))) + # Cumulative gates G_t (P1 / P2) and the S0 / committed-state boxes (P0 / P4). + sE = cute.make_tensor(sRegion.iterator + OFF_E, cute.make_layout((t_max * K,))) + sEc = cute.make_tensor(sRegion.iterator + OFF_EC, cute.make_layout((K,))) + sBoxf = cute.make_tensor(sRegion.iterator, cute.make_layout((V * K,))) + sBoxh = cute.make_tensor( + cute.recast_ptr(sRegion.iterator, dtype=cutlass.BFloat16), cute.make_layout((V * K,)) + ) + + # ---- load phase (v2's schedule): slot-independent loads under the slot load; then the + # slot-dependent ones. Window and replay raw conv rows and replay gates go to smem through + # cp.async (no registers); replay k / v / beta (converted on placement) and the new + # tokens through registers. + ch = tidx % 128 + quad = tidx // 128 + k4 = 4 * in_warp_tid + r_w = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + if quad == 0: + ld_row4(w_q, (hk_off + ch, 0), r_w, 0) + if quad == 1: + ld_row4(w_k, (hk_off + ch, 0), r_w, 0) + if quad == 2: + ld_row4(w_v, (hv_off + ch, 0), r_w, 0) + # gate warps (REPLAY_WARP0 + j): this lane's dt_bias channels [4 l, 4 l + 4) and A_log + j_gate = warp_idx - REPLAY_WARP0 + tok_row_g = i_n * NEW_TOK + j_gate + r_dtb = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + r_gate = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + r_gb = cute.make_rmem_tensor(cute.make_layout((2,)), cutlass.Float32) + if (warp_idx >= REPLAY_WARP0) & (j_gate < NEW_TOK): + ld_row4(dt_bias, (hk_off + k4,), r_dtb, 0) + r_gb[1] = cutlass.Float32(A_log[i_h]) + r_new = cute.make_rmem_tensor(cute.make_layout((3 * 4,)), cutlass.Int32) + tok_row = i_n * NEW_TOK + warp_idx + if cutlass.const_expr(not PDL): + if warp_idx < NEW_TOK: + _load_new_token_qkv( + x_q, x_k, x_v, tok_row, hk_off, hv_off, k4, r_new, USE_FLAT_LAYOUT, i_h, i_hv + ) + if (warp_idx >= REPLAY_WARP0) & (j_gate < NEW_TOK): + ld_row4(g, (0, tok_row_g, i_hv, k4), r_gate, 0) + r_gb[0] = cutlass.Float32(beta[0, tok_row_g, i_hv]) + # The TMA warp, the first to use an mbarrier, initializes them; every other warp first touches + # one after the placement barrier. + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mbar1a, 1) + cute.arch.mbarrier_init(mbar1, 1) + cute.arch.mbarrier_init(mbar3, 1) + cute.arch.mbarrier_init(mbar2, 1) + cute.arch.mbarrier_init(mbar2b, 1) + cute.arch.mbarrier_init(mbar_ld, 1) + cute.arch.mbarrier_init_fence() + cute.arch.fence_view_async_shared() + slot = Int64(slot32) + h0_idx = slot * HV + i_hv + # byte step between consecutive window positions (mode 2) of the conv states + cs_step = [cute.crd2idx((0, 0, 1), t.layout) * 2 for t in (cs_q, cs_k, cs_v)] + rep = warp_idx - REPLAY_WARP0 + r_rep = cute.make_rmem_tensor(cute.make_layout((2 * 4,)), cutlass.Int32) + r_repb = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Float32) + if (warp_idx >= REPLAY_WARP0) & (rep < NUM_SPEC): + ld_row4(k_cache, (slot, rep, hk_off + k4), r_rep, 0) + ld_row4(v_cache, (slot, rep, hv_off + k4), r_rep, 4) + r_repb[0] = cutlass.Float32(beta_cache[slot, rep, i_hv]) + # every pending entry lands in its row; rows from the accepted count on are overwritten + # by the new tokens' precompute before they are read + _cp_async_16(saddr(sG, rep * K + k4), gaddr(g_cache, (slot, rep, hk_off + k4))) + rr = KERNEL_WIDTH - 1 + rep + _cp_async_8(saddr(sXrq, rr * K + k4), gaddr(cs_q, (slot, hk_off + k4, rr))) + _cp_async_8(saddr(sXrk, rr * K + k4), gaddr(cs_k, (slot, hk_off + k4, rr))) + _cp_async_8(saddr(sXrv, rr * V + k4), gaddr(cs_v, (slot, hv_off + k4, rr))) + _cp_async_commit() + if warp_idx == TMA_WARP: + # committed window rows 0 .. KERNEL_WIDTH - 2: one base per conv state + wq_base = gaddr(cs_q, (slot, hk_off + k4, 0)) + wk_base = gaddr(cs_k, (slot, hk_off + k4, 0)) + wv_base = gaddr(cs_v, (slot, hv_off + k4, 0)) + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + _cp_async_8(saddr(sXrq, w * K + k4), wq_base + w * cs_step[0]) + _cp_async_8(saddr(sXrk, w * K + k4), wk_base + w * cs_step[1]) + _cp_async_8(saddr(sXrv, w * V + k4), wv_base + w * cs_step[2]) + _cp_async_commit() + commit_len = num_accepted_tokens[slot] + # S0 into the region's boxes (consumed right after the precompute) + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ld, V * K * (STATE_DTYPE.width // 8)) + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sTl_h = cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gS_h_all = cute.local_tile(mS_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gS_h = gS_h_all[(None, None, 0, h0_idx)] + else: + gS_h = gS_h_all[(None, None, 0, (i_hv, slot))] + tSl_h, tGl_h = cpasync.tma_partition( + tma_atom_l, + 0, + cute.make_layout(1), + cute.group_modes(sTl_h, 0, 2), + cute.group_modes(gS_h, 0, 2), + ) + if warp_idx == TMA_WARP: + cute.copy(tma_atom_l, tGl_h, tSl_h, tma_bar_ptr=mbar_ld) + if cutlass.const_expr(PDL): + griddepcontrol_wait() + if cutlass.const_expr(EARLY_TRIGGER): + griddepcontrol_launch_dependents() + if warp_idx < NEW_TOK: + _load_new_token_qkv( + x_q, x_k, x_v, tok_row, hk_off, hv_off, k4, r_new, USE_FLAT_LAYOUT, i_h, i_hv + ) + if (warp_idx >= REPLAY_WARP0) & (j_gate < NEW_TOK): + ld_row4(g, (0, tok_row_g, i_hv, k4), r_gate, 0) + r_gb[0] = cutlass.Float32(beta[0, tok_row_g, i_hv]) + if warp_idx == 0: + cute.arch.alloc_tmem(TMEM_COLS, tmem_holder) + # the SM's other CTA allocates its 256 columns as soon as we hold ours + cute.arch.relinquish_tmem_alloc_permit() + # ---- placement + if quad < 3: + wv = unpack4(w_q.element_type, r_w, 0) + for w in cutlass.range_constexpr(KERNEL_WIDTH): + sConvW[quad * (KERNEL_WIDTH * K) + w * K + ch] = wv[w] + if commit_len > NUM_SPEC: + commit_len = cutlass.Int32(NUM_SPEC) + T_loop = commit_len + NEW_TOK + if warp_idx < NEW_TOK: + # raw rows of new token j stay bf16 and are indexed by j + _st_shared_b32x2(saddr(sXnq, warp_idx * K + k4), r_new[0], r_new[1]) + _st_shared_b32x2(saddr(sXnk, warp_idx * K + k4), r_new[4], r_new[5]) + _st_shared_b32x2(saddr(sXnv, warp_idx * V + k4), r_new[8], r_new[9]) + if (warp_idx >= REPLAY_WARP0) & (rep < commit_len): + st_shared_f32x4(saddr(sK, rep * K + k4), *unpack4(K_CACHE_DTYPE, r_rep, 0)) + st_shared_f32x4(saddr(sVall, rep * V + k4), *unpack4(V_CACHE_DTYPE, r_rep, 4)) + if in_warp_tid == 0: + sBeta[rep] = r_repb[0] + if ((warp_idx >= REPLAY_WARP0) & (rep < NUM_SPEC)) | (warp_idx == TMA_WARP): + _cp_async_wait_all() + if tidx < K: + for _z in cutlass.range(T_loop, t_max, 1): + sG[_z, tidx] = cutlass.Float32(0.0) + sVall[_z * V + tidx] = cutlass.Float32(0.0) + if tidx == 0: + sBeta[_z] = cutlass.Float32(0.0) + wg = warp_idx // 4 + wg_tidx = tidx % 128 + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + cute.arch.barrier() + + # ---- precompute (v2's arithmetic; raw rows read as bf16 from the split staging) + half = NUM_THREADS // 64 + role = warp_idx // half + j = warp_idx % half + if j < NEW_TOK: + t = commit_len + j + # raw row of tap w of token j: j + w - 3 >= 0 -> new token row, else window / replay row + x_row = [] + for w in cutlass.range_constexpr(KERNEL_WIDTH): + r_new_idx = j + w - (KERNEL_WIDTH - 1) + old_row = commit_len + j + w + is_new = cutlass.Int32(r_new_idx >= 0) + x_row.append(old_row + is_new * (NUM_OLD + r_new_idx - old_row)) + if role == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_conv_0 = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(KERNEL_WIDTH): + r_conv_0 += cutlass.Float32(sXq[x_row[w] * K + k_idx]) * sConvW[w * K + k_idx] + r_conv = ( + cutlass.Float32(sXk[x_row[0] * K + k_idx]) + * sConvW[k_weight_base + 0 * K + k_idx] + ) + for w in cutlass.range_constexpr(1, KERNEL_WIDTH): + r_conv += ( + cutlass.Float32(sXk[x_row[w] * K + k_idx]) + * sConvW[k_weight_base + w * K + k_idx] + ) + e0 = cute.math.exp(-r_conv_0, fastmath=True) + r_q[i] = r_conv_0 * cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) + r_conv = r_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) + ) + r_k[i] = r_conv + sum_q = 0.0 + sum_k = 0.0 + for i in cutlass.range(vec_size, unroll_full=True): + sum_q += r_q[i] * r_q[i] + sum_k += r_k[i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=-1, mask_and_clamp=31 + ) + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_q_scaled = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + sQn[j * K + k_idx] = r_q[i] * rnorm_q_scaled + sK[t, k_idx] = r_k[i] * rnorm_k + else: + # gate of channels [4 l, 4 l + 4) from this warp's registers (v2's per-channel math) + r_exp_A = cute.math.exp(r_gb[1], fastmath=True) + g_raw = unpack4(g.element_type, r_gate, 0) + dtb = unpack4(dt_bias.element_type, r_dtb, 0) + gates = [] + for i in cutlass.range_constexpr(4): + r_g_raw = g_raw[i] + r_g_raw = r_g_raw + dtb[i] + exp_A_x = r_exp_A * r_g_raw + sigmoid_val = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) + ) + gates.append(lower_bound * sigmoid_val) + st_shared_f32x4(saddr(sG, t * K + k4), *gates) + r_beta_t = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_gb[0], fastmath=True) + ) + if in_warp_tid == 0: + sBeta[t] = r_beta_t + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid + _v_conv = 0.0 + for w in cutlass.range_constexpr(KERNEL_WIDTH): + _v_conv += ( + cutlass.Float32(sXv[x_row[w] * V + v_idx]) + * sConvW[v_weight_base + w * V + v_idx] + ) + _v_conv = _v_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) + ) + sVall[t * V + v_idx] = _v_conv + + # ================================================================= + # Chunk (WY-form) recurrence on tensor cores. + # ================================================================= + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + pool = cutlass.utils.TmemBufferPool(tmem_ptr, TMEM_COLS) + tD4_frag = tiled_mma4.make_fragment_C(tiled_mma4.partition_shape_C((V, T_PAD))) + tD4 = pool.allocate_tensor(tD4_frag.layout, cutlass.Float32) + tD2h_frag = tiled_mma2h.make_fragment_C(tiled_mma2h.partition_shape_C((V, K // 2))) + tD2h0 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) + tD2h1 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) + tD2q_frag = tiled_mma2q.make_fragment_C(tiled_mma2q.partition_shape_C((V, T_PAD + K // 2))) + tD2q = cute.make_tensor(tD4.iterator, tD2q_frag.layout) + tA_frag = tiled_mma_f32a.make_fragment_A(tiled_mma_f32a.partition_shape_A((V, K))) + # S0 hi, then S0 lo, then the GEMM2 accumulator: the same 128 columns + tX = cute.make_tensor(tD2h0.iterator, tA_frag.layout) + tD1_frag = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C((V, N1))) + tD1 = pool.allocate_tensor(tD1_frag.layout, cutlass.Float32) + tD3_frag = tiled_mma3.make_fragment_C(tiled_mma3.partition_shape_C((4 * T_PAD, T_PAD))) + tD3 = pool.allocate_tensor(tD3_frag.layout, cutlass.Float32) + + st8 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x8), cutlass.Float32) + ld8 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x8), cutlass.Float32) + ld16 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_st8 = tcgen05.make_tmem_copy(st8, tX).get_slice(wg_tidx) + tX_st8 = thr_st8.partition_D(tX) + r8 = cute.make_rmem_tensor(thr_st8.partition_S(tX).shape[:-1], cutlass.Float32) + thr_ld8 = tcgen05.make_tmem_copy(ld8, tX).get_slice(wg_tidx) + tX_ld8 = thr_ld8.partition_S(tX) + tS_list = [] + tG_list = [] + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sT_h = cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gT_h_all = cute.local_tile(mT_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gT_h = gT_h_all[(None, None, 0, h0_idx)] + else: + gT_h = gT_h_all[(None, None, 0, (i_hv, slot))] + tS_h, tG_h = cpasync.tma_partition( + tma_atom_t, + 0, + cute.make_layout(1), + cute.group_modes(sT_h, 0, 2), + cute.group_modes(gT_h, 0, 2), + ) + tS_list.append(tS_h) + tG_list.append(tG_h) + + # ---- S0 row v = wg_tidx, columns [32 wg, +32): from its TMA box (v2's conflict-free 16 B + # chunks) into registers, its hi part into X; the boxes are free after the barrier, which + # also publishes the precompute + v_row = wg_tidx + vx = 4 * (wg_tidx % 8) + r_state = cute.make_rmem_tensor(cute.make_layout((32,)), cutlass.Float32) + # a bf16 state stays packed: bf16x2 words, value 2 i in the low half of word i + r_state_w = cute.make_rmem_tensor(cute.make_layout((16,)), cutlass.Int32) + cute.arch.mbarrier_wait(mbar_ld, 0) + if cutlass.const_expr(STATE_BF16): + # bf16 box: 64 columns per 128 B row, so this warpgroup's 32 columns are the 16 B + # chunks [4 (wg % 2), +4) of box wg // 2 (v2), kept as raw bf16x2 words + for c in cutlass.range_constexpr(4): + chunk = (4 * (wg % 2) + c) ^ (wg_tidx % 8) + off_l = (wg // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + words = _ld_shared_f32x4(saddr(sBoxh, off_l)) + for i in cutlass.range_constexpr(4): + r_state_w[4 * c + i] = words[i].bitcast(cutlass.Int32) + else: + for c in cutlass.range(8, unroll_full=True): + off_l = wg * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range(4, unroll_full=True): + r_state[4 * c + i] = sBoxf[off_l + i] + for jj in cutlass.range_constexpr(4): + if wg == jj: + for c in cutlass.range_constexpr(4): + for i in cutlass.range_constexpr(8): + if cutlass.const_expr(S0_LO): + r8[i] = round_tf32(r_state[8 * c + i]) + else: + r8[i] = _s0_value(r_state, r_state_w, 8 * c + i, STATE_BF16) + cute.copy(st8, r8, tX_st8[(None, None, None, 4 * jj + c)]) + cute.arch.fence_view_async_tmem_store() + cute.arch.barrier() + + # ---- cumulative gates (v2), q of the replayed drafts is zero + kq = tidx % K + quarter = tidx // K + t_mid = T_loop // 2 + r_G = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + g_run = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t < t_max: + g_run = g_run + sG[t, kq] + r_G[t] = g_run + g_mid = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t == t_mid: + g_mid = r_G[t] + exp_clamp = cutlass.Float32(87.0) + ka = kq // 32 + kc = kq % 32 + b1_base = ka * (3 * T_PAD * 32) + a3_base = ka * (4 * T_PAD * 32) + b3_base = ka * (T_PAD * 32) + b1k_base = ka * (T_PAD * 32) + r_b1 = cute.make_rmem_tensor(cute.make_layout((3 * (T_PAD // 4),)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + kt_hi = cutlass.Float32(0.0) + kt_lo = cutlass.Float32(0.0) + qt_r = cutlass.Float32(0.0) + kf = cutlass.Float32(0.0) + qf_hi = cutlass.Float32(0.0) + qf_lo = cutlass.Float32(0.0) + kb_hi = cutlass.Float32(0.0) + kb_lo = cutlass.Float32(0.0) + if t < T_loop: + d_t = r_G[t] - g_mid + e_t = cute.math.exp(-cute.arch.fmax(-d_t, -exp_clamp), fastmath=True) + e_bt = cute.math.exp(-cute.arch.fmax(d_t, -exp_clamp), fastmath=True) + eg_t = cute.math.exp(r_G[t], fastmath=True) + sE[t * K + kq] = r_G[t] + if t == commit_len: + sEc[kq] = eg_t + k_t = sK[t, kq] + q_t = cutlass.Float32(0.0) + if t >= commit_len: + q_t = sQn[(t - commit_len) * K + kq] + kt = k_t * eg_t + kt_hi = round_tf32(kt) + qt_r = round_tf32(q_t * eg_t) + kf = round_tf32(k_t * e_t) + qf = q_t * e_t + qf_hi = round_tf32(qf) + kb = k_t * e_bt + kb_hi = round_tf32(kb) + kt_lo = kt - kt_hi + qf_lo = qf - qf_hi + kb_lo = kb - kb_hi + xc = kc ^ (4 * (t % 8)) + sA3f[a3_base + 32 * t + xc] = kf + sA3f[a3_base + 32 * (T_PAD + t) + xc] = qf_hi + sA3f[a3_base + 32 * (2 * T_PAD + t) + xc] = qf_lo + sB3hif[b3_base + 32 * t + xc] = kb_hi + sB3lof[b3_base + 32 * t + xc] = kb_lo + r_b1[3 * (t // 4)] = kt_hi + r_b1[3 * (t // 4) + 1] = qt_r + r_b1[3 * (t // 4) + 2] = kt_lo + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + tCrA3 = tiled_mma3.make_fragment_A( + cute.make_tensor( + cute.recast_ptr(sA3.iterator, smem_layout_a3.inner, dtype=cutlass.TFloat32), + sA3.layout, + ) + ) + tCrB3hi = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3hi.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3hi.layout, + ) + ) + tCrB3lo = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3lo.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3lo.layout, + ) + ) + if warp_idx == 0: + tiled_mma3.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3hi[(None, None, kb_i, 0)], tD3 + ) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3lo[(None, None, kb_i, 0)], tD3 + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar3) + + # GEMM1 operand rows over the GEMM3 tiles once GEMM3 has read them; D3 -> smem. One warp per + # warpgroup polls the mbarrier, the others wait at the warpgroup's barrier. + if warp_idx % 4 == 1: + cute.arch.mbarrier_wait(mbar3, 0) + # one compile-time barrier id per warpgroup: a barrier id held in a register makes the CTA + # reserve all 16 named barriers, which leaves room for one CTA per SM on SM107 + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.arch.barrier(barrier_id=BAR_WG0 + jj, number_of_threads=128) + tcgen05_fence_after_thread_sync() + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + xc = kc ^ (4 * (t % 8)) + sB1f[b1_base + 32 * t + xc] = r_b1[3 * (t // 4)] + if cutlass.const_expr(S0_LO): + sB1kf[b1k_base + 32 * t + xc] = r_b1[3 * (t // 4)] + sB1f[b1_base + 32 * (T_PAD + t) + xc] = r_b1[3 * (t // 4) + 1] + sB1f[b1_base + 32 * (2 * T_PAD + t) + xc] = r_b1[3 * (t // 4) + 2] + ld3 = cute.make_copy_atom( + tcgen05.Ld16x128bOp(tcgen05.Repetition(T_PAD * 32 // 128)), cutlass.Float32 + ) + thr_ld3 = tcgen05.make_tmem_copy(ld3, tD3).get_slice(wg_tidx) + tD3_ld = thr_ld3.partition_S(tD3) + sLB3 = cute.make_tensor( + sLB.iterator, + cute.make_layout(((4 * T_PAD, T_PAD), 1, 1), stride=((T_PAD, 1), 0, 0)), + ) + tD3_sm = thr_ld3.partition_D(sLB3) + rD3 = cute.make_rmem_tensor(tD3_sm.shape, cutlass.Float32) + if wg == 0: + cute.copy(ld3, tD3_ld, rD3) + cute.arch.fence_view_async_tmem_load() + cute.autovec_copy(rD3, tD3_sm) + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + tCrB1 = tiled_mma1.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1.iterator, smem_layout_b1.inner, dtype=cutlass.TFloat32), + sB1.layout, + ) + ) + tCrB1k = tiled_mma1k.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1k.iterator, smem_layout_b1k.inner, dtype=cutlass.TFloat32), + sB1k.layout, + ) + ) + tD1k = cute.make_tensor( + tD1.iterator, + tiled_mma1k.make_fragment_C(tiled_mma1k.partition_shape_C((V, T_PAD))).layout, + ) + tX_tf = cute.make_tensor(cute.recast_ptr(tX.iterator, dtype=cutlass.TFloat32), tX.layout) + if warp_idx == 0: + tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm(tiled_mma1, tD1, tX_tf[(None, None, kb_i)], tCrB1[(None, None, kb_i, 0)], tD1) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, True) + with cute.arch.elect_one(): + if cutlass.const_expr(S0_LO): + tcgen05.commit(mbar1a) + else: + tcgen05.commit(mbar1) + # ---- under GEMM1: D4 zero, Bcoef tiles (warpgroups 1..2, from sLB), Kbar tiles, the new + # tokens' replay-cache rows and conv-window commits + st16 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_stD4 = tcgen05.make_tmem_copy(st16, tD4).get_slice(wg_tidx) + tD4_st = thr_stD4.partition_D(tD4) + r_zero = cute.make_rmem_tensor(thr_stD4.partition_S(tD4).shape, cutlass.Float32) + for i in cutlass.range(cute.size(r_zero), unroll_full=True): + r_zero[i] = cutlass.Float32(0.0) + if wg == 0: + cute.copy(st16, r_zero, tD4_st) + if (tidx >= 128) & (tidx < 128 + T_PAD * T_PAD): + t4 = (tidx - 128) // T_PAD + s4 = (tidx - 128) % T_PAD + bc = cutlass.Float32(0.0) + if (t4 < T_loop) & (s4 <= t4): + bc = sLB[T_PAD + t4, s4] + sLB[2 * T_PAD + t4, s4] + bc_hi = round_tf32(bc) + o_b4 = swz64_k16(t4, s4) + sB4hif[o_b4] = bc_hi + sB4lof[o_b4] = bc - bc_hi + # Kbar[k, s] = K[s, k] e^{G_c[k] - G_s[k]} (s <= c, else 0): thread (k = kq, quarter) owns + # columns s = 4 quarter + {0..3}, i.e. one 16 B chunk of row k of the SW64 hi / lo tiles + sB2hi_row = saddr(sRegion, OFF_B2HI) + sB2lo_row = saddr(sRegion, OFF_B2LO) + # (quarters whose first column is past the committed count hold zeros only) + kb_0 = cutlass.Float32(0.0) + kb_1 = cutlass.Float32(0.0) + kb_2 = cutlass.Float32(0.0) + kb_3 = cutlass.Float32(0.0) + s_q = 4 * quarter + if s_q <= commit_len: + e_c = sE[commit_len * K + kq] + kb_0 = sK[s_q, kq] * cute.math.exp(e_c - sE[s_q * K + kq], fastmath=True) + if s_q + 1 <= commit_len: + kb_1 = sK[s_q + 1, kq] * cute.math.exp(e_c - sE[(s_q + 1) * K + kq], fastmath=True) + if s_q + 2 <= commit_len: + kb_2 = sK[s_q + 2, kq] * cute.math.exp(e_c - sE[(s_q + 2) * K + kq], fastmath=True) + if s_q + 3 <= commit_len: + kb_3 = sK[s_q + 3, kq] * cute.math.exp(e_c - sE[(s_q + 3) * K + kq], fastmath=True) + kb_vals = [kb_0, kb_1, kb_2, kb_3] + kb_his = [round_tf32(v) for v in kb_vals] + o_b2 = swz64_k16(kq, 4 * quarter) + st_shared_f32x4(sB2hi_row + 4 * o_b2, *kb_his) + st_shared_f32x4(sB2lo_row + 4 * o_b2, *[kb_vals[i] - kb_his[i] for i in range(4)]) + if (warp_idx >= 1) & (warp_idx <= NUM_SPEC): + p = warp_idx - 1 + t_c = commit_len + 1 + p + st_row4(k_cache, (slot, p, hk_off + k4), _ld_shared_f32x4(saddr(sK, t_c * K + k4))) + st_row4(g_cache, (slot, p, hk_off + k4), _ld_shared_f32x4(saddr(sG, t_c * K + k4))) + st_row4(v_cache, (slot, p, hv_off + k4), _ld_shared_f32x4(saddr(sVall, t_c * V + k4))) + # draft p is new token 1 + p + wq = _ld_shared_b32x2(saddr(sXnq, (1 + p) * K + k4)) + st_global_b32x2(gaddr(cs_q, (slot, hk_off + k4, KERNEL_WIDTH - 1 + p)), wq[0], wq[1]) + wk = _ld_shared_b32x2(saddr(sXnk, (1 + p) * K + k4)) + st_global_b32x2(gaddr(cs_k, (slot, hk_off + k4, KERNEL_WIDTH - 1 + p)), wk[0], wk[1]) + wvv = _ld_shared_b32x2(saddr(sXnv, (1 + p) * V + k4)) + st_global_b32x2(gaddr(cs_v, (slot, hv_off + k4, KERNEL_WIDTH - 1 + p)), wvv[0], wvv[1]) + if in_warp_tid == 0: + beta_cache[slot, p, i_hv] = sBeta[t_c].to(BETA_CACHE_DTYPE) + # committed window: raw rows commit_len + 1 + w (w = 2 is new token 0) + for sec in cutlass.range_constexpr(3): + if warp_idx == REPLAY_WARP0 + sec: + cs_sec = (cs_q, cs_k, cs_v)[sec] + old_sec = (sXrq, sXrk, sXrv)[sec] + new_sec = (sXnq, sXnk, sXnv)[sec] + off_sec = hv_off if sec == 2 else hk_off + sec_base = gaddr(cs_sec, (slot, off_sec + k4, 0)) + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + if cutlass.const_expr(w + 1 >= KERNEL_WIDTH - 1): + src = saddr(new_sec, (w + 1 - (KERNEL_WIDTH - 1)) * K + k4) + else: + src = saddr(old_sec, (commit_len + 1 + w) * K + k4) + ww = _ld_shared_b32x2(src) + st_global_b32x2(sec_base + w * cs_step[sec], ww[0], ww[1]) + # ---- GEMM1's S0 lo sweep once the hi pass has read X + if cutlass.const_expr(S0_LO): + cute.arch.mbarrier_wait(mbar1a, 0) + tcgen05_fence_after_thread_sync() + for jj in cutlass.range_constexpr(4): + if wg == jj: + for c in cutlass.range_constexpr(4): + for i in cutlass.range_constexpr(8): + s_hi = round_tf32(r_state[8 * c + i]) + r8[i] = r_state[8 * c + i] - s_hi + cute.copy(st8, r8, tX_st8[(None, None, None, 4 * jj + c)]) + cute.arch.fence_view_async_tmem_store() + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + if warp_idx == 0: + tiled_mma1k.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1k, + tD1k, + tX_tf[(None, None, kb_i)], + tCrB1k[(None, None, kb_i, 0)], + tD1k, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar1) + # ---- GEMM2 accumulator seed S0 * e^{G_c} into X once GEMM1 is done + cute.arch.mbarrier_wait(mbar1, 0) + tcgen05_fence_after_thread_sync() + for jj in cutlass.range_constexpr(4): + if wg == jj: + for c in cutlass.range_constexpr(4): + for i in cutlass.range_constexpr(8): + r8[i] = ( + _s0_value(r_state, r_state_w, 8 * c + i, STATE_BF16) + * sEc[jj * 32 + 8 * c + i] + ) + cute.copy(st8, r8, tX_st8[(None, None, None, 4 * jj + c)]) + cute.arch.fence_view_async_tmem_store() + # ---- warpgroup 0: V' by forward substitution (row v per thread) straight into the A2 tiles + thr_ld1 = tcgen05.make_tmem_copy(ld16, tD1).get_slice(wg_tidx) + tD1_ld = thr_ld1.partition_S(tD1) + # 16-column chunks of D1 are mode 1 of the ld16 partition: [K~hi | Q~ | K~lo] + rD1_shape = thr_ld1.partition_D(tD1)[(None, 0, None, None)].shape + rD1p = cute.make_rmem_tensor(rD1_shape, cutlass.Float32) + rD1q = cute.make_rmem_tensor(rD1_shape, cutlass.Float32) + rD1p_flat = cute.make_tensor(rD1p.iterator, cute.make_layout((T_PAD,))) + rD1q_flat = cute.make_tensor(rD1q.iterator, cute.make_layout((T_PAD,))) + if wg == 0: + cute.copy(ld16, tD1_ld[(None, 0, None, None)], rD1p) + cute.copy(ld16, tD1_ld[(None, 2, None, None)], rD1q) + cute.arch.fence_view_async_tmem_load() + r_Y = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for s in cutlass.range(T_PAD, unroll_full=True): + y = cutlass.Float32(0.0) + if s < t_max: + u0 = rD1p_flat[s] + rD1q_flat[s] + y = sBeta[s] * (sVall[s * V + v_row] - u0) + r_Y[s] = y + r_V = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + vp = cutlass.Float32(0.0) + if t < t_max: + acc0 = cutlass.Float32(0.0) + acc1 = cutlass.Float32(0.0) + acc2 = cutlass.Float32(0.0) + acc3 = cutlass.Float32(0.0) + for s in range(t): + if s % 4 == 0: + acc0 = acc0 + sLB[t, s] * r_V[s] + elif s % 4 == 1: + acc1 = acc1 + sLB[t, s] * r_V[s] + elif s % 4 == 2: + acc2 = acc2 + sLB[t, s] * r_V[s] + else: + acc3 = acc3 + sLB[t, s] * r_V[s] + vp = r_Y[t] - sBeta[t] * ((acc0 + acc1) + (acc2 + acc3)) + r_V[t] = vp + sA2hi_row = saddr(sRegion, OFF_A2HI) + sA2lo_row = saddr(sRegion, OFF_A2LO) + for c in cutlass.range_constexpr(4): + his = [round_tf32(r_V[4 * c + i]) for i in range(4)] + o_a2 = swz64_k16(v_row, 4 * c) + st_shared_f32x4(sA2hi_row + 4 * o_a2, *his) + st_shared_f32x4(sA2lo_row + 4 * o_a2, *[r_V[4 * c + i] - his[i] for i in range(4)]) + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + sA2hi_tf = cute.make_tensor( + cute.recast_ptr(sA2hi.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), sA2hi.layout + ) + sA2lo_tf = cute.make_tensor( + cute.recast_ptr(sA2lo.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), sA2lo.layout + ) + tCrA2hi = tiled_mma2h.make_fragment_A(sA2hi_tf) + tCrA2lo = tiled_mma2h.make_fragment_A(sA2lo_tf) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B2HI + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrB2lo_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B2LO + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2hi_tf) + tCrA2lo_q = tiled_mma2q.make_fragment_A(sA2lo_tf) + tCrB2q_hi = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B4HI, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + tCrB2q_lo = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B4LO, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + if warp_idx == 0: + tiled_mma2q.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_lo[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2lo_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2) + tiled_mma2h.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2lo_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2lo[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2b) + # output-gate operand of the epilogue (warps 2, 3 of each warpgroup: one token each) + token_idx = (warp_idx // 4) * 2 + (warp_idx % 4 - 2) + r_og = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + r_ow = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + if ((warp_idx % 4) // 2 == 1) & (token_idx < NEW_TOK): + for i in cutlass.range(V // 32, unroll_full=True): + r_og[i] = cutlass.Float32( + onorm_g[0, i_n * NEW_TOK + token_idx, i_hv, i * 32 + in_warp_tid] + ) + r_ow[i] = cutlass.Float32(onorm_weight[i * 32 + in_warp_tid]) + # ---- outputs of the new tokens: o_t = Q-part of D1 + D4 (v2), bf16-rounded into sVall rows; + # warpgroup wg reads column t of both for its tokens wg and wg + 4 (lane = row v) + cute.arch.mbarrier_wait(mbar2, 0) + tcgen05_fence_after_thread_sync() + lane_hi = (warp_idx % 4) << 21 + d4_col = cutlass.Int32(tD4.iterator.toint()) + lane_hi + d1q_col = cutlass.Int32(tD1.iterator.toint()) + lane_hi + T_PAD + for jw in cutlass.range_constexpr((NEW_TOK + 3) // 4): + j_out = wg + 4 * jw + if j_out < NEW_TOK: + t_col = commit_len + j_out + o_q = _tmem_ld_x1(d1q_col + t_col) + o_d = _tmem_ld_x1(d4_col + t_col) + cute.arch.fence_view_async_tmem_load() + sVall[j_out * V + v_row] = cutlass.Float32(cutlass.BFloat16(o_q + o_d)) + cute.arch.barrier() + # ---- committed state: warpgroup wg stages columns [32 wg, +32) into box wg of the region + # (free once GEMM2 is done), then one thread per half issues the TMA stores (v2) + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + for jj in cutlass.range_constexpr(4): + if wg == jj: + for cc in cutlass.range_constexpr(4): + cute.copy(ld8, tX_ld8[(None, None, None, 4 * jj + cc)], r8) + cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr(STATE_BF16): + # bf16 box jj // 2, 16 B chunk 4 (jj % 2) + cc, round to nearest (v2's + # cvt.rn per value, two values per cvt.rn.bf16x2) + chunk = (4 * (jj % 2) + cc) ^ (wg_tidx % 8) + off_b = (jj // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + st_shared_f32x4( + saddr(sBoxh, off_b), + *[cvt_bf16x2(r8[2 * i + 1], r8[2 * i]) for i in range(4)], + ) + else: + for h in cutlass.range_constexpr(2): + c = 2 * cc + h + off_b = jj * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range_constexpr(4): + sBoxf[off_b + i] = r8[4 * h + i] + tcgen05_fence_before_thread_sync() + if warp_idx > 0: + cute.arch.barrier_arrive(barrier_id=BAR_TMEM_FREE, number_of_threads=NUM_THREADS) + cute.arch.fence_view_async_shared() + if wg < 2: + cute.arch.barrier(barrier_id=BAR_STORE_LO, number_of_threads=256) + if tidx == 0: + for sl in cutlass.range_constexpr(NUM_SLABS // 2): + cute.copy(tma_atom_t, tS_list[sl], tG_list[sl]) + cute.arch.cp_async_bulk_commit_group() + else: + cute.arch.barrier(barrier_id=BAR_STORE_HI, number_of_threads=256) + if tidx == 256: + for sl in cutlass.range_constexpr(NUM_SLABS // 2, NUM_SLABS): + cute.copy(tma_atom_t, tS_list[sl], tG_list[sl]) + cute.arch.cp_async_bulk_commit_group() + if warp_idx == 0: + cute.arch.barrier(barrier_id=BAR_TMEM_FREE, number_of_threads=NUM_THREADS) + tcgen05_fence_after_thread_sync() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS) + # ---- gated RMSNorm epilogue (v2) with kda_decode_mtp_kernel's QUANTIZE_OUTPUT epilogue: + # BF16-rounded normalized values, UE8M0 scale of bf16(max |p_v| * inv_rms) / 448 (bf16 + # rounding is monotonic), E4M3 values, the scale byte replicated four times in the R128c4 + # layout + if ((warp_idx % 4) // 2 == 1) & (token_idx < NEW_TOK): + pre_values = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + rms_partial = cutlass.Float32(0.0) + pre_max = cutlass.Float32(0.0) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + value = sVall[token_idx * V + v_idx] + gate = r_og[i] + sigmoid_gate = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) + ) + p_v = value * r_ow[i] * sigmoid_gate + pre_values[i] = p_v + rms_partial += value * value + pre_max = cute.arch.fmax(pre_max, cute.arch.fmax(p_v, -p_v)) + for offset in [16, 8, 4, 2, 1]: + rms_partial += cute.arch.shuffle_sync_bfly( + rms_partial, offset=offset, mask=-1, mask_and_clamp=31 + ) + pre_max = cute.arch.fmax( + pre_max, + cute.arch.shuffle_sync_bfly(pre_max, offset=offset, mask=-1, mask_and_clamp=31), + ) + inv_rms = cute.math.rsqrt( + rms_partial / cutlass.Float32(V) + cutlass.Float32(onorm_eps), fastmath=True + ) + output_row = i_n * NEW_TOK + token_idx + output_amax = cutlass.Float32(cutlass.BFloat16(pre_max * inv_rms)) + output_amax = cute.arch.fmax(output_amax, cutlass.Float32(1.0e-10)) + output_sf = (output_amax * cutlass.Float32(1.0 / 448.0)).to(cutlass.Float8E8M0FNU) + quant_scale = cute.arch.rcp_approx(cutlass.Float32(output_sf)) + for i in cutlass.range(V // 32, unroll_full=True): + normalized = cutlass.Float32(cutlass.BFloat16(pre_values[i] * inv_rms)) + o8[output_row, hv_off + i * 32 + in_warp_tid] = cutlass.Float8E4M3FN( + normalized * quant_scale + ) + if in_warp_tid == 0: + scale_bytes = cute.make_tensor( + cute.recast_ptr(osf.iterator, dtype=cutlass.Float8E8M0FNU), + cute.make_layout(osf.shape), + ) + scale_offset = ( + (output_row // 128) * HV * 512 + + i_hv * 512 + + (output_row % 32) * 16 + + ((output_row % 128) // 32) * 4 + ) + for replica in cutlass.range_constexpr(4): + scale_bytes[scale_offset + replica] = output_sf + # the state stores must have read their smem boxes before the CTA exits + if (tidx == 0) | (tidx == 256): + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode_v2.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode_v2.py new file mode 100644 index 000000000000..1fad31030b82 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/kda_mtp_decode_v2.py @@ -0,0 +1,1614 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Latency-scheduled CuTe DSL kernel for KDA multi-token speculative verify (packed rows). + +Computes exactly what ``kda_mtp_decode.kda_decode_mtp_kernel`` computes on its runtime +contract -- packed token rows, fused gated RMSNorm with a bf16 output, 3xTF32 tcgen05 chunk +recurrence, fp32 or bf16 state, K = V = 128, conv width 4, NUM_SPEC <= 7 -- with the same +arithmetic in the same order, so outputs, committed state, conv pools and replay caches are +bitwise identical. See that module for the math (replay of accepted drafts, chunk WY form, +GEMM1/GEMM3/forward substitution/GEMM2, in-place state commit after the golden token, the +cluster-pair ``SPLIT_V == 2`` form). The difference is the schedule of the latency-bound +CTA: + +* **Load phase.** The state slot is the first load. The slot-independent operands (new-token + projections, conv weights, per-head constants, output-norm operands) are issued under it; the + slot-dependent ones (the S0 TMA, the committed conv windows and every pending replay entry, + accepted or not) as soon as it arrives. Each lane moves four consecutive channels of a row + (8 B / 16 B vector loads) and keeps the raw words in registers until the accepted count fixes + the row's smem position, so the CTA waits for one round trip per dependency level (slot, then + operands) instead of one per staging step. +* **Precompute** works in smem only, with the q and k conv / SiLU / L2-norm chains of a token + interleaved in one warp. +* **Tensor-core ordering.** The GEMM3 operand rows lie clear of the S0 boxes, so GEMM3 is issued + before S0 moves into TMEM and runs under that copy; GEMM1 follows once its rows are written over + the read boxes. While the tensor cores run, warps 1..15 stage Kbar, seed the GEMM2 + accumulator, build the B2 tiles and write the new tokens' replay-cache rows and conv-window + commits (vectorized, from smem). + +Optional (compile-time) variants, both bitwise identical to the plain kernel fed by +``k3_kda_fab``: + +* **PDL.** Launched as a programmatic dependent of the projection GEMM: the slot, accepted count, + S0 TMA, replay entries, conv windows and weights are loaded before ``griddepcontrol.wait``, + only the new-token rows (and the TMEM allocation) after it; the dependent grid is released + right after the wait. +* **Fused forget gate.** The kernel reads ``f_a`` (``k3_kda_fab`` in f_a mode) instead of the + gate and computes ``bf16(f_a @ W_fb[h]^T)`` itself: warps 8..15 each own 16 ``W_fb`` rows + (TMA-loaded into free scratch words up front) and run eight ``mma.sync`` k16 steps on f_a + fragments loaded with the other new-token rows, before the precompute barrier. +""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import Int64 +from cutlass.cutlass_dsl import T, dsl_user_op + +NUM_THREADS = 512 +TILE_K = 128 +# Chunk length padding (MMA N granularity); NUM_SPEC <= 7 -> T_loop <= 15. +T_PAD = 16 +# TMEM columns: D4 (16) + D2 halves (2 x 64) + S0 hi (128) + S0 lo (128) + D1 (48) + D3 (16) +TMEM_COLS = 512 +# One 114 KB scratch region hosts, at different times, the TMA boxes of S0 / the committed state +# and the swizzled MMA operand tiles. Offsets in floats, all multiples of 256 (1024 B) for the +# 128B-swizzle atoms. +SCRATCH_FLOATS = 29184 +OFF_B4HI = 10240 # 16 x 16 Bcoef hi, immediately followed by Kbar hi -> N=80 GEMM2 half-0 tile +OFF_B4LO = 12544 # 16 x 16 Bcoef lo, immediately followed by Kbar lo +OFF_B1K = 6144 # 16 x 128 K~hi (duplicate, N=16 tile for the S0_lo sweep; A2 reuses it) +OFF_B1 = 0 # 48 x 128 [K~hi; Q~; K~lo] +OFF_A2HI = 6144 # 128 x 16 V'^T hi +OFF_A2LO = 8192 # 128 x 16 V'^T lo +OFF_B2HI = 10496 # 128 x 16 Kbar hi +OFF_B2LO = 12800 # 128 x 16 Kbar lo +OFF_A3 = 16384 # 64 x 128 [Kf; Qf hi; Qf lo; -] (clear of the S0 boxes) +OFF_B3HI = 24576 # 16 x 128 Kb hi +OFF_B3LO = 26624 # 16 x 128 Kb lo +TP_STRIDE = 17 # padded row stride of the transposed (v|k, s) staging buffers +# FUSED_FB: this head's W_fb rows (128 x 128 BF16 = 32 KB, two 128B-swizzled TMA boxes of 64 +# columns) in scratch words that stay unused until the cumulative-gate stage writes the GEMM3 +# tiles (OFF_A3 onwards). +OFF_FB = 16384 +FB_BOX_COLS = 64 + + +def swz64_k16(n, k): + """Element offset in a K-major 64B-swizzled fp32 tile with 16 columns.""" + off = n * 16 + k + return off ^ (((off >> 5) & 3) << 2) + + +@dsl_user_op +def cvt_bf16x2(hi: cutlass.Float32, lo: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Float32: + """Pack two fp32 values into one bf16x2 word (round to nearest even, `lo` in the low half).""" + w = cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + return w.bitcast(cutlass.Float32) + + +@dsl_user_op +def tcgen05_fence_before_thread_sync(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "tcgen05.fence::before_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def tcgen05_fence_after_thread_sync(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "tcgen05.fence::after_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def round_tf32(x: cutlass.Float32) -> cutlass.Float32: + """Round to the tf32 grid (rna on mantissa bit 13) so hi/lo splits are exact.""" + b = x.bitcast(cutlass.Int32) + b = (b + cutlass.Int32(0x1000)) & cutlass.Int32(-8192) + return b.bitcast(cutlass.Float32) + + +def _asm_words(n, asm, constraints, operands, *, loc, ip): + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * n), + operands, + asm, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return [cutlass.Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(n)] + + +@dsl_user_op +def ld_global_b32x2(addr: Int64, *, loc=None, ip=None): + """Two 32-bit words (8 B, 8-byte aligned) from global memory, unconverted: the load has no + consumer until its value is placed, so a warp keeps all of its loads in flight.""" + return _asm_words( + 2, + "ld.global.v2.b32 {$0, $1}, [$2];", + "=r,=r,l", + [Int64(addr).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_global_b32x4(addr: Int64, *, loc=None, ip=None): + """Four 32-bit words (16 B, 16-byte aligned) from global memory, unconverted.""" + return _asm_words( + 4, + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + [Int64(addr).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_shared_f32x4(addr: cutlass.Int32, *, loc=None, ip=None): + """Four consecutive fp32 (16-byte aligned) from shared memory.""" + words = _asm_words( + 4, + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + [cutlass.Int32(addr).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + return [w.bitcast(cutlass.Float32) for w in words] + + +@dsl_user_op +def _st_asm(asm, constraints, operands, *, loc=None, ip=None): + llvm.inline_asm( + None, + operands, + asm, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def st_global_b32x2(addr, a, b): + _st_asm( + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + [Int64(addr).ir_value(), cutlass.Int32(a).ir_value(), cutlass.Int32(b).ir_value()], + ) + + +def st_global_b32x4(addr, a, b, c, d): + _st_asm( + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + [Int64(addr).ir_value()] + [cutlass.Int32(x).ir_value() for x in (a, b, c, d)], + ) + + +def st_shared_f32x4(addr, a, b, c, d): + _st_asm( + "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", + "r,f,f,f,f", + [cutlass.Int32(addr).ir_value()] + [cutlass.Float32(x).ir_value() for x in (a, b, c, d)], + ) + + +@dsl_user_op +def griddepcontrol_wait(*, loc=None, ip=None): + _st_asm("griddepcontrol.wait;", "", [], loc=loc, ip=ip) + + +@dsl_user_op +def griddepcontrol_launch_dependents(*, loc=None, ip=None): + _st_asm("griddepcontrol.launch_dependents;", "", [], loc=loc, ip=ip) + + +@dsl_user_op +def ldmatrix_x4(addr: cutlass.Int32, *, loc=None, ip=None): + """Four 8x8 b16 matrices from shared memory (``ldmatrix.sync.aligned.m8n8.x4``).""" + return _asm_words( + 4, + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + [cutlass.Int32(addr).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_global_b32(addr: Int64, *, loc=None, ip=None): + """One 32-bit word (4-byte aligned) from global memory, unconverted.""" + return _asm_words( + 1, "ld.global.b32 $0, [$1];", "=r,l", [Int64(addr).ir_value(loc=loc, ip=ip)], loc=loc, ip=ip + )[0] + + +@dsl_user_op +def mma_bf16_16816(a, b0, b1, c, *, loc=None, ip=None): + """``d[16x8] = a[16x16] @ b[16x8] + c`` (``mma.sync`` m16n8k16, BF16 inputs, FP32 accumulate).""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32()] * 4), + [cutlass.Int32(x).ir_value(loc=loc, ip=ip) for x in (*a, b0, b1)] + + [cutlass.Float32(x).ir_value(loc=loc, ip=ip) for x in c], + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return [cutlass.Float32(llvm.extractvalue(T.f32(), res, [i], loc=loc, ip=ip)) for i in range(4)] + + +def round_bf16(x: cutlass.Float32) -> cutlass.Float32: + """``x`` rounded to BF16 (nearest even, ``cvt.rn``) and widened back to FP32.""" + w = cvt_bf16x2(x, x).bitcast(cutlass.Int32) + return (w & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + + +def _load_fb_b_fragments(fa, i_n, NEW_TOK, lane, r_fb_b): + """B fragments of the f_b MMA (``m16n8k16``, K = 128 in eight k16 steps): column ``n`` holds + the f_a row of this request's token whose packed row index is ``n`` mod 8 (the column + k3_kda_fab's tail gives it); lane ``l`` holds rows ``k = 16 kk + 2 (l % 4) + {0, 1}`` and + ``+ 8`` of column ``l / 4``. Columns without a token of this request read the request's + last row here and are zeroed before the MMA.""" + jt = (lane // 4 - (i_n * NEW_TOK) % 8 + 8) % 8 + over = cutlass.Int32(jt >= NEW_TOK) + tok = i_n * NEW_TOK + jt - over * (jt - (NEW_TOK - 1)) + kq = 2 * (lane % 4) + for kk in range(8): + r_fb_b[2 * kk] = ld_global_b32(gaddr(fa, (0, tok, 16 * kk + kq))) + r_fb_b[2 * kk + 1] = ld_global_b32(gaddr(fa, (0, tok, 16 * kk + 8 + kq))) + + +def _load_new_token_row( + x_q, + x_k, + x_v, + g, + onorm_g, + beta, + tok_row, + i_h, + i_hv, + hk_off, + hv_off, + k4, + r_new, + r_newb, + use_flat_layout, + fused_fb, +): + """Issue one lane's loads of a new token's projections: q / k / v / gate (unless the gate is + computed in the kernel) / output gate four channels each into ``r_new``, beta into + ``r_newb``.""" + if cutlass.const_expr(use_flat_layout): + ld_row4(x_q, (0, tok_row, hk_off + k4), r_new, 0) + ld_row4(x_k, (0, tok_row, hk_off + k4), r_new, 4) + ld_row4(x_v, (0, tok_row, hv_off + k4), r_new, 8) + else: + ld_row4(x_q, (0, tok_row, i_h, k4), r_new, 0) + ld_row4(x_k, (0, tok_row, i_h, k4), r_new, 4) + ld_row4(x_v, (0, tok_row, i_hv, k4), r_new, 8) + if cutlass.const_expr(not fused_fb): + ld_row4(g, (0, tok_row, i_hv, k4), r_new, 12) + ld_row4(onorm_g, (0, tok_row, i_hv, k4), r_new, 16) + r_newb[0] = cutlass.Float32(beta[0, tok_row, i_hv]) + + +def gaddr(tensor, coord): + """Global address of `tensor[coord]`.""" + return (tensor.iterator + cute.crd2idx(coord, tensor.layout)).toint() + + +def saddr(tensor, offset): + """Shared-memory address of element `offset` of an smem tensor (row-major offset).""" + return (tensor.iterator + offset).toint() + + +def ld_row4(tensor, coord, words, base): + """Issue the load of four consecutive elements of `tensor` (16- or 32-bit) at `coord` into the + raw-word buffer `words[base : base + 4]` (two packed words for a 16-bit type).""" + if cutlass.const_expr(tensor.element_type.width == 16): + w = ld_global_b32x2(gaddr(tensor, coord)) + words[base] = w[0] + words[base + 1] = w[1] + else: + w = ld_global_b32x4(gaddr(tensor, coord)) + for i in range(4): + words[base + i] = w[i] + + +def unpack4(dtype, words, base): + """Four fp32 values from the raw words of `ld_row4` (bf16 -> fp32 is exact).""" + if cutlass.const_expr(dtype.width == 16): + assert dtype == cutlass.BFloat16, "16-bit operands are bf16" + w0 = words[base] + w1 = words[base + 1] + return [ + (w0 << 16).bitcast(cutlass.Float32), + (w0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + (w1 << 16).bitcast(cutlass.Float32), + (w1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + ] + return [words[base + i].bitcast(cutlass.Float32) for i in range(4)] + + +def st_row4(tensor, coord, v): + """Store four fp32 values `v` as consecutive elements of `tensor` (rounded to nearest even + for bf16, like a scalar `.to(BFloat16)`).""" + if cutlass.const_expr(tensor.element_type.width == 16): + assert tensor.element_type == cutlass.BFloat16, "16-bit caches are bf16" + st_global_b32x2( + gaddr(tensor, coord), + cvt_bf16x2(v[1], v[0]).bitcast(cutlass.Int32), + cvt_bf16x2(v[3], v[2]).bitcast(cutlass.Int32), + ) + else: + st_global_b32x4(gaddr(tensor, coord), *[x.bitcast(cutlass.Int32) for x in v]) + + +@cute.kernel +def kda_decode_mtp_v2_kernel( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + smem_qk_layout: cute.Layout, + tiled_mma1: cute.TiledMma, + tiled_mma3: cute.TiledMma, + tiled_mma1k: cute.TiledMma, + smem_layout_b1: cute.ComposedLayout, + smem_layout_a2: cute.ComposedLayout, + smem_layout_a3: cute.ComposedLayout, + smem_layout_b3: cute.ComposedLayout, + smem_layout_b1k: cute.ComposedLayout, + tiled_mma4: cute.TiledMma, + tma_atom_t: cute.CopyAtom, + mT_tma: cute.Tensor, + box_layout: cute.ComposedLayout, + tma_atom_l: cute.CopyAtom, + mS_tma: cute.Tensor, + tiled_mma2h: cute.TiledMma, + smem_layout_b2h: cute.ComposedLayout, + box_layout_t: cute.ComposedLayout, + tiled_mma2q: cute.TiledMma, + smem_layout_b2q: cute.ComposedLayout, + tiled_mma_f32a: cute.TiledMma, + fa: cute.Tensor, + tma_atom_fb: cute.CopyAtom, + mFb_tma: cute.Tensor, + fb_layout: cute.ComposedLayout, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + SPLIT_V: cutlass.Constexpr[int], + PDL: cutlass.Constexpr[bool], + FUSED_FB: cutlass.Constexpr[bool], +): + """KDA MTP verify, packed-token / fused-norm contract (see the module docstring). + + Grid ``(HV, N, SPLIT_V)``: one (head, request) tile per CTA, or per cluster pair of CTAs + with ``SPLIT_V == 2`` (each CTA seeds and stores half of the state rows). + + ``PDL``: launched as a programmatic dependent; everything that does not depend on this + layer's projections (slot, accepted count, state TMA, replay entries, conv windows, weights) + is loaded before ``griddepcontrol.wait``, the new-token rows and TMEM after it. + ``FUSED_FB``: ``fa`` holds the BF16 ``f_a`` rows ``[1, T, 128]`` instead of the forget gate; + the CTA computes its head's gate ``bf16(f_a @ W_fb[h]^T)`` with ``mma.sync`` (``W_fb`` rows + TMA-loaded via ``tma_atom_fb`` before ``griddepcontrol.wait``, f_a fragments loaded with the + other new-token rows): FP32 accumulation, the operands, positions and k order of + ``k3_kda_fab``'s tail, so bitwise its forget gate. + """ + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + i_hv, i_n, vh = cute.arch.block_idx() + i_h = i_hv + # The request's state slot is the first load of the kernel; its round trip overlaps the + # slot-independent loads issued below. + slot32 = ssm_state_indices[i_n] + hk_off = i_h * K + hv_off = i_hv * V + # Recurrent-state element type (fp32 or bf16). The TMA boxes have 128 B swizzled rows, + # i.e. BOX_COLS = 32 fp32 or 64 bf16 columns; the state tile is NUM_SLABS such boxes. + STATE_DTYPE = h0.element_type + STATE_BF16 = STATE_DTYPE == cutlass.BFloat16 + # 3xTF32 recurrence (hi/lo splits on every product that feeds the state); S0_LO: a fp32 + # state needs the lo part of S0 (a bf16 state is exact in tf32). + S0_LO = not STATE_BF16 + N1 = 3 * T_PAD # GEMM1 N: [K~hi; Q~; K~lo] + CONV_DTYPE = cs_q.element_type + K_CACHE_DTYPE = k_cache.element_type + V_CACHE_DTYPE = v_cache.element_type + BETA_CACHE_DTYPE = beta_cache.element_type + BOX_COLS = 128 // (STATE_DTYPE.width // 8) + NUM_SLABS = K // BOX_COLS + t_max = 2 * NUM_SPEC + 1 + NEW_TOK = 1 + NUM_SPEC + vec_size = TILE_K // 32 + # Load-phase warp roles: warp j < NEW_TOK stages new token j, warp REPLAY_WARP0 + r the + # pending draft r (all NUM_SPEC entries are loaded, the accepted ones placed), the last warp + # the committed conv windows and the S0 TMA. + REPLAY_WARP0 = NUM_THREADS // 64 + TMA_WARP = NUM_THREADS // 32 - 1 + assert NUM_SPEC <= 7 and t_max <= T_PAD, "chunk form pads the token axis to 16" + assert NEW_TOK <= REPLAY_WARP0 and REPLAY_WARP0 + NUM_SPEC <= TMA_WARP + assert K == 128 and V == 128 and KERNEL_WIDTH == 4 and NUM_THREADS == 512, ( + "the thread mappings assume K = V = 128, conv width 4 and 512 threads" + ) + k_weight_base = KERNEL_WIDTH * K + v_weight_base = 2 * KERNEL_WIDTH * K + conv_weight_elems = 2 * KERNEL_WIDTH * K + KERNEL_WIDTH * V + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sK = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sG = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) + sConvW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((conv_weight_elems,)), 16) + sOnG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((NEW_TOK * V,)), 16) + sOnW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V,)), 16) + # Staged per-token inputs (raw conv sequence with the committed window + # prepended in rows 0..W-2, raw gates, raw betas). + sXQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXK = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXV = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * V,)), 16) + sGin = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * K,)), 16) + sBetaRaw = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + sDtBias = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + sAlog = smem.allocate_tensor(cutlass.Float32, cute.make_layout((32,)), 16) + tmem_holder = smem.allocate_array(cutlass.Int32, 1) + mbar1 = smem.allocate_array(cutlass.Int64, 1) + mbar2 = smem.allocate_array(cutlass.Int64, 1) + mbar_ld = smem.allocate_array(cutlass.Int64, 1) + mbar2b = smem.allocate_array(cutlass.Int64, 1) + mbar_fb = smem.allocate_array(cutlass.Int64, 1) + sScratch = smem.allocate_tensor(cutlass.Float32, cute.make_layout((SCRATCH_FLOATS,)), 1024) + sB1 = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B1, smem_layout_b1.inner, dtype=cutlass.Float32), + smem_layout_b1.outer, + ) + sA2hi = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_A2HI, smem_layout_a2.inner, dtype=cutlass.Float32), + smem_layout_a2.outer, + ) + sA2lo = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_A2LO, smem_layout_a2.inner, dtype=cutlass.Float32), + smem_layout_a2.outer, + ) + sA3 = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_A3, smem_layout_a3.inner, dtype=cutlass.Float32), + smem_layout_a3.outer, + ) + sB3hi = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B3HI, smem_layout_b3.inner, dtype=cutlass.Float32), + smem_layout_b3.outer, + ) + sB3lo = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B3LO, smem_layout_b3.inner, dtype=cutlass.Float32), + smem_layout_b3.outer, + ) + sB1k = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B1K, smem_layout_b1k.inner, dtype=cutlass.Float32), + smem_layout_b1k.outer, + ) + sB1kf = cute.make_tensor(sScratch.iterator + OFF_B1K, cute.make_layout((T_PAD * 128,))) + # Flat (unswizzled-pointer) views for element writes with closed-form swizzle offsets. + sB1f = cute.make_tensor(sScratch.iterator + OFF_B1, cute.make_layout((48 * 128,))) + sA2hif = cute.make_tensor(sScratch.iterator + OFF_A2HI, cute.make_layout((128 * 16,))) + sA2lof = cute.make_tensor(sScratch.iterator + OFF_A2LO, cute.make_layout((128 * 16,))) + sB2hif = cute.make_tensor(sScratch.iterator + OFF_B2HI, cute.make_layout((128 * 16,))) + sB2lof = cute.make_tensor(sScratch.iterator + OFF_B2LO, cute.make_layout((128 * 16,))) + sB4hif = cute.make_tensor(sScratch.iterator + OFF_B4HI, cute.make_layout((T_PAD * T_PAD,))) + sB4lof = cute.make_tensor(sScratch.iterator + OFF_B4LO, cute.make_layout((T_PAD * T_PAD,))) + sA3f = cute.make_tensor(sScratch.iterator + OFF_A3, cute.make_layout((64 * 128,))) + sB3hif = cute.make_tensor(sScratch.iterator + OFF_B3HI, cute.make_layout((16 * 128,))) + sB3lof = cute.make_tensor(sScratch.iterator + OFF_B3LO, cute.make_layout((16 * 128,))) + # S0 / committed-state TMA boxes at the start of the scratch region (fp32 and bf16 views) + sBoxf = cute.make_tensor(sScratch.iterator, cute.make_layout((V * K,))) + sBoxh = cute.make_tensor( + cute.recast_ptr(sScratch.iterator, dtype=cutlass.BFloat16), cute.make_layout((V * K,)) + ) + # GEMM3 result [L; B hi; B lo; -] as a plain (64 x 16) row-major matrix. + sLB = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((4 * T_PAD, T_PAD), stride=(T_PAD, 1)), 16 + ) + sEc = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + # V'^T (v, s) and Kbar^T (k, s) staged with a padded row stride (conflict-free for both + # the row-per-thread producers and the 16-lanes-per-row tile writers). + sVpT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V * TP_STRIDE,)), 16) + sKbT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K * TP_STRIDE,)), 16) + # Cumulative gates G_t (per k) reuse the raw-gate staging buffer (no longer needed after the + # precompute). + sE = cute.make_tensor(sGin.iterator, cute.make_layout((t_max * K,))) + + # ---- load phase. Every load is issued before any of its consumers, so the CTA waits for its + # round trips in parallel: the slot-independent operands (new-token projections, conv + # weights, per-head constants, output-norm operands) under the slot load, the + # slot-dependent ones (S0 TMA, committed conv windows, every replay entry) as soon as the + # slot arrives; only the smem placement of the new tokens and accepted drafts waits for + # the accepted count. Each lane moves four consecutive channels (8 B for bf16, 16 B for + # fp32) per row, held as raw words until the row's smem position is known. + ch = tidx % 128 + quad = tidx // 128 + k4 = 4 * in_warp_tid + # (1) slot-independent: conv weights ([ch][w] rows of 16 B: threads [0, 384) one row each), + # dt_bias / output-norm weight / A_log (threads [384, 512)), new token j = warp + # (warps < NEW_TOK) + r_w = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + if quad == 0: + ld_row4(w_q, (hk_off + ch, 0), r_w, 0) + if quad == 1: + ld_row4(w_k, (hk_off + ch, 0), r_w, 0) + if quad == 2: + ld_row4(w_v, (hv_off + ch, 0), r_w, 0) + r_c = cute.make_rmem_tensor(cute.make_layout((3,)), cutlass.Float32) + if quad == 3: + r_c[0] = cutlass.Float32(dt_bias[hk_off + ch]) + r_c[1] = cutlass.Float32(onorm_weight[ch]) + r_c[2] = cutlass.Float32(A_log[i_h]) + r_new = cute.make_rmem_tensor(cute.make_layout((5 * 4,)), cutlass.Int32) + r_newb = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Float32) + tok_row = i_n * NEW_TOK + warp_idx + new_rows = (x_q, x_k, x_v, g, onorm_g, beta) + # FUSED_FB: warp REPLAY_WARP0 + mi computes the forget gate of W_fb rows [16 mi, 16 mi + 16) + # of this head; its f_a B fragments are loaded with the other new-token rows. + r_fb_b = cute.make_rmem_tensor(cute.make_layout((16,)), cutlass.Int32) + if cutlass.const_expr(not PDL): + if warp_idx < NEW_TOK: + _load_new_token_row( + *new_rows, + tok_row, + i_h, + i_hv, + hk_off, + hv_off, + k4, + r_new, + r_newb, + USE_FLAT_LAYOUT, + FUSED_FB, + ) + if cutlass.const_expr(FUSED_FB): + if warp_idx >= REPLAY_WARP0: + _load_fb_b_fragments(fa, i_n, NEW_TOK, in_warp_tid, r_fb_b) + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mbar1, 1) + cute.arch.mbarrier_init(mbar2, 1) + cute.arch.mbarrier_init(mbar2b, 1) + cute.arch.mbarrier_init(mbar_ld, 1) + if cutlass.const_expr(FUSED_FB): + cute.arch.mbarrier_init(mbar_fb, 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + if cutlass.const_expr(FUSED_FB): + # This head's W_fb rows into scratch (independent of the previous grid). + sFb = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_FB, fb_layout.inner, dtype=cutlass.BFloat16), + fb_layout.outer, + ) + gFb = cute.local_tile(mFb_tma, (TILE_K, FB_BOX_COLS), (None, None)) + tFbs, tFbg = cpasync.tma_partition( + tma_atom_fb, + 0, + cute.make_layout(1), + cute.group_modes(sFb, 0, 2), + cute.group_modes(gFb, 0, 2), + ) + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_fb, TILE_K * TILE_K * 2) + for b in cutlass.range_constexpr(TILE_K // FB_BOX_COLS): + cute.copy(tma_atom_fb, tFbg[(None, i_hv, b)], tFbs[(None, b)], tma_bar_ptr=mbar_fb) + # V2 state views may place many layer pages between adjacent slots. Widen + # the slot before indexing: slot * state stride can exceed INT32_MAX. + slot = Int64(slot32) + h0_idx = slot * HV + i_hv + # (2) slot-dependent: replay entry r = warp - REPLAY_WARP0 (every pending entry, accepted or + # not), the committed conv windows (TMA warp: three columns of q, k and v) + rep = warp_idx - REPLAY_WARP0 + r_rep = cute.make_rmem_tensor(cute.make_layout((6 * 4,)), cutlass.Int32) + r_repb = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Float32) + if (warp_idx >= REPLAY_WARP0) & (rep < NUM_SPEC): + ld_row4(k_cache, (slot, rep, hk_off + k4), r_rep, 0) + ld_row4(g_cache, (slot, rep, hk_off + k4), r_rep, 4) + ld_row4(v_cache, (slot, rep, hv_off + k4), r_rep, 8) + ld_row4(cs_q, (slot, hk_off + k4, KERNEL_WIDTH - 1 + rep), r_rep, 12) + ld_row4(cs_k, (slot, hk_off + k4, KERNEL_WIDTH - 1 + rep), r_rep, 16) + ld_row4(cs_v, (slot, hv_off + k4, KERNEL_WIDTH - 1 + rep), r_rep, 20) + r_repb[0] = cutlass.Float32(beta_cache[slot, rep, i_hv]) + r_win = cute.make_rmem_tensor(cute.make_layout((3 * (KERNEL_WIDTH - 1) * 4,)), cutlass.Int32) + if warp_idx == TMA_WARP: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + ld_row4(cs_q, (slot, hk_off + k4, w), r_win, 4 * w) + ld_row4(cs_k, (slot, hk_off + k4, w), r_win, 4 * (KERNEL_WIDTH - 1 + w)) + ld_row4(cs_v, (slot, hv_off + k4, w), r_win, 4 * (2 * (KERNEL_WIDTH - 1) + w)) + commit_len = num_accepted_tokens[slot] + # ---- TMA bulk load of S0 into the boxes; completes under the token staging / + # precompute and is consumed (mbar_ld) at the start of the chunk stage. + # Issued by the whole of one warp (a TMA load issued from a single thread never + # completes on this DSL build; the copy elects its issuing lane internally). + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ld, V * K * (STATE_DTYPE.width // 8)) + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sTl_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gS_h_all = cute.local_tile(mS_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gS_h = gS_h_all[(None, None, 0, h0_idx)] + else: + gS_h = gS_h_all[(None, None, 0, (i_hv, slot))] + tSl_h, tGl_h = cpasync.tma_partition( + tma_atom_l, + 0, + cute.make_layout(1), + cute.group_modes(sTl_h, 0, 2), + cute.group_modes(gS_h, 0, 2), + ) + if warp_idx == TMA_WARP: + cute.copy(tma_atom_l, tGl_h, tSl_h, tma_bar_ptr=mbar_ld) + if cutlass.const_expr(PDL): + # Everything above only reads this layer's pools, metadata and weights; the new-token + # rows are the previous grids' outputs. + griddepcontrol_wait() + # Release the dependent grid (the o_proj GEMM) right after the wait instead of at exit, + # so it can become resident on the SMs this grid leaves free and stream its weights + # during the recurrence. + # + # Why it is safe in the Kimi-K3 decode graph. The trigger follows this grid's own wait, + # so the projection GEMM, ``k3_kda_fab`` and everything before them have completed when + # a dependent starts. The grids that can start before this one completes are the o_proj + # GEMM (low-M or mid-M FP8 kernel), which before its wait only streams its weights and + # weight scales into shared memory (its weight warp triggers at once), and behind it the + # o_proj MNNVL all-reduce, which calls ``cudaGridDependencySynchronize`` before it reads + # or writes anything and triggers only after it. Neither writes global memory this grid + # reads (state pool, conv windows, replay rows, new-token rows) before its wait, and this + # grid's outputs (o, the committed state and conv windows) are read only after a wait: + # A_mid, which reads the prefix and snapshot rows before its own wait, starts only after + # the all-reduce's wait, i.e. after this grid completed. + griddepcontrol_launch_dependents() + if warp_idx < NEW_TOK: + _load_new_token_row( + *new_rows, + tok_row, + i_h, + i_hv, + hk_off, + hv_off, + k4, + r_new, + r_newb, + USE_FLAT_LAYOUT, + FUSED_FB, + ) + if cutlass.const_expr(FUSED_FB): + if warp_idx >= REPLAY_WARP0: + _load_fb_b_fragments(fa, i_n, NEW_TOK, in_warp_tid, r_fb_b) + # TMEM allocation overlaps the in-flight loads; the pointer is retrieved after the + # precompute barrier (first TMEM use is the S0 staging in the chunk stage). With PDL it + # follows griddepcontrol.wait, so a waiting CTA never holds tensor memory. + if warp_idx == 0: + cute.arch.alloc_tmem(TMEM_COLS, tmem_holder) + # (3) placement: the slot-independent operands first, the token rows once the accepted count + # is in. Smem rows keep the [row][channel] fp32 layout of the precompute. + if quad < 3: + wv = unpack4(w_q.element_type, r_w, 0) + for w in cutlass.range_constexpr(KERNEL_WIDTH): + sConvW[quad * (KERNEL_WIDTH * K) + w * K + ch] = wv[w] + if quad == 3: + sDtBias[ch] = r_c[0] + sOnW[ch] = r_c[1] + if ch == 0: + sAlog[0] = r_c[2] + if warp_idx < NEW_TOK: + st_shared_f32x4(saddr(sOnG, warp_idx * V + k4), *unpack4(onorm_g.element_type, r_new, 16)) + if warp_idx == TMA_WARP: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_shared_f32x4(saddr(sXQ, w * K + k4), *unpack4(CONV_DTYPE, r_win, 4 * w)) + st_shared_f32x4( + saddr(sXK, w * K + k4), *unpack4(CONV_DTYPE, r_win, 4 * (KERNEL_WIDTH - 1 + w)) + ) + st_shared_f32x4( + saddr(sXV, w * V + k4), + *unpack4(CONV_DTYPE, r_win, 4 * (2 * (KERNEL_WIDTH - 1) + w)), + ) + # Only NUM_SPEC drafts can be pending from the previous round. Clamp + # so a malformed count cannot drive T_loop past the t_max-sized SMEM + # buffers or the num_spec extents of the replay caches. + if commit_len > NUM_SPEC: + commit_len = cutlass.Int32(NUM_SPEC) + T_loop = commit_len + NEW_TOK + if warp_idx < NEW_TOK: + t_new = commit_len + warp_idx + st_shared_f32x4(saddr(sXQ, (t_new + 3) * K + k4), *unpack4(x_q.element_type, r_new, 0)) + st_shared_f32x4(saddr(sXK, (t_new + 3) * K + k4), *unpack4(x_k.element_type, r_new, 4)) + st_shared_f32x4(saddr(sXV, (t_new + 3) * V + k4), *unpack4(x_v.element_type, r_new, 8)) + if cutlass.const_expr(not FUSED_FB): + st_shared_f32x4(saddr(sGin, t_new * K + k4), *unpack4(g.element_type, r_new, 12)) + if in_warp_tid == 0: + sBetaRaw[t_new] = r_newb[0] + if (warp_idx >= REPLAY_WARP0) & (rep < commit_len): + # replayed drafts produce no output: their q rows are unused + z = cutlass.Float32(0.0) + st_shared_f32x4(saddr(sQ, rep * K + k4), z, z, z, z) + st_shared_f32x4(saddr(sK, rep * K + k4), *unpack4(K_CACHE_DTYPE, r_rep, 0)) + st_shared_f32x4(saddr(sG, rep * K + k4), *unpack4(g_cache.element_type, r_rep, 4)) + st_shared_f32x4(saddr(sVall, rep * V + k4), *unpack4(V_CACHE_DTYPE, r_rep, 8)) + st_shared_f32x4(saddr(sXQ, (rep + 3) * K + k4), *unpack4(CONV_DTYPE, r_rep, 12)) + st_shared_f32x4(saddr(sXK, (rep + 3) * K + k4), *unpack4(CONV_DTYPE, r_rep, 16)) + st_shared_f32x4(saddr(sXV, (rep + 3) * V + k4), *unpack4(CONV_DTYPE, r_rep, 20)) + if in_warp_tid == 0: + sBeta[rep] = r_repb[0] + if cutlass.const_expr(FUSED_FB): + # ---- forget gate of the new tokens, gate[t, row] = bf16(f_a[t] @ W_fb[h * 128 + row]^T) + # (MMA M = this warp's 16 W_fb rows, N = the token columns, K = 128 in eight + # sequential k16 steps from zero) into the raw-gate rows read by the precompute. + if warp_idx >= REPLAY_WARP0: + col0 = (i_n * NEW_TOK) % 8 + if (in_warp_tid // 4 - col0 + 8) % 8 >= NEW_TOK: + for i in cutlass.range_constexpr(16): + r_fb_b[i] = cutlass.Int32(0) + # A fragments: W_fb rows 16 mi + {0..7, 8..15} x k16 chunk kk through ldmatrix.x4 on + # the 128B-swizzled boxes (row r & 7 == lane & 7 selects the 16 B unit XOR) + mi = warp_idx - REPLAY_WARP0 + mat = in_warp_tid >> 3 + r8 = in_warp_tid & 7 + a_row = saddr(sScratch, OFF_FB) + (mi * 16 + (mat & 1) * 8 + r8) * 128 + cute.arch.mbarrier_wait(mbar_fb, 0) + acc = [cutlass.Float32(0.0)] * 4 + for kk in cutlass.range_constexpr(TILE_K // 16): + a = ldmatrix_x4( + a_row + (kk // 4) * (TILE_K * 128) + ((((kk % 4) * 2 + (mat >> 1)) ^ r8) << 4) + ) + acc = mma_bf16_16816(a, r_fb_b[2 * kk], r_fb_b[2 * kk + 1], acc) + for e in cutlass.range_constexpr(4): + jt = (2 * (in_warp_tid & 3) + (e & 1) - col0 + 8) % 8 + if jt < NEW_TOK: + row = mi * 16 + (in_warp_tid >> 2) + (e >> 1) * 8 + sGin[(commit_len + jt) * K + row] = round_bf16(acc[e]) + # Padding rows [T_loop, t_max) of the gates, values and betas are zeroed so the chunk + # stage's r_Y / cumulative-gate / V' loops are branch-free. + if tidx < K: + for _z in cutlass.range(T_loop, t_max, 1): + sG[_z, tidx] = cutlass.Float32(0.0) + sVall[_z * V + tidx] = cutlass.Float32(0.0) + if tidx == 0: + sBeta[_z] = cutlass.Float32(0.0) + wg = warp_idx // 4 + wg_tidx = tidx % 128 + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + r_state = cute.make_rmem_tensor(cute.make_layout((32,)), cutlass.Float32) + cute.arch.barrier() + + # ---- per-token-parallel precompute (smem only; the new tokens' cache and conv-window writes + # are issued from smem while the tensor cores run GEMM1/GEMM3): warp j < NEW_TOK computes + # the q and k conv + SiLU + L2 norm of new token j (the two chains interleaved), warp 8 + j + # its gate, beta and v conv. + half = NUM_THREADS // 64 + role = warp_idx // half + j = warp_idx % half + if j < NEW_TOK: + t = commit_len + j + if role == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_conv_0 = cutlass.Float32(0.0) + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + r_conv_0 += sXQ[(t + w) * K + k_idx] * sConvW[w * K + k_idx] + r_conv = sXK[t * K + k_idx] * sConvW[k_weight_base + 0 * K + k_idx] + for w in cutlass.range(1, KERNEL_WIDTH, unroll_full=True): + r_conv += sXK[(t + w) * K + k_idx] * sConvW[k_weight_base + w * K + k_idx] + e0 = cute.math.exp(-r_conv_0, fastmath=True) + r_q[i] = r_conv_0 * cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) + r_conv = r_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) + ) + r_k[i] = r_conv + sum_q = 0.0 + sum_k = 0.0 + for i in cutlass.range(vec_size, unroll_full=True): + sum_q += r_q[i] * r_q[i] + sum_k += r_k[i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=-1, mask_and_clamp=31 + ) + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_q_scaled = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + sQ[t, k_idx] = r_q[i] * rnorm_q_scaled + sK[t, k_idx] = r_k[i] * rnorm_k + else: + r_exp_A = cute.math.exp(sAlog[0], fastmath=True) + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_g_raw = sGin[t * K + k_idx] + r_g_raw = r_g_raw + sDtBias[k_idx] + exp_A_x = r_exp_A * r_g_raw + sigmoid_val = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) + ) + sG[t, k_idx] = lower_bound * sigmoid_val + r_beta_t = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-sBetaRaw[t], fastmath=True) + ) + if in_warp_tid == 0: + sBeta[t] = r_beta_t + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid + _v_conv = 0.0 + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + _v_conv += sXV[(t + w) * V + v_idx] * sConvW[v_weight_base + w * V + v_idx] + _v_conv = _v_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) + ) + sVall[t * V + v_idx] = _v_conv + cute.arch.barrier() + + # ================================================================= + # Chunk (WY-form) recurrence on tensor cores. + # ================================================================= + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + pool = cutlass.utils.TmemBufferPool(tmem_ptr, TMEM_COLS) + tD4_frag = tiled_mma4.make_fragment_C(tiled_mma4.partition_shape_C((V, T_PAD))) + tD4 = pool.allocate_tensor(tD4_frag.layout, cutlass.Float32) # outputs (columns 0..15) + tD2h_frag = tiled_mma2h.make_fragment_C(tiled_mma2h.partition_shape_C((V, K // 2))) + tD2h0 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) # GEMM2 half 0: k cols 0..63 + tD2h1 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) # GEMM2 half 1: k cols 64..127 + tD2q_frag = tiled_mma2q.make_fragment_C(tiled_mma2q.partition_shape_C((V, T_PAD + K // 2))) + tD2q = cute.make_tensor(tD4.iterator, tD2q_frag.layout) # [D4 | dS half 0] accumulator view + tA_frag = tiled_mma_f32a.make_fragment_A(tiled_mma_f32a.partition_shape_A((V, K))) + tAhi = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tAlo = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tD1_frag = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C((V, N1))) + tD1 = pool.allocate_tensor(tD1_frag.layout, cutlass.Float32) + tD3_frag = tiled_mma3.make_fragment_C(tiled_mma3.partition_shape_C((4 * T_PAD, T_PAD))) + tD3 = pool.allocate_tensor(tD3_frag.layout, cutlass.Float32) + tD2 = cute.make_tensor(tD2h0.iterator, tA_frag.layout) # A-view layout (all 128 columns) + + st32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + ld16 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + ld32 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + thr_stA = tcgen05.make_tmem_copy(st32, tD2).get_slice(wg_tidx) + tD2_st = thr_stA.partition_D(tD2) + rA_shape = thr_stA.partition_S(tD2).shape + tAhi_st = thr_stA.partition_D(tAhi) + tAlo_st = thr_stA.partition_D(tAlo) + thr_ldD2 = tcgen05.make_tmem_copy(ld32, tD2).get_slice(wg_tidx) + tD2_ld = thr_ldD2.partition_S(tD2) + rS = cute.make_rmem_tensor(thr_ldD2.partition_D(tD2).shape[:-1], cutlass.Float32) + rS_flat = cute.make_tensor(rS.iterator, cute.make_layout((32,))) + # TMA store partitions for the NUM_SLABS boxes of this CTA's committed-state tile. + tS_list = [] + tG_list = [] + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sT_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + sl_h * ((V // SPLIT_V) * 32), + box_layout_t.inner, + dtype=STATE_DTYPE, + ), + box_layout_t.outer, + ) + gT_h_all = cute.local_tile(mT_tma, (V // SPLIT_V, BOX_COLS, 1), (vh, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gT_h = gT_h_all[(None, None, 0, h0_idx)] + else: + gT_h = gT_h_all[(None, None, 0, (i_hv, slot))] + tS_h, tG_h = cpasync.tma_partition( + tma_atom_t, + 0, + cute.make_layout(1), + cute.group_modes(sT_h, 0, 2), + cute.group_modes(gT_h, 0, 2), + ) + tS_list.append(tS_h) + tG_list.append(tG_h) + r_hi = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_lo = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_hi_flat = cute.make_tensor(r_hi.iterator, cute.make_layout((32,))) + r_lo_flat = cute.make_tensor(r_lo.iterator, cute.make_layout((32,))) + + # ---- cumulative gates per k column (thread kq, quarter q owns tokens t = q mod 4): + # G_t -> sE, the T x T coupling factors, and the operand rows written with closed-form + # 128B-swizzle offsets (lanes = consecutive k: conflict-free). The GEMM3 rows (A3 = + # [Kf; Qf hi; Qf lo], B3 = Kb hi / lo) lie clear of the S0 boxes and are written first, + # so GEMM3 runs while S0 moves into TMEM; the GEMM1 rows (K~ hi / lo, Q~) stay in + # registers until the S0 boxes have been read. + kq = tidx % K + quarter = tidx // K + t_mid = T_loop // 2 + r_G = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + g_run = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t < t_max: + g_run = g_run + sG[t, kq] + r_G[t] = g_run + g_mid = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t == t_mid: + g_mid = r_G[t] + # Factors of the T x T coupling: e^{G_t - G_mid} (t side) and e^{G_mid - G_s} (s side), + # clamped at e^87 so a chunk whose decay exceeds the fp32 range yields ~0 instead of + # inf/NaN (the affected products are < e^{-|g|} and negligible either way). + exp_clamp = cutlass.Float32(87.0) + ka = kq // 32 + kc = kq % 32 + b1_base = ka * (3 * T_PAD * 32) + a3_base = ka * (4 * T_PAD * 32) + b3_base = ka * (T_PAD * 32) + b1k_base = ka * (T_PAD * 32) + # GEMM1 rows of this thread's tokens t = quarter + 4 i: (K~ hi, Q~, K~ lo) per token + r_b1 = cute.make_rmem_tensor(cute.make_layout((3 * (T_PAD // 4),)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + kt_hi = cutlass.Float32(0.0) + kt_lo = cutlass.Float32(0.0) + qt_r = cutlass.Float32(0.0) + kf = cutlass.Float32(0.0) + qf_hi = cutlass.Float32(0.0) + qf_lo = cutlass.Float32(0.0) + kb_hi = cutlass.Float32(0.0) + kb_lo = cutlass.Float32(0.0) + if t < T_loop: + d_t = r_G[t] - g_mid + e_t = cute.math.exp(-cute.arch.fmax(-d_t, -exp_clamp), fastmath=True) + e_bt = cute.math.exp(-cute.arch.fmax(d_t, -exp_clamp), fastmath=True) + eg_t = cute.math.exp(r_G[t], fastmath=True) + sE[t * K + kq] = r_G[t] + if t == commit_len: + sEc[kq] = eg_t + k_t = sK[t, kq] + q_t = sQ[t, kq] + kt = k_t * eg_t + kt_hi = round_tf32(kt) + qt_r = round_tf32(q_t * eg_t) + kf = round_tf32(k_t * e_t) + qf = q_t * e_t + qf_hi = round_tf32(qf) + kb = k_t * e_bt + kb_hi = round_tf32(kb) + kt_lo = kt - kt_hi + qf_lo = qf - qf_hi + kb_lo = kb - kb_hi + # rows t, 16+t, 32+t share t % 8 -> the same 16-byte-chunk XOR + xc = kc ^ (4 * (t % 8)) + sA3f[a3_base + 32 * t + xc] = kf + sA3f[a3_base + 32 * (T_PAD + t) + xc] = qf_hi + sA3f[a3_base + 32 * (2 * T_PAD + t) + xc] = qf_lo + sB3hif[b3_base + 32 * t + xc] = kb_hi + sB3lof[b3_base + 32 * t + xc] = kb_lo + r_b1[3 * (t // 4)] = kt_hi + r_b1[3 * (t // 4) + 1] = qt_r + r_b1[3 * (t // 4) + 2] = kt_lo + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + # ---- GEMM3: D3 = A3 Kb_hi^T + A3 Kb_lo^T (warp 0; completion is tracked with GEMM1's) + tCrA3 = tiled_mma3.make_fragment_A( + cute.make_tensor( + cute.recast_ptr(sA3.iterator, smem_layout_a3.inner, dtype=cutlass.TFloat32), + sA3.layout, + ) + ) + tCrB3hi = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3hi.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3hi.layout, + ) + ) + tCrB3lo = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3lo.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3lo.layout, + ) + ) + if warp_idx == 0: + tiled_mma3.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3hi[(None, None, kb_i, 0)], tD3 + ) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3lo[(None, None, kb_i, 0)], tD3 + ) + + # ---- S0 row v = wg_tidx, columns [32 wg, +32), from the TMA-loaded swizzled box + # (16 B chunks: conflict-free LDS.128 under the 128B swizzle) into TMEM (hi / lo) + cute.arch.mbarrier_wait(mbar_ld, 0) + # in-place state update: the peer CTA must not store its half before our S0 load + # completed -> arrive here, wait right before the state store. + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_arrive_relaxed() + vx = 4 * (wg_tidx % 8) + if cutlass.const_expr(STATE_BF16): + # bf16 box: 64 columns per 128 B row, so this warpgroup's 32 columns are the 16 B + # chunks [4 (wg % 2), +4) of box wg // 2. A bf16 value is exact in tf32: no lo part, + # and the S0_lo sweep of GEMM1 is skipped. + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (wg % 2) + c) ^ (wg_tidx % 8) + off_l = (wg // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + for i in cutlass.range(8, unroll_full=True): + r_state[8 * c + i] = cutlass.Float32(sBoxh[off_l + i]) + for i in cutlass.range(32, unroll_full=True): + r_hi_flat[i] = r_state[i] + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, jj)]) + else: + for c in cutlass.range(8, unroll_full=True): + off_l = wg * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range(4, unroll_full=True): + r_state[4 * c + i] = sBoxf[off_l + i] + for i in cutlass.range(32, unroll_full=True): + s_hi = round_tf32(r_state[i]) + r_hi_flat[i] = s_hi + r_lo_flat[i] = r_state[i] - s_hi + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, jj)]) + cute.copy(st32, r_lo, tAlo_st[(None, None, None, jj)]) + cute.arch.fence_view_async_tmem_store() + # the S0 boxes have been read: the GEMM1 operand rows go over them + cute.arch.barrier() + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + xc = kc ^ (4 * (t % 8)) + sB1f[b1_base + 32 * t + xc] = r_b1[3 * (t // 4)] + if cutlass.const_expr(S0_LO): + sB1kf[b1k_base + 32 * t + xc] = r_b1[3 * (t // 4)] + sB1f[b1_base + 32 * (T_PAD + t) + xc] = r_b1[3 * (t // 4) + 1] + sB1f[b1_base + 32 * (2 * T_PAD + t) + xc] = r_b1[3 * (t // 4) + 2] + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + + # ---- operand fragments of GEMM1 (D1 = S0_hi B1^T + S0_lo K~hi^T) + tCrB1 = tiled_mma1.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1.iterator, smem_layout_b1.inner, dtype=cutlass.TFloat32), + sB1.layout, + ) + ) + tCrB1k = tiled_mma1k.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1k.iterator, smem_layout_b1k.inner, dtype=cutlass.TFloat32), + sB1k.layout, + ) + ) + tD1k = cute.make_tensor( + tD1.iterator, + tiled_mma1k.make_fragment_C(tiled_mma1k.partition_shape_C((V, T_PAD))).layout, + ) + tAhi_tf = cute.make_tensor(cute.recast_ptr(tAhi.iterator, dtype=cutlass.TFloat32), tAhi.layout) + tAlo_tf = cute.make_tensor(cute.recast_ptr(tAlo.iterator, dtype=cutlass.TFloat32), tAlo.layout) + # ---- warp 0 issues GEMM1 right away; warps 1..15 stage Kbar^T (sKbT[k, s], lanes = + # consecutive k: conflict-free), seed the GEMM2 accumulator, build the B2 tiles and write + # the new tokens' cache rows while the tensor cores run. + if warp_idx > 0: + for i in cutlass.range( + (K * T_PAD + NUM_THREADS - 33) // (NUM_THREADS - 32), unroll_full=True + ): + e_kb = i * (NUM_THREADS - 32) + (tidx - 32) + if e_kb < K * T_PAD: + k_kb = e_kb % K + s_idx = e_kb // K + b_val = cutlass.Float32(0.0) + if s_idx <= commit_len: + b_val = sK[s_idx, k_kb] * cute.math.exp( + sE[commit_len * K + k_kb] - sE[s_idx * K + k_kb], fastmath=True + ) + sKbT[k_kb * TP_STRIDE + s_idx] = b_val + if warp_idx == 0: + tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1, tD1, tAhi_tf[(None, None, kb_i)], tCrB1[(None, None, kb_i, 0)], tD1 + ) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(S0_LO): + tiled_mma1k.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1k, + tD1k, + tAlo_tf[(None, None, kb_i)], + tCrB1k[(None, None, kb_i, 0)], + tD1k, + ) + # tracks GEMM3 and GEMM1 + with cute.arch.elect_one(): + tcgen05.commit(mbar1) + else: + # sKbT complete before the B2 tiles read it + cute.arch.barrier(barrier_id=5, number_of_threads=NUM_THREADS - 32) + # ---- GEMM2 accumulator seed (overlaps the GEMM1/GEMM3 execution) + # D4 (output-combination) columns start from zero: GEMM2 half 0 accumulates into [D4 | dS] + st16 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_stD4 = tcgen05.make_tmem_copy(st16, tD4).get_slice(wg_tidx) + tD4_st = thr_stD4.partition_D(tD4) + r_zero = cute.make_rmem_tensor(thr_stD4.partition_S(tD4).shape, cutlass.Float32) + for i in cutlass.range(cute.size(r_zero), unroll_full=True): + r_zero[i] = cutlass.Float32(0.0) + if wg == 0: + cute.copy(st16, r_zero, tD4_st) + # only this CTA's state rows (TMEM lanes [vh * V/2, +V/2)) are seeded and stored + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + for i in cutlass.range(32, unroll_full=True): + r_hi_flat[i] = r_state[i] * sEc[wg * 32 + i] + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.copy(st32, r_hi, tD2_st[(None, None, None, jj)]) + cute.arch.fence_view_async_tmem_store() + # ---- B2 tiles (k rows x s cols, K-major SW64) from sKbT while the MMAs run: + # lanes 0..15 -> row 2w, lanes 16..31 -> row 2w+1, four 32-row rounds + s_col = in_warp_tid % T_PAD + if warp_idx > 0: + for r in cutlass.range(4, unroll_full=True): + k_row = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + b_val2 = sKbT[k_row * TP_STRIDE + s_col] + b_hi = round_tf32(b_val2) + o_b2 = swz64_k16(k_row, s_col) + sB2hif[o_b2] = b_hi + sB2lof[o_b2] = b_val2 - b_hi + # warp 0 is issuing the MMAs: warps 1..4 cover its rows 32 r + {0, 1} + if warp_idx <= 4: + k_row0 = 32 * (warp_idx - 1) + in_warp_tid // T_PAD + b_val0 = sKbT[k_row0 * TP_STRIDE + s_col] + b_hi0 = round_tf32(b_val0) + o_b20 = swz64_k16(k_row0, s_col) + sB2hif[o_b20] = b_hi0 + sB2lof[o_b20] = b_val0 - b_hi0 + # ---- the new tokens' replay-cache rows and conv-window commits, from smem: warp 1 + p + # writes draft p's k / gate / v cache rows, its raw conv columns and beta, warps + # REPLAY_WARP0 + {0, 1, 2} commit the q / k / v conv windows (the raw inputs after the + # golden token). One CTA of a pair writes them. + if vh == 0: + if (warp_idx >= 1) & (warp_idx <= NUM_SPEC): + p = warp_idx - 1 + t_c = commit_len + 1 + p + st_row4(k_cache, (slot, p, hk_off + k4), ld_shared_f32x4(saddr(sK, t_c * K + k4))) + st_row4(g_cache, (slot, p, hk_off + k4), ld_shared_f32x4(saddr(sG, t_c * K + k4))) + st_row4(v_cache, (slot, p, hv_off + k4), ld_shared_f32x4(saddr(sVall, t_c * V + k4))) + st_row4( + cs_q, + (slot, hk_off + k4, KERNEL_WIDTH - 1 + p), + ld_shared_f32x4(saddr(sXQ, (t_c + 3) * K + k4)), + ) + st_row4( + cs_k, + (slot, hk_off + k4, KERNEL_WIDTH - 1 + p), + ld_shared_f32x4(saddr(sXK, (t_c + 3) * K + k4)), + ) + st_row4( + cs_v, + (slot, hv_off + k4, KERNEL_WIDTH - 1 + p), + ld_shared_f32x4(saddr(sXV, (t_c + 3) * V + k4)), + ) + if in_warp_tid == 0: + beta_cache[slot, p, i_hv] = sBeta[t_c].to(BETA_CACHE_DTYPE) + if warp_idx == REPLAY_WARP0: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_row4( + cs_q, + (slot, hk_off + k4, w), + ld_shared_f32x4(saddr(sXQ, (commit_len + 1 + w) * K + k4)), + ) + if warp_idx == REPLAY_WARP0 + 1: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_row4( + cs_k, + (slot, hk_off + k4, w), + ld_shared_f32x4(saddr(sXK, (commit_len + 1 + w) * K + k4)), + ) + if warp_idx == REPLAY_WARP0 + 2: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_row4( + cs_v, + (slot, hv_off + k4, w), + ld_shared_f32x4(saddr(sXV, (commit_len + 1 + w) * V + k4)), + ) + # ---- wait for GEMM1/GEMM3 + cute.arch.mbarrier_wait(mbar1, 0) + tcgen05_fence_after_thread_sync() + thr_ld1 = tcgen05.make_tmem_copy(ld16, tD1).get_slice(wg_tidx) + rD1 = cute.make_rmem_tensor(thr_ld1.partition_D(tD1).shape, cutlass.Float32) + cute.copy(ld16, thr_ld1.partition_S(tD1), rD1) + cute.arch.fence_view_async_tmem_load() + rD1_flat = cute.make_tensor(rD1.iterator, cute.make_layout((N1,))) + v_row = wg_tidx + # warpgroup 0: D3 (64 x 16) -> smem [L; B_hi; B_lo; -] + ld3 = cute.make_copy_atom( + tcgen05.Ld16x128bOp(tcgen05.Repetition(T_PAD * 32 // 128)), cutlass.Float32 + ) + thr_ld3 = tcgen05.make_tmem_copy(ld3, tD3).get_slice(wg_tidx) + tD3_ld = thr_ld3.partition_S(tD3) + sLB3 = cute.make_tensor( + sLB.iterator, + cute.make_layout(((4 * T_PAD, T_PAD), 1, 1), stride=((T_PAD, 1), 0, 0)), + ) + tD3_sm = thr_ld3.partition_D(sLB3) + rD3 = cute.make_rmem_tensor(tD3_sm.shape, cutlass.Float32) + if wg == 0: + cute.copy(ld3, tD3_ld, rD3) + cute.arch.fence_view_async_tmem_load() + cute.autovec_copy(rD3, tD3_sm) + cute.arch.barrier() + # ---- warpgroup 0: V' by forward substitution (row v per thread, four accumulators); + # warpgroups 1-2: output-combination operand Bcoef[t, s] = B_hi + B_lo (rows t < T_loop, + # s <= t; zero elsewhere), one element per thread, K-major SW64 16 x 16 tiles (hi / lo) + # that head the GEMM2 half-0 B tile + if wg == 0: + r_Y = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for s in cutlass.range(T_PAD, unroll_full=True): + y = cutlass.Float32(0.0) + if s < t_max: + u0 = rD1_flat[s] + rD1_flat[2 * T_PAD + s] + y = sBeta[s] * (sVall[s * V + v_row] - u0) + r_Y[s] = y + r_V = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + vp = cutlass.Float32(0.0) + if t < t_max: + acc0 = cutlass.Float32(0.0) + acc1 = cutlass.Float32(0.0) + acc2 = cutlass.Float32(0.0) + acc3 = cutlass.Float32(0.0) + for s in range(t): + if s % 4 == 0: + acc0 = acc0 + sLB[t, s] * r_V[s] + elif s % 4 == 1: + acc1 = acc1 + sLB[t, s] * r_V[s] + elif s % 4 == 2: + acc2 = acc2 + sLB[t, s] * r_V[s] + else: + acc3 = acc3 + sLB[t, s] * r_V[s] + vp = r_Y[t] - sBeta[t] * ((acc0 + acc1) + (acc2 + acc3)) + r_V[t] = vp + for t in cutlass.range(T_PAD, unroll_full=True): + sVpT[v_row * TP_STRIDE + t] = r_V[t] + else: + if tidx < 128 + T_PAD * T_PAD: + t4 = (tidx - 128) // T_PAD + s4 = (tidx - 128) % T_PAD + bc = cutlass.Float32(0.0) + if (t4 < T_loop) & (s4 <= t4): + bc = sLB[T_PAD + t4, s4] + sLB[2 * T_PAD + t4, s4] + bc_hi = round_tf32(bc) + o_b4 = swz64_k16(t4, s4) + sB4hif[o_b4] = bc_hi + sB4lof[o_b4] = bc - bc_hi + cute.arch.barrier() + # ---- A2 tiles (v rows x s cols) from sVpT, same 16-lanes-per-row mapping; s > c -> 0 + for r in cutlass.range(4, unroll_full=True): + v_row2 = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + a_val = sVpT[v_row2 * TP_STRIDE + s_col] + a_hi = round_tf32(a_val) + o_a2 = swz64_k16(v_row2, s_col) + sA2hif[o_a2] = a_hi + sA2lof[o_a2] = a_val - a_hi + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + sA2hi_tf = cute.make_tensor( + cute.recast_ptr(sA2hi.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2hi.layout, + ) + sA2lo_tf = cute.make_tensor( + cute.recast_ptr(sA2lo.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2lo.layout, + ) + tCrA2hi = tiled_mma2h.make_fragment_A(sA2hi_tf) + tCrA2lo = tiled_mma2h.make_fragment_A(sA2lo_tf) + # B2 halves: k rows 0..63 / 64..127 of the SW64 K-major tiles (a 4 KB shift keeps the swizzle) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B2HI + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrB2lo_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B2LO + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2hi_tf) + tCrA2lo_q = tiled_mma2q.make_fragment_A(sA2lo_tf) + tCrB2q_hi = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B4HI, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + tCrB2q_lo = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B4LO, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + if warp_idx == 0: + # GEMM2 half 0 fused with the output combination: [D4 | dS_0..63] += A2 x [Bcoef; Kbar_0..63] + tiled_mma2q.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_lo[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2lo_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2) + tiled_mma2h.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2lo_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2lo[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2b) + + # ---- outputs of the new tokens: o_t = O0_t + D4[v, t] (columns 0..15 of GEMM2 half 0), + # warpgroup w writes new tokens j = w mod 4, bf16-rounded, into sVall rows j + cute.arch.mbarrier_wait(mbar2, 0) + tcgen05_fence_after_thread_sync() + thr_ld4 = tcgen05.make_tmem_copy(ld16, tD4).get_slice(wg_tidx) + rD4 = cute.make_rmem_tensor(thr_ld4.partition_D(tD4).shape, cutlass.Float32) + cute.copy(ld16, thr_ld4.partition_S(tD4), rD4) + cute.arch.fence_view_async_tmem_load() + rD4_flat = cute.make_tensor(rD4.iterator, cute.make_layout((T_PAD,))) + for j_new in cutlass.range(NEW_TOK, unroll_full=True): + if wg == j_new % 4: + t_out = commit_len + j_new + o_val = cutlass.Float32(0.0) + for tt in cutlass.range(T_PAD, unroll_full=True): + if tt == t_out: + o_val = rD1_flat[T_PAD + tt] + rD4_flat[tt] + sVall[j_new * V + v_row] = cutlass.Float32(cutlass.BFloat16(o_val)) + # ---- committed-state store through TMA: D2 slab -> swizzled smem box (row per thread, + # 16 B chunks land in distinct banks under the 128B swizzle) -> one bulk tensor store + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_wait() + # All 16 warps produced disjoint value rows above; the barrier makes them visible to the + # norm epilogue, which assigns one complete token to each warp. + cute.arch.barrier() + # The state-owning warps of warpgroups 0/1 stage k columns 0..63 (GEMM2 half 0, mbar2), those + # of warpgroups 2/3 columns 64..127 (half 1, mbar2b) into the swizzled smem boxes; the first + # staging thread of each pair issues the TMA stores once its half is visible (named barriers + # 6 / 7) and waits for the bulk groups at the very end. The warps without state rows run the + # norm epilogue. + if wg < 2: + cute.arch.mbarrier_wait(mbar2, 0) + else: + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + for sl in cutlass.range_constexpr(4): + if wg == sl: + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + cute.copy(ld32, tD2_ld[(None, None, None, sl)], rS) + cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr(STATE_BF16): + # bf16 box: chunks [4 (sl % 2), +4) of box sl // 2, round to nearest + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (sl % 2) + c) ^ (wg_tidx % 8) + off_b = ( + (sl // 2) * ((V // SPLIT_V) * 64) + + 64 * (wg_tidx - (V // SPLIT_V) * vh) + + 8 * chunk + ) + for i in cutlass.range(8, unroll_full=True): + sBoxh[off_b + i] = cutlass.BFloat16(rS_flat[8 * c + i]) + else: + for c in cutlass.range(8, unroll_full=True): + off_b = ( + sl * ((V // SPLIT_V) * 32) + + 32 * (wg_tidx - (V // SPLIT_V) * vh) + + ((4 * c) ^ vx) + ) + for i in cutlass.range(4, unroll_full=True): + sBoxf[off_b + i] = rS_flat[4 * c + i] + # last TMEM reads of the kernel: warps 1..15 signal barrier 8, warp 0 frees TMEM below + tcgen05_fence_before_thread_sync() + if warp_idx > 0: + cute.arch.barrier_arrive(barrier_id=8, number_of_threads=NUM_THREADS) + cute.arch.fence_view_async_shared() + st_tid0 = (V // 2) * vh + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + if wg < 2: + cute.arch.barrier(barrier_id=6, number_of_threads=256 // SPLIT_V) + if tidx == st_tid0: + for s in cutlass.range_constexpr(NUM_SLABS // 2): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + else: + cute.arch.barrier(barrier_id=7, number_of_threads=256 // SPLIT_V) + if tidx == 256 + st_tid0: + for s in cutlass.range_constexpr(NUM_SLABS // 2, NUM_SLABS): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + if warp_idx == 0: + # all TMEM reads are done (barrier 8): free the columns while the store drains + cute.arch.barrier(barrier_id=8, number_of_threads=NUM_THREADS) + tcgen05_fence_after_thread_sync() + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS) + # ---- gated RMSNorm epilogue (CTA 0 of a pair) on the warps (warp % 4) in {2, 3}: one token + # per warp, p_v = o_v * w_v * sigmoid(g_v), out = bf16(p_v * rsqrt(mean(o_v^2) + eps)) + token_idx = (warp_idx // 4) * 2 + (warp_idx % 4 - 2) + if (vh == 0) & ((warp_idx % 4) // 2 == 1): + while token_idx < NEW_TOK: + pre_values = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + rms_partial = cutlass.Float32(0.0) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + value = sVall[token_idx * V + v_idx] + gate = sOnG[token_idx * V + v_idx] + sigmoid_gate = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) + ) + pre_values[i] = value * sOnW[v_idx] * sigmoid_gate + rms_partial += value * value + for offset in [16, 8, 4, 2, 1]: + rms_partial += cute.arch.shuffle_sync_bfly( + rms_partial, offset=offset, mask=-1, mask_and_clamp=31 + ) + inv_rms = cute.math.rsqrt( + rms_partial / cutlass.Float32(V) + cutlass.Float32(onorm_eps), + fastmath=True, + ) + output_row = i_n * NEW_TOK + token_idx + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[0, output_row, i_hv, v_idx] = cutlass.BFloat16(pre_values[i] * inv_rms) + token_idx = token_idx + NUM_THREADS // 32 + # the state stores must have read their smem boxes before the CTA exits + if (tidx == st_tid0) | (tidx == 256 + st_tid0): + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/mtp_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/mtp_ops.py new file mode 100644 index 000000000000..220bf1d3d022 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/kda/mtp_ops.py @@ -0,0 +1,1881 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL host launcher of the Kimi K3 KDA multi-token speculative verify. + +``kda_mtp_decode_impl`` launches the source-integrated ``kda_decode_mtp_kernel`` (see +``kernels/kda/kda_mtp_decode.py``) and its v2 / two-CTA forms; the model calls it through +the ``ops/ssm/kda_mtp_decode`` entry. One launch fuses, per generation +request: replay of previously-accepted draft tokens from the ``kg/v/beta`` +caches, causal conv + SiLU, Q/K L2 norm, beta sigmoid, lower-bound gate, and +the KDA delta-rule recurrence over the ``1 + num_spec`` new tokens. The +recurrent state and base conv windows are committed **in place** after the +first new (golden) token; the spec tokens are cached for the next round. + +State-management contract (differs from the legacy intermediate-buffer + +``update_mamba_states`` promotion flow): after this op returns, the pools +hold the state as of the last golden token, with the new spec tokens pending +in the replay caches. The next round passes ``num_accepted_tokens`` (how +many of those pending drafts the sampler accepted) and the kernel replays +them before the new tokens. No host-side promotion of KDA SSM/conv state is +required or allowed. + +Productization deltas vs the drop's host wrapper (``reference.py``): + +* ``zero_accepted_hint`` / ``regular_metadata_hint`` are explicit caller + arguments. The drop derived them by *reading the device tensors* + (``torch.count_nonzero(...).item()`` / ``torch.equal``) behind + ``id()``-keyed caches — a host-device sync per novel tensor object plus a + stale-cache hazard on id reuse. Callers that statically know the pattern + (benchmarks, the first verify round after prefill) may pass the hints; + the runtime default (``False``/``False``) is always correct and never + syncs. +* The compile cache is keyed purely by dtype/shape/stride layouts and + constexpr flags — no ``id()`` or ``data_ptr`` keys. +* Bias / ``pad_slot_id`` arguments (unsupported by the specialized kernel, + previously validated-then-rejected) are dropped from the signature. Gated + output norm and MXFP8 quantization are optional packed-layout epilogues. + +Kernel shape contract: ``K == V == 128``, conv width ``W == 4``, +``HV == H``, one V pass (TILE_V=128, 512 threads). ``num_spec`` is a +compile-time constant per cache allocation; every shape compiles the same +kernel structure (the drop's benchmark-only setmaxreg / register-weight +specializations were retired with the latency restructuring). + +Kernel selection: calls on the runtime contract (packed token rows, fused gated +RMSNorm with a bf16 output, 3xTF32 recurrence, vector-aligned operands) run +``kda_decode_mtp_v2_kernel`` (``kda_mtp_decode_v2.py``: the same arithmetic in +the same order, bitwise identical results, shorter critical path) unless +``TRTLLM_K3_KDA_MTP_V2=0``. On SM107 the same contract with the MXFP8 output runs +``kda_decode_mtp_2cta_kernel`` with its MXFP8 epilogue (one CTA per (head, request) +tile, two CTAs per SM, bitwise identical results). Every other variant runs +``kda_decode_mtp_kernel``, which on SM107 uses its compact two-CTAs-per-SM layout on +multi-wave grids with at most three drafts. + +The launcher runs the private kernels in ``kernels/kda`` and is not a registered torch op (main's +``custom_ops`` registers its own ``trtllm::kda_mtp_decode`` / ``kda_mtp_decode_fp8`` over main's +kernel, with the replay-only commit the cache manager uses). It reads and writes the replay caches +and the slot-indexed accepted counts main's KDA cache manager allocates, in main's layout. The +switches above are read per call. ``kernels/kda/_kda_kernels.py`` imports the entry (and through it this +module) on first use; importing it raises ``ImportError`` without CuTe DSL. +""" + +import contextlib +import os +from math import gcd +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm.logger import logger + +if IS_CUTLASS_DSL_AVAILABLE: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + import cutlass.utils.blackwell_helpers as sm100_utils + from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 + from cutlass.cute.runtime import from_dlpack + + from .kda_mtp_decode import NUM_THREADS, SPLIT_MAX_TILES, T_PAD, TILE_K, kda_decode_mtp_kernel + from .kda_mtp_decode_2cta import kda_decode_mtp_2cta_kernel + from .kda_mtp_decode_v2 import kda_decode_mtp_v2_kernel +else: + raise ImportError("Kimi K3 KDA MTP decode requires NVIDIA CUTLASS DSL") + +# CTAs per SM the two-CTA verify kernel is compiled for (64 registers per thread). +_TWO_CTA_MIN_BLOCKS = 2 +# CTAs per SM the compact layout of ``kda_decode_mtp_kernel`` is launched for (64 registers per +# thread; see ``COMPACT_TMEM_COLS`` in kda_mtp_decode.py). +_COMPACT_MIN_BLOCKS = 2 +# One V pass: the kernel runs 16 warps x 8 rows over the full 128-row state. +_TILE_V = 128 + + +if IS_CUTLASS_DSL_AVAILABLE: + + @cute.jit + def _run_kda_decode_mtp( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + output_scale: cute.Tensor, + ht: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + stage_timing: cute.Tensor, + ssm_state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + num_accepted_tokens: cute.Tensor, + precompute_control: cute.Tensor, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + N: cutlass.Int32, + NUM_SPEC: cutlass.Constexpr[int], + TILE_V: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + scale_leading_dim: cutlass.Constexpr[int], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + USE_SETMAXREG: cutlass.Constexpr[bool], + USE_REGULAR_METADATA: cutlass.Constexpr[bool], + USE_PACKED_TOKEN_LAYOUT: cutlass.Constexpr[bool], + USE_REG_Q_WEIGHTS: cutlass.Constexpr[bool], + USE_ZERO_ACCEPTED: cutlass.Constexpr[bool], + FUSE_PRECOMPUTE: cutlass.Constexpr[bool], + RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], + FUSE_OUTPUT_NORM: cutlass.Constexpr[bool], + QUANTIZE_OUTPUT: cutlass.Constexpr[bool], + PROFILE_STAGES: cutlass.Constexpr[bool], + SPLIT_V: cutlass.Constexpr[int], + BF16_MMA: cutlass.Constexpr[bool], + stream: cuda.CUstream, + COMPACT: cutlass.Constexpr[bool] = False, + ): + if cutlass.const_expr(USE_ZERO_ACCEPTED): + t_max = 1 + NUM_SPEC + else: + t_max = 2 * NUM_SPEC + 1 + smem_qk_layout = cute.make_layout((t_max, K), stride=(K, 1)) + # BF16_MMA: bf16 tensor-core operands ([K~; Q~] only, K = 16 per instruction). + OP_DTYPE = cutlass.BFloat16 if BF16_MMA else cutlass.Float32 + N1 = 2 * T_PAD if BF16_MMA else 3 * T_PAD + # GEMM1: M = V rows of the state (TMEM), N = operand rows ([K~; Q~; K~lo] for tf32). + tiled_mma1 = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, N1), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1 = sm100_utils.make_smem_layout_b(tiled_mma1, (V, N1, K), OP_DTYPE, 1) + # fp32 A-fragment layout of the (V x K) state in TMEM: the GEMM2 accumulator seed / state + # store views use it in every mode. + tiled_mma_f32a = sm100_utils.make_trivial_tiled_mma( + cutlass.Float32, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, 3 * T_PAD), + tcgen05.OperandSource.TMEM, + ) + # GEMM1b: the S0_lo sweep only needs K~hi (N = T_PAD) -> a third of the N=48 work. + tiled_mma1k = sm100_utils.make_trivial_tiled_mma( + cutlass.Float32, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1k = sm100_utils.make_smem_layout_b( + tiled_mma1k, (V, T_PAD, K), cutlass.Float32, 1 + ) + # TMA bulk store of the committed state: (V, K, tiles) view of ht, one 128B-swizzled + # (V/SPLIT_V x 32 fp32 | 64 bf16) box per 128 B-wide column slab. The view takes ht's own strides (the + # per-layer state cache is a strided slice of a larger buffer); the tile mode is the pool + # index (flat layout) or the (head, slot) pair. + if cutlass.const_expr(cute.rank(ht.shape) == 3): + mT_vkl = cute.make_tensor( + ht.iterator, + cute.make_layout( + (V, K, ht.shape[0]), + stride=(ht.layout.stride[1], ht.layout.stride[2], ht.layout.stride[0]), + ), + ) + else: + mT_vkl = cute.make_tensor( + ht.iterator, + cute.make_layout( + (V, K, (ht.shape[1], ht.shape[0])), + stride=( + ht.layout.stride[2], + ht.layout.stride[3], + (ht.layout.stride[1], ht.layout.stride[0]), + ), + ), + ) + # 128 B swizzled box rows: 32 fp32 or 64 bf16 state columns per box + box_cols = 128 // (h0.element_type.width // 8) + box_atom = sm100_utils.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, h0.element_type + ) + box_layout = cute.tile_to_shape(box_atom, (V, box_cols), order=(0, 1)) + # committed-state store: each CTA writes its V/SPLIT_V rows through (V/SPLIT_V x box_cols) boxes + box_layout_t = cute.tile_to_shape(box_atom, (V // SPLIT_V, box_cols), order=(0, 1)) + tma_atom_t, mT_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mT_vkl, box_layout_t, (V // SPLIT_V, box_cols) + ) + # TMA bulk load of the initial state S0 into the same box layout (issued at kernel start). + if cutlass.const_expr(cute.rank(h0.shape) == 3): + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, h0.shape[0]), + stride=(h0.layout.stride[1], h0.layout.stride[2], h0.layout.stride[0]), + ), + ) + else: + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, (h0.shape[1], h0.shape[0])), + stride=( + h0.layout.stride[2], + h0.layout.stride[3], + (h0.layout.stride[1], h0.layout.stride[0]), + ), + ), + ) + tma_atom_l, mS_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mS_vkl, box_layout, (V, box_cols) + ) + # GEMM2: (V x T_PAD) x (T_PAD x K), both operands in smem, tf32; issued as two N = K/2 + # column halves so the state store can start on half 0. + tiled_mma2h = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2h = sm100_utils.make_smem_layout_b( + tiled_mma2h, (V, K // 2, T_PAD), OP_DTYPE, 1 + ) + smem_layout_a2 = sm100_utils.make_smem_layout_a( + tiled_mma2h, (V, K // 2, T_PAD), OP_DTYPE, 1 + ) + # GEMM2 half 0 also carries the output-combination GEMM: N = T_PAD (Bcoef) + K/2 (Kbar rows 0..63). + tiled_mma2q = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, T_PAD + K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2q = sm100_utils.make_smem_layout_b( + tiled_mma2q, (V, T_PAD + K // 2, T_PAD), OP_DTYPE, 1 + ) + # GEMM3: T x T coupling matrices, A3 = [Kf; Qf] padded to 64 rows, B3 = Kb (16 rows). + tiled_mma3 = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (4 * T_PAD, T_PAD), + tcgen05.OperandSource.SMEM, + ) + smem_layout_a3 = sm100_utils.make_smem_layout_a( + tiled_mma3, (4 * T_PAD, T_PAD, K), OP_DTYPE, 1 + ) + smem_layout_b3 = sm100_utils.make_smem_layout_b( + tiled_mma3, (4 * T_PAD, T_PAD, K), OP_DTYPE, 1 + ) + # Output combination D4[v, t] = sum_s V'[s, v] Bcoef[t, s]: computed as the first 16 columns + # of GEMM2 half 0; this tiled MMA only provides the accumulator layout of those columns. + tiled_mma4 = sm100_utils.make_trivial_tiled_mma( + cutlass.Float32, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.SMEM, + ) + launched = kda_decode_mtp_kernel( + h0, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + o, + output_scale, + ht, + k_cache, + g_cache, + v_cache, + beta_cache, + smem_qk_layout, + ssm_state_indices, + cu_seqlens, + num_accepted_tokens, + precompute_control, + TILE_V, + scale, + HV, + K, + V, + NUM_SPEC, + KERNEL_WIDTH, + lower_bound, + onorm_eps, + scale_leading_dim, + USE_FLAT_LAYOUT, + USE_SETMAXREG, + USE_REGULAR_METADATA, + USE_PACKED_TOKEN_LAYOUT, + USE_REG_Q_WEIGHTS, + USE_ZERO_ACCEPTED, + FUSE_PRECOMPUTE, + RUNTIME_PRECOMPUTE_FLAG, + FUSE_OUTPUT_NORM, + QUANTIZE_OUTPUT, + stage_timing, + PROFILE_STAGES, + tiled_mma1, + tiled_mma3, + tiled_mma1k, + smem_layout_b1, + smem_layout_a2, + smem_layout_a3, + smem_layout_b3, + smem_layout_b1k, + tiled_mma4, + tma_atom_t, + mT_tma, + box_layout, + tma_atom_l, + mS_tma, + tiled_mma2h, + smem_layout_b2h, + box_layout_t, + tiled_mma2q, + smem_layout_b2q, + SPLIT_V, + BF16_MMA, + tiled_mma_f32a, + COMPACT, + ) + if cutlass.const_expr(SPLIT_V == 2): + # cluster pair: the two CTAs of a (head, request) must be co-scheduled (in-place state) + launched.launch( + grid=(HV, N, 2), block=[NUM_THREADS, 1, 1], cluster=(1, 1, 2), stream=stream + ) + elif cutlass.const_expr(COMPACT): + # Two co-resident 512-thread CTAs need <= 64 registers per thread (the sm_107a build + # allocates 75 without the launch bound, which leaves room for one CTA per SM). + launched.launch( + grid=(HV, N, 1), + block=[NUM_THREADS, 1, 1], + stream=stream, + min_blocks_per_mp=_COMPACT_MIN_BLOCKS, + ) + else: + launched.launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) + + +if IS_CUTLASS_DSL_AVAILABLE: + + @cute.jit + def _run_kda_decode_mtp_v2( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + ht: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + fa: cute.Tensor, + w_fb: cute.Tensor, + N: cutlass.Int32, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + SPLIT_V: cutlass.Constexpr[int], + PDL: cutlass.Constexpr[bool], + FUSED_FB: cutlass.Constexpr[bool], + stream: cuda.CUstream, + ): + """Host side of ``kda_decode_mtp_v2_kernel`` (packed rows, fused gated RMSNorm, bf16 + output, 3xTF32): the tf32 MMA atoms, operand smem layouts and state TMA atoms of + ``_run_kda_decode_mtp``, then the launch. With ``FUSED_FB``, ``fa`` is the BF16 + ``[1, T, 128]`` f_a view and ``w_fb`` the BF16 ``[HV * 128, 128]`` f_b weight (else both + are unused placeholders).""" + f32 = cutlass.Float32 + smem_qk_layout = cute.make_layout((2 * NUM_SPEC + 1, K), stride=(K, 1)) + tiled_mma1 = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, 3 * T_PAD), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1 = sm100_utils.make_smem_layout_b(tiled_mma1, (V, 3 * T_PAD, K), f32, 1) + tiled_mma_f32a = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, 3 * T_PAD), + tcgen05.OperandSource.TMEM, + ) + tiled_mma1k = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1k = sm100_utils.make_smem_layout_b(tiled_mma1k, (V, T_PAD, K), f32, 1) + if cutlass.const_expr(cute.rank(ht.shape) == 3): + mT_vkl = cute.make_tensor( + ht.iterator, + cute.make_layout( + (V, K, ht.shape[0]), + stride=(ht.layout.stride[1], ht.layout.stride[2], ht.layout.stride[0]), + ), + ) + else: + mT_vkl = cute.make_tensor( + ht.iterator, + cute.make_layout( + (V, K, (ht.shape[1], ht.shape[0])), + stride=( + ht.layout.stride[2], + ht.layout.stride[3], + (ht.layout.stride[1], ht.layout.stride[0]), + ), + ), + ) + box_cols = 128 // (h0.element_type.width // 8) + box_atom = sm100_utils.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, h0.element_type + ) + box_layout = cute.tile_to_shape(box_atom, (V, box_cols), order=(0, 1)) + box_layout_t = cute.tile_to_shape(box_atom, (V // SPLIT_V, box_cols), order=(0, 1)) + tma_atom_t, mT_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mT_vkl, box_layout_t, (V // SPLIT_V, box_cols) + ) + if cutlass.const_expr(cute.rank(h0.shape) == 3): + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, h0.shape[0]), + stride=(h0.layout.stride[1], h0.layout.stride[2], h0.layout.stride[0]), + ), + ) + else: + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, (h0.shape[1], h0.shape[0])), + stride=( + h0.layout.stride[2], + h0.layout.stride[3], + (h0.layout.stride[1], h0.layout.stride[0]), + ), + ), + ) + tma_atom_l, mS_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mS_vkl, box_layout, (V, box_cols) + ) + tiled_mma2h = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2h = sm100_utils.make_smem_layout_b(tiled_mma2h, (V, K // 2, T_PAD), f32, 1) + smem_layout_a2 = sm100_utils.make_smem_layout_a(tiled_mma2h, (V, K // 2, T_PAD), f32, 1) + tiled_mma2q = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, T_PAD + K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2q = sm100_utils.make_smem_layout_b( + tiled_mma2q, (V, T_PAD + K // 2, T_PAD), f32, 1 + ) + tiled_mma3 = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (4 * T_PAD, T_PAD), + tcgen05.OperandSource.SMEM, + ) + smem_layout_a3 = sm100_utils.make_smem_layout_a(tiled_mma3, (4 * T_PAD, T_PAD, K), f32, 1) + smem_layout_b3 = sm100_utils.make_smem_layout_b(tiled_mma3, (4 * T_PAD, T_PAD, K), f32, 1) + tiled_mma4 = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.SMEM, + ) + if cutlass.const_expr(FUSED_FB): + # W_fb rows of one head: two 128 x 64 BF16 boxes with 128B swizzle. + fb_atom = sm100_utils.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + mFb = cute.make_tensor( + w_fb.iterator, cute.make_layout((w_fb.shape[0], K), stride=(K, 1)) + ) + tma_atom_fb, mFb_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mFb, + cute.tile_to_shape(fb_atom, (K, 64), order=(0, 1)), + (K, 64), + ) + fb_layout = cute.tile_to_shape(fb_atom, (K, 64, K // 64), (0, 1, 2)) + else: + tma_atom_fb, mFb_tma, fb_layout = tma_atom_l, mS_tma, box_layout + launched = kda_decode_mtp_v2_kernel( + h0, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + o, + k_cache, + g_cache, + v_cache, + beta_cache, + ssm_state_indices, + num_accepted_tokens, + smem_qk_layout, + tiled_mma1, + tiled_mma3, + tiled_mma1k, + smem_layout_b1, + smem_layout_a2, + smem_layout_a3, + smem_layout_b3, + smem_layout_b1k, + tiled_mma4, + tma_atom_t, + mT_tma, + box_layout, + tma_atom_l, + mS_tma, + tiled_mma2h, + smem_layout_b2h, + box_layout_t, + tiled_mma2q, + smem_layout_b2q, + tiled_mma_f32a, + fa, + tma_atom_fb, + mFb_tma, + fb_layout, + scale, + HV, + K, + V, + NUM_SPEC, + KERNEL_WIDTH, + lower_bound, + onorm_eps, + USE_FLAT_LAYOUT, + SPLIT_V, + PDL, + FUSED_FB, + ) + if cutlass.const_expr(SPLIT_V == 2): + # cluster pair: the two CTAs of a (head, request) must be co-scheduled (in-place state) + launched.launch( + grid=(HV, N, 2), + block=[NUM_THREADS, 1, 1], + cluster=(1, 1, 2), + stream=stream, + use_pdl=PDL, + ) + else: + launched.launch( + grid=(HV, N, 1), + block=[NUM_THREADS, 1, 1], + stream=stream, + use_pdl=PDL, + ) + + +if IS_CUTLASS_DSL_AVAILABLE: + + @cute.jit + def _run_kda_decode_mtp_2cta( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + o8: cute.Tensor, + osf: cute.Tensor, + N: cutlass.Int32, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + PDL: cutlass.Constexpr[bool], + MIN_BLOCKS: cutlass.Constexpr[int], + stream: cuda.CUstream, + EARLY_TRIGGER: cutlass.Constexpr[bool] = False, + ): + """Host side of ``kda_decode_mtp_2cta_kernel``: v2's tf32 MMA atoms, operand smem + layouts and state TMA atoms (fp32 or bf16 state), launched for ``MIN_BLOCKS`` CTAs per SM + (64 registers per thread at 2). The kernel writes the E4M3 output ``o8`` and its R128c4 + UE8M0 scales ``osf``.""" + f32 = cutlass.Float32 + tiled_mma1 = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, 3 * T_PAD), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1 = sm100_utils.make_smem_layout_b(tiled_mma1, (V, 3 * T_PAD, K), f32, 1) + tiled_mma_f32a = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, 3 * T_PAD), + tcgen05.OperandSource.TMEM, + ) + tiled_mma1k = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1k = sm100_utils.make_smem_layout_b(tiled_mma1k, (V, T_PAD, K), f32, 1) + tiled_mma2h = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2h = sm100_utils.make_smem_layout_b(tiled_mma2h, (V, K // 2, T_PAD), f32, 1) + smem_layout_a2 = sm100_utils.make_smem_layout_a(tiled_mma2h, (V, K // 2, T_PAD), f32, 1) + tiled_mma2q = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, T_PAD + K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2q = sm100_utils.make_smem_layout_b( + tiled_mma2q, (V, T_PAD + K // 2, T_PAD), f32, 1 + ) + tiled_mma3 = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (4 * T_PAD, T_PAD), + tcgen05.OperandSource.SMEM, + ) + smem_layout_a3 = sm100_utils.make_smem_layout_a(tiled_mma3, (4 * T_PAD, T_PAD, K), f32, 1) + smem_layout_b3 = sm100_utils.make_smem_layout_b(tiled_mma3, (4 * T_PAD, T_PAD, K), f32, 1) + tiled_mma4 = sm100_utils.make_trivial_tiled_mma( + f32, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.SMEM, + ) + if cutlass.const_expr(cute.rank(h0.shape) == 3): + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, h0.shape[0]), + stride=(h0.layout.stride[1], h0.layout.stride[2], h0.layout.stride[0]), + ), + ) + else: + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, (h0.shape[1], h0.shape[0])), + stride=( + h0.layout.stride[2], + h0.layout.stride[3], + (h0.layout.stride[1], h0.layout.stride[0]), + ), + ), + ) + # 128 B swizzled box rows: 32 fp32 or 64 bf16 state columns per box + box_cols = 128 // (h0.element_type.width // 8) + box_atom = sm100_utils.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, h0.element_type + ) + box_layout = cute.tile_to_shape(box_atom, (V, box_cols), order=(0, 1)) + tma_atom_l, mS_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mS_vkl, box_layout, (V, box_cols) + ) + tma_atom_t, mT_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mS_vkl, box_layout, (V, box_cols) + ) + launched = kda_decode_mtp_2cta_kernel( + h0, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + k_cache, + g_cache, + v_cache, + beta_cache, + ssm_state_indices, + num_accepted_tokens, + o8, + osf, + tiled_mma1, + tiled_mma3, + tiled_mma1k, + smem_layout_b1, + smem_layout_a2, + smem_layout_a3, + smem_layout_b3, + smem_layout_b1k, + tiled_mma4, + tiled_mma2h, + smem_layout_b2h, + tiled_mma2q, + smem_layout_b2q, + tiled_mma_f32a, + tma_atom_l, + mS_tma, + box_layout, + tma_atom_t, + mT_tma, + scale, + HV, + K, + V, + NUM_SPEC, + KERNEL_WIDTH, + lower_bound, + onorm_eps, + USE_FLAT_LAYOUT, + PDL, + EARLY_TRIGGER, + ) + launched.launch( + grid=(HV, N, 1), + block=[NUM_THREADS, 1, 1], + stream=stream, + use_pdl=PDL, + min_blocks_per_mp=MIN_BLOCKS, + ) + + +def _require_stride_layout( + *, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + g, + beta, + A_log, + dt_bias, + recurrent_state, + k_cache, + g_cache, + v_cache, + beta_cache, + ssm_state_indices, + cu_seqlens, + num_accepted_tokens, + out, + onorm_g, + onorm_weight, + H, + HV, + K, + V, + W, + num_spec, + T_total, + packed_token_layout, + fuse_output_norm, + quantize_output, + fa=None, + w_fb=None, +): + if x_q.ndim != 4 or x_k.ndim != 4 or x_v.ndim != 4: + raise ValueError("Expected x_q/x_k/x_v to have shape [1, T, H, D].") + if x_q.shape != (1, T_total, H, K) or x_k.shape != (1, T_total, H, K): + raise ValueError(f"Expected x_q/x_k shape [1, {T_total}, {H}, {K}].") + if x_v.shape != (1, T_total, HV, V): + raise ValueError(f"Expected x_v shape [1, {T_total}, {HV}, {V}].") + if fa is not None: + if fa.ndim != 3 or fa.shape != (1, T_total, K) or fa.dtype != torch.bfloat16: + raise ValueError(f"Expected bf16 fa shape [1, {T_total}, {K}].") + if w_fb.shape != (HV * K, K) or w_fb.dtype != torch.bfloat16 or not w_fb.is_contiguous(): + raise ValueError(f"Expected a contiguous bf16 w_fb of shape [{HV * K}, {K}].") + if w_fb.data_ptr() % 16: + raise ValueError("Expected w_fb to be 16-byte aligned (TMA source).") + elif g.ndim != 4 or g.shape != (1, T_total, HV, K): + raise ValueError(f"Expected g shape [1, {T_total}, {HV}, {K}].") + if beta.ndim != 3 or beta.shape != (1, T_total, HV): + raise ValueError(f"Expected beta shape [1, {T_total}, {HV}].") + if quantize_output: + if out.ndim != 2 or out.shape != (T_total, HV * V): + raise ValueError(f"Expected FP8 out shape [{T_total}, {HV * V}].") + if out.dtype is not torch.float8_e4m3fn: + raise TypeError("Expected FP8 out dtype torch.float8_e4m3fn.") + elif out.ndim != 4 or out.shape != (1, T_total, HV, V): + raise ValueError(f"Expected BF16 out shape [1, {T_total}, {HV}, {V}].") + if fuse_output_norm: + if not packed_token_layout: + raise ValueError("fuse_output_norm requires packed_token_layout=True.") + if onorm_g is None or onorm_weight is None: + raise ValueError("fuse_output_norm requires onorm_g and onorm_weight.") + if onorm_g.ndim != 4 or onorm_g.shape != (1, T_total, HV, V): + raise ValueError(f"Expected onorm_g shape [1, {T_total}, {HV}, {V}].") + if onorm_g.stride(-1) != 1 or onorm_g.stride(-2) != V: + raise ValueError("Expected onorm_g to have contiguous head/value blocks.") + if onorm_weight.ndim != 1 or onorm_weight.shape[0] != V: + raise ValueError(f"Expected onorm_weight shape [{V}].") + if onorm_weight.stride(0) != 1: + raise ValueError("Expected onorm_weight to be contiguous.") + if quantize_output and not fuse_output_norm: + raise ValueError("quantize_output requires fuse_output_norm=True.") + + last_dim_tensors = { + "x_q": x_q, + "x_k": x_k, + "x_v": x_v, + "g": g if fa is None else fa, + "beta": beta, + "recurrent_state": recurrent_state, + "k_cache": k_cache, + "g_cache": g_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + } + if not quantize_output: + last_dim_tensors["out"] = out + for name, tensor in last_dim_tensors.items(): + if tensor.stride(-1) != 1: + raise ValueError(f"Expected {name} to be contiguous in its last dimension.") + + if w_q.shape != (H * K, W) or w_k.shape != (H * K, W) or w_v.shape != (HV * V, W): + raise ValueError(f"Expected w_q/w_k shape [{H * K}, {W}] and w_v shape [{HV * V}, {W}].") + if w_q.stride(1) != 1 or w_k.stride(1) != 1 or w_v.stride(1) != 1: + raise ValueError("Expected w_q/w_k/w_v to be contiguous in the kernel-width axis.") + + if A_log.ndim != 1 or A_log.shape[0] != H: + raise ValueError(f"Expected A_log shape [{H}].") + if dt_bias.ndim != 1 or dt_bias.shape[0] != H * K: + raise ValueError(f"Expected dt_bias shape [{H * K}].") + + state_s = W - 1 + num_spec + if cs_q.ndim != 3 or cs_k.ndim != 3 or cs_v.ndim != 3: + raise ValueError("Expected cs_q/cs_k/cs_v to have shape [pool, dim, S].") + if cs_q.shape[1] != H * K or cs_k.shape[1] != H * K: + raise ValueError(f"Expected cs_q/cs_k shape [pool, {H * K}, S].") + if cs_v.shape[1] != HV * V: + raise ValueError(f"Expected cs_v shape [pool, {HV * V}, S].") + if cs_q.shape[2] < state_s or cs_k.shape[2] < state_s or cs_v.shape[2] < state_s: + raise ValueError(f"Expected conv-state S dimension to be at least {state_s}.") + if cs_q.stride(1) != 1 or cs_k.stride(1) != 1 or cs_v.stride(1) != 1: + raise ValueError( + "Expected cs_q/cs_k/cs_v to use dim-contiguous layout " + "(allocate as [pool, S, dim] and transpose(1, 2))." + ) + + pool_size = recurrent_state.shape[0] + if recurrent_state.ndim != 4 or recurrent_state.shape[1:] != (HV, V, K): + raise ValueError( + f"Expected recurrent_state shape [pool, {HV}, {V}, {K}] (V-first pool layout)." + ) + if recurrent_state.dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"recurrent_state must be fp32 or bf16, got {recurrent_state.dtype}") + if k_cache.ndim != 3 or k_cache.shape[1:] != (num_spec, H * K): + raise ValueError(f"Expected k_cache shape [pool, {num_spec}, {H * K}].") + if g_cache.ndim != 3 or g_cache.shape[1:] != (num_spec, H * K): + raise ValueError(f"Expected g_cache shape [pool, {num_spec}, {H * K}].") + if g_cache.dtype != torch.float32: + raise ValueError("g_cache (the replayed gates) must be fp32.") + if v_cache.ndim != 3 or v_cache.shape[1:] != (num_spec, HV * V): + raise ValueError(f"Expected v_cache shape [pool, {num_spec}, {HV * V}].") + if beta_cache.ndim != 3 or beta_cache.shape[1:] != (num_spec, HV): + raise ValueError(f"Expected beta_cache shape [pool, {num_spec}, {HV}].") + if ( + k_cache.shape[0] < pool_size + or g_cache.shape[0] < pool_size + or v_cache.shape[0] < pool_size + or beta_cache.shape[0] < pool_size + ): + raise ValueError("Expected cache pool dimensions to cover recurrent_state rows.") + + if ssm_state_indices.ndim != 1 or num_accepted_tokens.ndim != 1: + raise ValueError("Expected ssm_state_indices and num_accepted_tokens to be 1D.") + num_requests = ssm_state_indices.shape[0] + if packed_token_layout: + if cu_seqlens is not None: + raise ValueError("packed_token_layout derives row offsets; cu_seqlens must be None.") + if T_total != num_requests * (num_spec + 1): + raise ValueError( + "packed_token_layout expects exactly num_spec + 1 new-token rows per request." + ) + if num_accepted_tokens.shape[0] < pool_size: + raise ValueError("Expected slot-indexed num_accepted_tokens to cover the state pool.") + else: + if cu_seqlens is None or cu_seqlens.ndim != 1: + raise ValueError("Expected cu_seqlens to be a 1D tensor.") + if cu_seqlens.shape[0] != num_requests + 1: + raise ValueError("Expected cu_seqlens length to be N + 1.") + if num_accepted_tokens.shape[0] != num_requests: + raise ValueError("Expected num_accepted_tokens length to match N.") + + +def _fits_32bit_stride(tensor: torch.Tensor) -> bool: + int32_max = 2**31 - 1 + max_offset = int(tensor.storage_offset()) + for size, stride in zip(tensor.shape, tensor.stride()): + stride = abs(int(stride)) + if stride > int32_max: + return False + if size: + max_offset += (int(size) - 1) * stride + if max_offset > int32_max: + return False + return True + + +def _from_dlpack_arg(tensor: torch.Tensor, *, assumed_align: int = 16): + return from_dlpack( + tensor, + assumed_align=assumed_align, + use_32bit_stride=_fits_32bit_stride(tensor), + ) + + +def _beta_cache_assumed_align(beta_cache: torch.Tensor) -> int: + """Return the alignment shared by KDA per-layer beta-cache views. + + The producer allocates ``[layers, slots, num_spec, local_heads]``. After + selecting a layer, ``slots * stride(0)`` is the physical layer span even + when the head rows are padded. Combine that span with the current pointer + and dtype, capped at the CuTe bridge's useful 16-byte guarantee. + """ + layer_span_bytes = beta_cache.shape[0] * beta_cache.stride(0) * beta_cache.element_size() + return gcd(16, beta_cache.data_ptr(), layer_span_bytes) + + +def _dlpack_arg(tensor: torch.Tensor, *, assumed_align: int): + # Alignment is deliberately mandatory: layout dynamism does not imply + # arbitrary pointer alignment. Each call site must state the guarantee + # provided by that tensor's producer and view pattern. + arg = _from_dlpack_arg(tensor, assumed_align=assumed_align) + for dim, stride in enumerate(tensor.stride()): + if stride == 1: + return arg.mark_layout_dynamic(dim) + return arg.mark_layout_dynamic() + + +def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False, *, assumed_align: int = 16): + arg = ( + _dlpack_arg(tensor, assumed_align=assumed_align) + if dynamic_layout + else _from_dlpack_arg(tensor, assumed_align=assumed_align) + ) + shape_mask = arg.dynamic_shapes_mask + stride_mask = arg.dynamic_strides_mask + shape = tuple(None if dynamic else size for size, dynamic in zip(tensor.shape, shape_mask)) + stride = tuple( + None if dynamic else value for value, dynamic in zip(tensor.stride(), stride_mask) + ) + return (tensor.dtype, shape, stride, _fits_32bit_stride(tensor), assumed_align) + + +# (device_index, enabled) -> persistent int32 [1] control tensor. Keys are +# plain values (not tensor identities), so entries never go stale. +_precompute_control_cache = {} + + +def _bf16_math_enabled() -> bool: + """Recurrence precision of the verify kernel, from ``TRTLLM_KDA_MTP_BF16_MATH``. + + Unset / 0 (default): 3xTF32 on every product that feeds the state. 1: bf16 tensor-core + operands (bf16 TMEM / smem tiles, K = 16 per instruction, fp32 accumulation -- the + prefill kernel's precision). + """ + value = os.environ.get("TRTLLM_KDA_MTP_BF16_MATH", "0") + if value not in ("0", "1"): + raise ValueError(f"TRTLLM_KDA_MTP_BF16_MATH must be 0 or 1, got {value!r}") + return value == "1" + + +# Verify-kernel schedule. True: calls on the runtime contract (packed rows, fused gated RMSNorm, +# bf16 output, 3xTF32 recurrence, vector-aligned operands) run ``kda_decode_mtp_v2_kernel``, which +# produces bitwise the same results with a shorter critical path. False (tests): every call runs +# ``kda_decode_mtp_kernel``. +_MTP_V2 = True + + +def _mtp_v2_enabled() -> bool: + """Whether calls on the runtime contract run ``kda_decode_mtp_v2_kernel`` (``_MTP_V2``).""" + return _MTP_V2 + + +_num_sms_cache = {} + + +def _num_sms(device: torch.device) -> int: + dev = torch.device(device) + idx = dev.index if dev.index is not None else torch.cuda.current_device() + if idx not in _num_sms_cache: + _num_sms_cache[idx] = torch.cuda.get_device_properties(idx).multi_processor_count + return _num_sms_cache[idx] + + +def _compact_layout_enabled(device: torch.device) -> bool: + """Whether ``kda_decode_mtp_kernel`` may use its compact (two CTAs per SM) layout: on SM107. + The layout produces bitwise the results of the default layout.""" + return torch.cuda.get_device_capability(device) == (10, 7) + + +def _mxfp8_two_cta_enabled(device: torch.device) -> bool: + """Whether MXFP8-output verify calls run the two-CTAs-per-SM kernel (``kda_mtp_decode_2cta`` + with its MXFP8 epilogue) instead of ``kda_decode_mtp_kernel``: on SM107. Both kernels produce + bitwise the same results.""" + return torch.cuda.get_device_capability(device) == (10, 7) + + +def verify_ctas(num_heads: int, num_requests: int, device: torch.device) -> int: + """CTAs of the v2 verify kernel's grid for these tiles (cluster pairs up to + ``SPLIT_MAX_TILES`` tiles, or one CTA per tile).""" + tiles = num_heads * num_requests + return tiles * (2 if tiles <= SPLIT_MAX_TILES else 1) + + +def _vec4_rows(tensor: torch.Tensor, dim: int = -1) -> bool: + """True when every group of four consecutive elements along the unit-stride dimension `dim` + that starts at a multiple of four is one aligned 8 B (bf16) or 16 B (fp32) access: the v2 + kernel moves its per-token rows with such vector loads and stores.""" + if tensor.dtype not in (torch.bfloat16, torch.float32) or tensor.stride(dim) != 1: + return False + if tensor.data_ptr() % (4 * tensor.element_size()): + return False + return all( + stride % 4 == 0 for d, stride in enumerate(tensor.stride()) if d != dim % tensor.ndim + ) + + +def _v2_operands_ok(*, w_q, w_k, w_v, rows, conv_caches) -> bool: + """Operand layouts / dtypes the v2 kernel's vector accesses support (see `_vec4_rows`): + token rows and replay caches vectorized along their last dimension, the dim-contiguous conv + caches along the channel dimension, conv weights as contiguous [channel][4] rows.""" + weights_ok = all( + w.dtype == w_q.dtype and w.is_contiguous() and w.shape[1] == 4 and _vec4_rows(w) + for w in (w_q, w_k, w_v) + ) + return ( + weights_ok + and all(_vec4_rows(t) for t in rows) + and all(_vec4_rows(t, dim=1) for t in conv_caches) + ) + + +def _precompute_control_tensor(device: torch.device, enabled: bool) -> torch.Tensor: + dev = torch.device(device) + key = (dev.index if dev.index is not None else torch.cuda.current_device(), bool(enabled)) + if key not in _precompute_control_cache: + _precompute_control_cache[key] = torch.tensor( + [1 if enabled else 0], dtype=torch.int32, device=dev + ) + return _precompute_control_cache[key] + + +def _try_flatten_args( + *, + recurrent_state: torch.Tensor, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + T_total: int, + H: int, + HV: int, + K: int, + V: int, +) -> Tuple[bool, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + try: + h0 = recurrent_state.view(-1, V, K) + x_q_flat = x_q.view(1, T_total, H * K) + x_k_flat = x_k.view(1, T_total, H * K) + x_v_flat = x_v.view(1, T_total, HV * V) + except RuntimeError: + return False, recurrent_state, x_q, x_k, x_v + return True, h0, x_q_flat, x_k_flat, x_v_flat + + +# Layout-and-constexpr-keyed compile cache. Request count and packed-token +# length are dynamic; batches sharing the same kernel variant reuse one +# artifact even when their launch grid and token-buffer extents differ. +_compiled_cache = {} + +# c240 (TRTLLM_K3_KDA_MX_PDL, set by the KDA mixer around one verify call through +# :func:`mxfp8_pdl_early_trigger`): an MXFP8-output verify launched with ``use_pdl`` runs the +# two-CTA kernel as a PDL launch that releases its dependent right after griddepcontrol.wait. +MXFP8_PDL_EARLY_TRIGGER = False + + +@contextlib.contextmanager +def mxfp8_pdl_early_trigger(enabled: bool): + """Within the block, MXFP8-output verify calls with ``use_pdl`` take the two-CTA kernel's PDL + launch with the early trigger (only for a consumer that waits before reading its output).""" + global MXFP8_PDL_EARLY_TRIGGER + saved, MXFP8_PDL_EARLY_TRIGGER = MXFP8_PDL_EARLY_TRIGGER, bool(enabled) + try: + yield + finally: + MXFP8_PDL_EARLY_TRIGGER = saved + + +def kda_mtp_decode_impl( + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + out: Optional[torch.Tensor] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + fuse_output_norm: bool = False, + quantize_output: bool = False, + fa: Optional[torch.Tensor] = None, + w_fb: Optional[torch.Tensor] = None, + use_pdl: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Launch the fused KDA MTP verify kernel. See the module docstring. + + Args (device tensors unless noted): + x_q/x_k/x_v: post-projection, pre-conv token states + ``[1, T_total, H, 128]`` bf16 — new tokens only, ``1 + + num_spec`` per request. Rows may be described by ``cu_seqlens`` + or packed uniformly by request. + w_q/w_k/w_v: conv weights ``[H*128, W]`` fp32, width-contiguous. + cs_q/cs_k/cs_v: extended conv caches ``[pool, H*128, >= W-1+M]`` + fp32, dim-contiguous. Columns ``[0, W-1)`` are the committed + window; tail columns hold raw pending-draft inputs. Mutated. + g, beta: raw gate ``[1, T, H, 128]`` and beta ``[1, T, H]`` bf16. + A_log, dt_bias: fp32 ``[H]`` / ``[H*128]``. + recurrent_state: pool ``[pool, H, V, K]`` fp32 or bf16, **V-first** layout + (matches the executor ssm pool and the single-token decode + kernel). Committed in place. + k_cache/g_cache/v_cache/beta_cache: replay caches ``[pool, M, H*K]`` (k, + fp32 or bf16) / ``[pool, M, H*K]`` (gate, fp32) / ``[pool, M, H*V]`` / + ``[pool, M, H]`` (fp32 or bf16). Mutated. + ssm_state_indices: per-request state-pool slots. On the legacy + metadata path, ``cu_seqlens`` supplies token offsets ``[N+1]`` + and accepted counts are request-indexed ``[N]``. On the packed + path, ``cu_seqlens`` is absent and accepted counts are pool-level. + zero_accepted_hint: caller asserts every ``num_accepted_tokens`` is + zero (compiles the smaller-smem no-replay variant). Wrong hints + produce wrong results — pass True only when statically known. + regular_metadata_hint: caller asserts ``cu_seqlens`` is the uniform + ``arange * (2*num_spec+1)`` pattern and ``ssm_state_indices`` + is ``arange(N)`` (benchmark identity layout). + packed_token_layout: input projections contain exactly ``1 + + num_spec`` rows per request. ``cu_seqlens`` must be ``None`` and + ``num_accepted_tokens`` is indexed by state-pool slot. The + kernel derives every new-token row from the request index. + fuse_output_norm: on the packed path, apply gated RMSNorm in the CuTe + epilogue using section-strided ``onorm_g`` and ``onorm_weight``. + quantize_output: return the normalized activation directly as MXFP8 + E4M3 plus packed UE8M0 1x128 scales for a prequantized projection. + fa, w_fb: fused forget gate (v2 contract only): ``fa`` is the BF16 + ``[1, T_total, 128]`` f_a projection (rows may be strided, e.g. the + first 128 columns of the ``[f_a | b]`` output) and ``w_fb`` the BF16 + ``[H * 128, 128]`` f_b weight; the kernel computes the forget gate + ``bf16(f_a @ w_fb^T)`` itself and ``g`` must be ``None``. + use_pdl: launch as a programmatic dependent of the previous grid (v2 + contract only); the kernel loads its state, replay entries and + weights before ``griddepcontrol.wait``. + + Returns either BF16 ``[1, T_total, H, V]`` or, with + ``quantize_output=True``, E4M3 ``[T_total, H*V]`` and packed UE8M0 scales + in the flat uint8 R128c4 layout consumed by the CuTe MXFP8 GEMM. + """ + _, T_total, _, D = x_q.shape + H = A_log.shape[0] + fused_fb = fa is not None + if fused_fb != (w_fb is not None) or fused_fb == (g is not None): + raise ValueError("pass either g, or fa and w_fb (fused forget gate)") + HV = H if fused_fb else (g.shape[2] if g.ndim == 4 else H) + K = D + V_dim = x_v.shape[-1] + W = w_q.shape[1] + if K != TILE_K or V_dim != 128 or W != 4: + raise ValueError("specialized kernel expects K=128, V=128, W=4") + if HV != H: + raise ValueError("specialized kernel expects HV == H") + if scale is None: + scale = K**-0.5 + + if packed_token_layout and regular_metadata_hint: + raise ValueError("packed_token_layout and regular_metadata_hint are mutually exclusive.") + if packed_token_layout: + N = ssm_state_indices.shape[0] + elif cu_seqlens is None: + raise ValueError("cu_seqlens may be None only with packed_token_layout=True.") + else: + N = cu_seqlens.shape[0] - 1 + if quantize_output: + if out is not None: + raise ValueError("quantize_output does not accept a caller-provided out tensor.") + out = torch.empty( + (T_total, HV * V_dim), + dtype=torch.float8_e4m3fn, + device=x_q.device, + ) + scale_leading_dim = (T_total + 127) // 128 * 128 + output_scale = torch.empty( + (scale_leading_dim * HV * 4,), + dtype=torch.uint8, + device=x_q.device, + ) + elif out is None: + output_shape = (1, T_total, HV, V_dim) + if packed_token_layout: + # Packed outputs contain new-token rows only, and the kernel + # writes every head/value element. Avoid a redundant zero-fill. + out = torch.empty(output_shape, dtype=x_q.dtype, device=x_q.device) + else: + # Legacy layouts include replay-position holes that remain zero. + out = torch.zeros(output_shape, dtype=x_q.dtype, device=x_q.device) + if num_accepted_tokens.dtype != torch.int32: + num_accepted_tokens = num_accepted_tokens.to(torch.int32) + if not quantize_output: + output_scale = num_accepted_tokens + scale_leading_dim = 1 + _require_stride_layout( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + k_cache=k_cache, + g_cache=g_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=ssm_state_indices, + cu_seqlens=cu_seqlens, + num_accepted_tokens=num_accepted_tokens, + out=out, + onorm_g=onorm_g, + onorm_weight=onorm_weight, + H=H, + HV=HV, + K=K, + V=V_dim, + W=W, + num_spec=num_spec, + T_total=T_total, + packed_token_layout=packed_token_layout, + fuse_output_norm=fuse_output_norm, + quantize_output=quantize_output, + fa=fa, + w_fb=w_fb, + ) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + precompute_control = _precompute_control_tensor(x_q.device, True) + # The packed variant proves at compile time that this tensor is never + # read. Reuse an existing metadata tensor as the CuTe placeholder. + cu_seqlens_arg = ssm_state_indices if cu_seqlens is None else cu_seqlens + # These tensors are compile-time dead in the raw-output variant. Reuse + # existing aligned arguments instead of allocating placeholders. + onorm_g_arg = onorm_g if onorm_g is not None else out + onorm_weight_arg = onorm_weight if onorm_weight is not None else A_log + + use_flat_layout, h0_arg, x_q_arg, x_k_arg, x_v_arg = _try_flatten_args( + recurrent_state=recurrent_state, + x_q=x_q, + x_k=x_k, + x_v=x_v, + T_total=T_total, + H=H, + HV=HV, + K=K, + V=V_dim, + ) + pool_size = h0_arg.shape[0] + # The restructured kernel keeps the conv weights in smem and prefetches + # the recurrent state before precompute; the drop's setmaxreg and + # register-weight specializations no longer apply to any shape. + use_setmaxreg = False + use_reg_q_weights = False + use_regular_metadata = bool(regular_metadata_hint) + # The kernel's USE_ZERO_ACCEPTED fast path unrolls exactly + # 1 + NUM_SPEC == 3 new tokens, so it is only valid for num_spec == 2. + # For other num_spec fall back to the generic loop (the hint implies + # num_accepted_tokens is all zeros, so the generic path computes the + # same result). + use_zero_accepted = bool(zero_accepted_hint) and num_spec == 2 + # stage_timing is unused (PROFILE_STAGES=False); pass `out` as the + # placeholder tensor argument like the drop's runner does. The alias is + # only valid while profiling stays off: with PROFILE_STAGES=True the + # kernel writes int64 stage deltas through this tensor, corrupting the + # bf16 output. Enabling profiling requires a dedicated int64 buffer of + # at least HV * N * 4 elements. + profile_stages = False + assert not profile_stages, ( + "stage_timing aliases `out`; allocate a dedicated int64 [HV * N * 4] " + "buffer before enabling PROFILE_STAGES" + ) + stage_timing_arg = out + beta_cache_assumed_align = _beta_cache_assumed_align(beta_cache) + output_scale_assumed_align = 16 if quantize_output else 4 + # Two CTAs per (head, request) while they fit in one wave (see SPLIT_MAX_TILES). + split_v = 2 if HV * N <= SPLIT_MAX_TILES else 1 + bf16_math = _bf16_math_enabled() + + gate_rows = (fa,) if fused_fb else (g,) + + def v2_operands_ok(): + return _v2_operands_ok( + w_q=w_q, + w_k=w_k, + w_v=w_v, + rows=(x_q_arg, x_k_arg, x_v_arg, *gate_rows, onorm_g, k_cache, g_cache, v_cache), + conv_caches=(cs_q, cs_k, cs_v), + ) + + def two_cta_operands_ok(): + return ( + h0_arg.dtype in (torch.float32, torch.bfloat16) + and all(t.dtype == torch.bfloat16 for t in (x_q, x_k, x_v, cs_q, cs_k, cs_v)) + and g_cache.dtype == torch.float32 + ) + + # MXFP8 output: the two-CTAs-per-SM kernel with kda_decode_mtp_kernel's quantization epilogue + # (the same arithmetic, see kda_mtp_decode_2cta.py), one CTA per tile at every grid size. Two + # of its tiles share an SM on multi-wave grids; on SM107 a tile alone on an SM also finishes + # sooner than in kda_decode_mtp_kernel, with or without its cluster pair. + # c240 (mxfp8_pdl_early_trigger): the MXFP8 two-CTA kernel also as a PDL launch that releases + # its dependent right after its griddepcontrol.wait, for a consumer that only streams weights + # before its own wait (the mid-M o_proj GEMM with the all-reduce push). + mx_pdl = bool(quantize_output and use_pdl and MXFP8_PDL_EARLY_TRIGGER) + use_mxfp8_2cta = ( + quantize_output + and _mxfp8_two_cta_enabled(x_q.device) + and packed_token_layout + and fuse_output_norm + and not fused_fb + and (not use_pdl or mx_pdl) + and not bf16_math + and num_spec <= 7 + and two_cta_operands_ok() + and v2_operands_ok() + ) + mx_pdl = mx_pdl and use_mxfp8_2cta + use_v2 = ( + (_mtp_v2_enabled() or fused_fb or use_pdl) + and packed_token_layout + and fuse_output_norm + and not quantize_output + and not bf16_math + and num_spec <= 7 + and v2_operands_ok() + ) + if (fused_fb or (use_pdl and not mx_pdl)) and not use_v2: + raise ValueError( + "the fused forget gate and PDL launch need the v2 kernel contract (packed rows, " + "fused output norm, bf16 output, 3xTF32, vector-aligned operands)" + ) + if use_mxfp8_2cta: + # E4M3 [T, HV * V] and the R128c4 UE8M0 scales (no BF16 output) + args_2cta = ( + (h0_arg, 16), + (x_q_arg, 16), + (x_k_arg, 16), + (x_v_arg, 16), + (w_q, 16), + (w_k, 16), + (w_v, 16), + (cs_q, 16), + (cs_k, 16), + (cs_v, 16), + (A_log, 16), + (g, 16), + (dt_bias, 16), + (beta, 16), + (onorm_g, 16), + (onorm_weight, 16), + (k_cache, 16), + (g_cache, 16), + (v_cache, 16), + (beta_cache, beta_cache_assumed_align), + (ssm_state_indices, 4), + (num_accepted_tokens, 4), + (out, 16), + (output_scale, 16), + ) + key_2cta = ( + "2cta", + tuple(_layout_key(t, dynamic_layout=True, assumed_align=a) for t, a in args_2cta), + scale, + HV, + K, + V_dim, + num_spec, + W, + float(lower_bound), + float(onorm_eps), + use_flat_layout, + bool(use_pdl), + _TWO_CTA_MIN_BLOCKS, + bool(mx_pdl), + ) + if key_2cta not in _compiled_cache: + logger.info( + f"kda_mtp_decode: compiling 2-CTA MXFP8 variant H={HV} num_spec={num_spec}" + f" pdl={bool(use_pdl)} early_trigger={bool(mx_pdl)}" + ) + _compiled_cache[key_2cta] = cute.compile( + _run_kda_decode_mtp_2cta, + *[_dlpack_arg(t, assumed_align=a) for t, a in args_2cta], + N, + scale=scale, + HV=HV, + K=K, + V=V_dim, + NUM_SPEC=num_spec, + KERNEL_WIDTH=W, + lower_bound=float(lower_bound), + onorm_eps=float(onorm_eps), + USE_FLAT_LAYOUT=use_flat_layout, + PDL=bool(use_pdl), + MIN_BLOCKS=_TWO_CTA_MIN_BLOCKS, + stream=stream, + EARLY_TRIGGER=bool(mx_pdl), + ) + _compiled_cache[key_2cta]( + *[_dlpack_arg(t, assumed_align=a) for t, a in args_2cta], N, stream + ) + return out, output_scale + if use_v2: + # Placeholders for the operands of the other gate mode (compile-time dead). + g_arg = onorm_g if fused_fb else g + fa_arg = fa if fused_fb else beta + w_fb_arg = w_fb if fused_fb else w_q + v2_args = ( + (h0_arg, 16), + (x_q_arg, 16), + (x_k_arg, 16), + (x_v_arg, 16), + (w_q, 16), + (w_k, 16), + (w_v, 16), + (cs_q, 16), + (cs_k, 16), + (cs_v, 16), + (A_log, 16), + (g_arg, 16), + (dt_bias, 16), + (beta, 16), + (onorm_g, 16), + (onorm_weight, 16), + (out, 16), + (h0_arg, 16), + (k_cache, 16), + (g_cache, 16), + (v_cache, 16), + (beta_cache, beta_cache_assumed_align), + (ssm_state_indices, 4), + (num_accepted_tokens, 4), + (fa_arg, 16 if fused_fb else 2), + (w_fb_arg, 16), + ) + v2_key = ( + "v2", + tuple(_layout_key(t, dynamic_layout=True, assumed_align=a) for t, a in v2_args), + scale, + HV, + K, + V_dim, + num_spec, + W, + float(lower_bound), + float(onorm_eps), + use_flat_layout, + split_v, + bool(use_pdl), + fused_fb, + ) + if v2_key not in _compiled_cache: + logger.info( + f"kda_mtp_decode: compiling v2 variant H={HV} num_spec={num_spec} split_v={split_v}" + f" pdl={bool(use_pdl)} fused_fb={fused_fb}" + ) + _compiled_cache[v2_key] = cute.compile( + _run_kda_decode_mtp_v2, + *[_dlpack_arg(t, assumed_align=a) for t, a in v2_args], + N, + scale=scale, + HV=HV, + K=K, + V=V_dim, + NUM_SPEC=num_spec, + KERNEL_WIDTH=W, + lower_bound=float(lower_bound), + onorm_eps=float(onorm_eps), + USE_FLAT_LAYOUT=use_flat_layout, + SPLIT_V=split_v, + PDL=bool(use_pdl), + FUSED_FB=fused_fb, + stream=stream, + ) + _compiled_cache[v2_key](*[_dlpack_arg(t, assumed_align=a) for t, a in v2_args], N, stream) + return out + + # Throughput regime (one CTA per tile, more tiles than SMs): the compact smem/TMEM layout + # lets two CTAs share an SM. Same arithmetic, so the results are bitwise identical. A + # single partial wave keeps the default layout (slightly lower latency). + compact = ( + split_v == 1 + and HV * N > _num_sms(x_q.device) + and h0_arg.dtype == torch.bfloat16 + and bool(fuse_output_norm) + and num_spec <= 3 + and not bf16_math + and _compact_layout_enabled(x_q.device) + ) + + key = ( + x_q.dtype, + scale, + HV, + K, + V_dim, + num_spec, + W, + pool_size, + lower_bound, + use_flat_layout, + _layout_key(h0_arg), + _layout_key(x_q_arg, dynamic_layout=True, assumed_align=16), + _layout_key(x_k_arg, dynamic_layout=True, assumed_align=16), + _layout_key(x_v_arg, dynamic_layout=True, assumed_align=16), + _layout_key(w_q), + _layout_key(w_k), + _layout_key(w_v), + _layout_key(cs_q), + _layout_key(cs_k), + _layout_key(cs_v), + _layout_key(A_log), + _layout_key(g, dynamic_layout=True, assumed_align=16), + _layout_key(dt_bias), + _layout_key(beta, dynamic_layout=True, assumed_align=16), + _layout_key(onorm_g_arg, dynamic_layout=True, assumed_align=16), + _layout_key(onorm_weight_arg), + _layout_key(out, dynamic_layout=True, assumed_align=16), + _layout_key(output_scale, assumed_align=output_scale_assumed_align), + _layout_key(k_cache), + _layout_key(g_cache), + _layout_key(v_cache), + _layout_key(beta_cache, assumed_align=beta_cache_assumed_align), + _layout_key(ssm_state_indices, dynamic_layout=True, assumed_align=4), + _layout_key(cu_seqlens_arg, dynamic_layout=True, assumed_align=4), + _layout_key(num_accepted_tokens, dynamic_layout=True, assumed_align=4), + use_setmaxreg, + use_regular_metadata, + bool(packed_token_layout), + use_reg_q_weights, + use_zero_accepted, + float(onorm_eps), + bool(fuse_output_norm), + bool(quantize_output), + split_v, + bf16_math, + compact, + ) + + if key not in _compiled_cache: + logger.info( + f"kda_mtp_decode: compiling variant N={N} H={HV} T={T_total} " + f"num_spec={num_spec} zero_accepted={use_zero_accepted} " + f"regular_metadata={use_regular_metadata} " + f"fuse_output_norm={fuse_output_norm} " + f"quantize_output={quantize_output} split_v={split_v} bf16_math={bf16_math} " + f"compact={compact}" + ) + _compiled_cache[key] = cute.compile( + _run_kda_decode_mtp, + _from_dlpack_arg(h0_arg), + _dlpack_arg(x_q_arg, assumed_align=16), + _dlpack_arg(x_k_arg, assumed_align=16), + _dlpack_arg(x_v_arg, assumed_align=16), + _from_dlpack_arg(w_q), + _from_dlpack_arg(w_k), + _from_dlpack_arg(w_v), + _from_dlpack_arg(cs_q), + _from_dlpack_arg(cs_k), + _from_dlpack_arg(cs_v), + _from_dlpack_arg(A_log), + _dlpack_arg(g, assumed_align=16), + _from_dlpack_arg(dt_bias), + _dlpack_arg(beta, assumed_align=16), + _dlpack_arg(onorm_g_arg, assumed_align=16), + _from_dlpack_arg(onorm_weight_arg), + _dlpack_arg(out, assumed_align=16), + _from_dlpack_arg(output_scale, assumed_align=output_scale_assumed_align), + _from_dlpack_arg(h0_arg), + _from_dlpack_arg(k_cache), + _from_dlpack_arg(g_cache), + _from_dlpack_arg(v_cache), + _from_dlpack_arg(beta_cache, assumed_align=beta_cache_assumed_align), + _dlpack_arg(stage_timing_arg, assumed_align=16), + _dlpack_arg(ssm_state_indices, assumed_align=4), + _dlpack_arg(cu_seqlens_arg, assumed_align=4), + _dlpack_arg(num_accepted_tokens, assumed_align=4), + _from_dlpack_arg(precompute_control), + scale=scale, + HV=HV, + K=K, + V=V_dim, + N=N, + NUM_SPEC=num_spec, + TILE_V=_TILE_V, + KERNEL_WIDTH=W, + lower_bound=lower_bound, + onorm_eps=float(onorm_eps), + scale_leading_dim=scale_leading_dim, + USE_FLAT_LAYOUT=use_flat_layout, + USE_SETMAXREG=use_setmaxreg, + USE_REGULAR_METADATA=use_regular_metadata, + USE_PACKED_TOKEN_LAYOUT=bool(packed_token_layout), + USE_REG_Q_WEIGHTS=use_reg_q_weights, + USE_ZERO_ACCEPTED=use_zero_accepted, + FUSE_PRECOMPUTE=True, + RUNTIME_PRECOMPUTE_FLAG=False, + FUSE_OUTPUT_NORM=bool(fuse_output_norm), + QUANTIZE_OUTPUT=bool(quantize_output), + PROFILE_STAGES=profile_stages, + SPLIT_V=split_v, + BF16_MMA=bf16_math, + stream=stream, + COMPACT=compact, + ) + + # Runtime alignment policy: + # - 16 B: model inputs/outputs, parameters, recurrent/conv/kg/v pools, + # output-norm tensors, and the persistent control scalar. + # - beta_cache: derived from its actual per-layer physical span and dtype. + # - 4 B: int32 scheduler metadata, which may be an element-offset view. + # - output_scale: 16 B when separately allocated, otherwise it aliases + # int32 num_accepted_tokens and inherits its 4-byte guarantee. + _compiled_cache[key]( + _dlpack_arg(h0_arg, assumed_align=16), + _dlpack_arg(x_q_arg, assumed_align=16), + _dlpack_arg(x_k_arg, assumed_align=16), + _dlpack_arg(x_v_arg, assumed_align=16), + _dlpack_arg(w_q, assumed_align=16), + _dlpack_arg(w_k, assumed_align=16), + _dlpack_arg(w_v, assumed_align=16), + _dlpack_arg(cs_q, assumed_align=16), + _dlpack_arg(cs_k, assumed_align=16), + _dlpack_arg(cs_v, assumed_align=16), + _dlpack_arg(A_log, assumed_align=16), + _dlpack_arg(g, assumed_align=16), + _dlpack_arg(dt_bias, assumed_align=16), + _dlpack_arg(beta, assumed_align=16), + _dlpack_arg(onorm_g_arg, assumed_align=16), + _dlpack_arg(onorm_weight_arg, assumed_align=16), + _dlpack_arg(out, assumed_align=16), + _dlpack_arg(output_scale, assumed_align=output_scale_assumed_align), + _dlpack_arg(h0_arg, assumed_align=16), + _dlpack_arg(k_cache, assumed_align=16), + _dlpack_arg(g_cache, assumed_align=16), + _dlpack_arg(v_cache, assumed_align=16), + _dlpack_arg(beta_cache, assumed_align=beta_cache_assumed_align), + _dlpack_arg(stage_timing_arg, assumed_align=16), + _dlpack_arg(ssm_state_indices, assumed_align=4), + _dlpack_arg(cu_seqlens_arg, assumed_align=4), + _dlpack_arg(num_accepted_tokens, assumed_align=4), + _dlpack_arg(precompute_control, assumed_align=16), + N, + stream, + ) + + if quantize_output: + return out, output_scale + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/low_m_bf16_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/low_m_bf16_splitk.py new file mode 100644 index 000000000000..bb06bfb87424 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/low_m_bf16_splitk.py @@ -0,0 +1,1037 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Adapted from FlashInfer (https://github.com/flashinfer-ai/flashinfer) under the Apache-2.0 +# license. Original file: flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py +# Original copyright: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# +# TODO(TRTLLM-15304): This file is a vendored copy of the FlashInfer split-K kernel. +# Once FlashInfer releases v0.6.18 (which packages these kernels as importable Python +# modules), remove this file and import directly from flashinfer instead. +# Tracking: https://github.com/flashinfer-ai/flashinfer/pull/4266 +"""SM10x (Blackwell) low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses +import functools + +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 +import torch as _torch +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.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +MAX_M = 32 + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES + base = SplitKTactic(mma_m, mma_n, split_k, min_stages) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(min_stages, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "MAX_M", + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", + "validate_tactic", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + ) -> None: + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel(a, b, c, bias).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=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + 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()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + 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): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + 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. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + 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, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + 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=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + 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, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + 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) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + 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_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + 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) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + 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, + ), + 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) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + 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() + 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): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + + 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]) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def _validate_runtime_tensors(a, b, bias, out, strided_a: bool = False) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any(tensor.dtype != a.dtype for tensor in tensors): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + dense = (b, out) if strided_a else (a, b, out) + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in dense): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if strided_a and a.stride(1) != 1: + raise ValueError("a strided row-major view needs a unit inner stride") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + # cute_ext.tma_load is used for a and b (not out, which uses partition_and_copy). + # TMA requires the non-innermost byte stride to be a multiple of 16 bytes. + # For a dense 2D tensor this is max(stride(0), stride(1)) * element_size. + _TMA_STRIDE_ALIGN = 16 + for tensor, tname in ((a, "a"), (b, "b")): + outer_stride_bytes = max(tensor.stride(0), tensor.stride(1)) * tensor.element_size() + if outer_stride_bytes % _TMA_STRIDE_ALIGN: + raise ValueError( + f"tensor '{tname}' outer byte stride {outer_stride_bytes} is not a " + f"multiple of {_TMA_STRIDE_ALIGN}; TMA load requires " + f"{_TMA_STRIDE_ALIGN}-byte-aligned strides " + f"(shape={tuple(tensor.shape)}, strides={tensor.stride()})." + ) + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a: _torch.Tensor, + b: _torch.Tensor, + bias: _torch.Tensor | None, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + strided_a: bool = False, +) -> _torch.Tensor: + """Run ``A[M,K] @ B[K,N]`` with the split-K BF16/FP16 GEMM kernel. + + Args: + a: Input tensor of shape ``[M, K]``, row-major BF16/FP16. + b: Weight tensor of shape ``[K, N]``, column-major BF16/FP16 + (i.e. the transpose of a ``[N, K]`` contiguous weight matrix). + bias: Optional bias of shape ``[N]``, or ``None``. + out: Pre-allocated output tensor of shape ``[M, N]``, BF16/FP16. + pdl: Whether to use Programmatic Dependent Launch. + tactic: The ``SplitKTactic`` specifying tile and pipeline parameters. + strided_a: Accept ``a`` as a row-major view with any 16-byte-aligned row stride (the + layouts are dynamic: the same compiled kernel). + + Returns: + The ``out`` tensor, filled with the GEMM result. + """ + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out, strided_a)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/__init__.py new file mode 100644 index 000000000000..30676f690dfb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuTe DSL / Triton kernels of the Kimi K3 MLA layer (FP8 KV cast, PDL query RMSNorm).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/mla_kv_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/mla_kv_fp8.py new file mode 100644 index 000000000000..dcac2cddceb3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/mla_kv_fp8.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""FP8 K and V of non-absorption MLA context attention, in one pass. + +The FP8 context MLA FMHA reads contiguous E4M3 K ``[T, H * (nope + rope)]`` and +V ``[T, H * v]``. :func:`mla_kv_to_fp8` writes them straight from the BF16 +kv_b_proj output and k_pe, with the bytes the attention op's context copy +(``quantizeCopyInputToFp8Kernel`` with its unit scale) writes for the +concatenated BF16 K and V: every element becomes +``cvt.rn.satfinite.e4m3(float(x))``. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _e4m3x2_from_bf16x2(pairs): + """Convert int32-packed BF16 pairs to int16-packed E4M3 pairs, first element in the low byte.""" + lo = (pairs << 16).to(tl.float32, bitcast=True) + hi = (pairs & -65536).to(tl.float32, bitcast=True) + return tl.inline_asm_elementwise( + "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", + "=h,f,f", + [lo, hi], + dtype=tl.int16, + is_pure=True, + pack=1, + ) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _mla_kv_fp8_kernel( + kv_ptr, + pe_ptr, + k_ptr, + v_ptr, + num_tokens, + kv_row_stride, + pe_row_stride, + NUM_HEADS: tl.constexpr, + NOPE_PAIRS: tl.constexpr, + ROPE_PAIRS: tl.constexpr, + V_PAIRS: tl.constexpr, + HEADS_PER_PROGRAM: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, +): + rows = tl.program_id(0).to(tl.int64) * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS) + mask = (rows < num_tokens)[:, None] + head0 = tl.program_id(1) * HEADS_PER_PROGRAM + nope = tl.arange(0, NOPE_PAIRS)[None, :] + rope = tl.arange(0, ROPE_PAIRS)[None, :] + v_cols = tl.arange(0, V_PAIRS)[None, :] + kv_rows = kv_ptr + rows[:, None] * kv_row_stride + k_rows = k_ptr + rows[:, None] * (NUM_HEADS * (NOPE_PAIRS + ROPE_PAIRS)) + v_rows = v_ptr + rows[:, None] * (NUM_HEADS * V_PAIRS) + pe = tl.load(pe_ptr + rows[:, None] * pe_row_stride + rope, mask=mask, other=0) + pe = _e4m3x2_from_bf16x2(pe) + for i in tl.static_range(HEADS_PER_PROGRAM): + head = head0 + i + k_head = k_rows + head * (NOPE_PAIRS + ROPE_PAIRS) + x = tl.load(kv_rows + head * NOPE_PAIRS + nope, mask=mask, other=0) + tl.store(k_head + nope, _e4m3x2_from_bf16x2(x), mask=mask) + tl.store(k_head + NOPE_PAIRS + rope, pe, mask=mask) + x = tl.load(kv_rows + NUM_HEADS * NOPE_PAIRS + head * V_PAIRS + v_cols, mask=mask, other=0) + tl.store(v_rows + head * V_PAIRS + v_cols, _e4m3x2_from_bf16x2(x), mask=mask) + + +def _pairs_view(x: torch.Tensor, name: str) -> torch.Tensor: + """int32 view of a 2-D BF16 tensor with unit column stride and even row stride / offset.""" + if ( + x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or x.stride(0) % 2 + or x.storage_offset() % 2 + or x.shape[1] % 2 + ): + raise ValueError( + f"{name} must be a 2-D BF16 view with unit column stride and even width, row stride " + f"and offset; got {x.dtype}, shape {tuple(x.shape)}, stride {x.stride()}." + ) + return x.view(torch.int32) + + +def mla_kv_to_fp8( + kv: torch.Tensor, + k_pe: torch.Tensor, + num_heads: int, + qk_nope_head_dim: int, + v_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """FP8 E4M3 K and V for the FP8 context FMHA of non-absorption MLA. + + Args: + kv: BF16 kv_b_proj output ``[T, H * (nope + v)]``: the nope columns of + all heads, then their V columns. + k_pe: BF16 RoPE key ``[T, rope]`` shared by all heads; unit column + stride and any even row stride (e.g. a ``latent_cache`` view). + num_heads: H. + qk_nope_head_dim: nope; a power of two. + v_head_dim: v; a power of two. + + Returns: + Contiguous K ``[T, H * (nope + rope)]`` (each head's nope columns + followed by k_pe) and V ``[T, H * v]``. + """ + num_tokens, rope_head_dim = k_pe.shape + nope_cols, v_cols = num_heads * qk_nope_head_dim, num_heads * v_head_dim + if kv.dim() != 2 or tuple(kv.shape) != (num_tokens, nope_cols + v_cols): + raise ValueError(f"kv must be [{num_tokens}, {nope_cols + v_cols}], got {tuple(kv.shape)}.") + for name, dim in ( + ("qk_nope_head_dim", qk_nope_head_dim), + ("qk_rope_head_dim", rope_head_dim), + ("v_head_dim", v_head_dim), + ): + if dim < 2 or dim & (dim - 1): + raise ValueError(f"{name} must be a power of two >= 2, got {dim}.") + heads_per_program = 4 if num_heads % 4 == 0 else 1 + kv_pairs = _pairs_view(kv, "kv") + pe_pairs = _pairs_view(k_pe, "k_pe") + k = torch.empty( + (num_tokens, num_heads * (qk_nope_head_dim + rope_head_dim)), + dtype=torch.float8_e4m3fn, + device=kv.device, + ) + v = torch.empty((num_tokens, v_cols), dtype=torch.float8_e4m3fn, device=kv.device) + if num_tokens == 0: + return k, v + block_tokens = 32 + grid = (triton.cdiv(num_tokens, block_tokens), num_heads // heads_per_program) + _mla_kv_fp8_kernel[grid]( + kv_pairs, + pe_pairs, + k.view(torch.int16), + v.view(torch.int16), + num_tokens, + kv_pairs.stride(0), + pe_pairs.stride(0), + NUM_HEADS=num_heads, + NOPE_PAIRS=qk_nope_head_dim // 2, + ROPE_PAIRS=rope_head_dim // 2, + V_PAIRS=v_head_dim // 2, + HEADS_PER_PROGRAM=heads_per_program, + BLOCK_TOKENS=block_tokens, + num_warps=4, + ) + return k, v diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/q_norm_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/q_norm_pdl.py new file mode 100644 index 000000000000..ed342a4320d3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/mla/q_norm_pdl.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RMSNorm of the MLA q_a latent that releases its dependent grid at launch (Kimi K3 decode). + +FlashInfer's CuTe DSL ``RMSNormKernel`` (what ``RMSNorm`` runs for this BF16 norm) with one +change: ``griddepcontrol.launch_dependents`` is issued when a CTA starts instead of after its +store. The arithmetic is the kernel's own code, so the output is bitwise the same. + +Why: the dependent is the q_b projection (``k3_fp8w_lowm_gemm``), a weight-streaming GEMM. With +the trigger at the end it launches only when this norm finishes, after the q_a GEMM. Released +at launch, its CTAs become resident on SMs the q_a GEMM leaves free and stream q_b's weights +while q_a and this norm still run. + +PDL argument for triggering before ``griddepcontrol.wait``: the trigger only lets the dependent +grid launch; what it reads before its own wait must not depend on this grid or anything +upstream that is not complete. The q_b GEMM reads only its weights (written at load time) +before its wait; its activation load waits for this grid's completion, and this grid's +completion implies the q_a GEMM's (this grid waits on it before reading). Grids further down +the chain (q absorb, MLA decode) are released by q_b's own trigger and read only weights +before their waits, as before; the decode additionally has a full graph dependency on the +KV append. So no read moves ahead of its writer. +""" + +from __future__ import annotations + +import functools +import math + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int64 + +try: + from flashinfer.norm.kernels.rmsnorm import RMSNormKernel + from flashinfer.norm.utils import ( + get_cutlass_dtype, + get_sm_version, + predicate_k, + row_reduce_sum_multirow, + ) + + IS_AVAILABLE = True +except ImportError: # pragma: no cover - FlashInfer without the CuTe DSL norm kernels + RMSNormKernel = object + IS_AVAILABLE = False + + +class _LaunchReleaseRMSNormKernel(RMSNormKernel): + """``RMSNormKernel`` whose PDL trigger comes first; everything else is unchanged.""" + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mW: cute.Tensor, + mY: cute.Tensor, + M: Int64, + eps: Float32, + enable_pdl: cutlass.Constexpr[bool], + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + if enable_pdl: + # See the module docstring for why the trigger may precede the wait. + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + H = self.H + cluster_n = self.cluster_n + weight_bias = self.weight_bias + copy_bits = self.copy_bits + threads_per_row = tv_layout.shape[0][0] + rows_per_block = tiler_mn[0] + warps_per_row = max(threads_per_row // 32, 1) + + if cutlass.const_expr(cluster_n > 1): + cluster_y = cute.arch.block_idx()[1] + else: + cluster_y = cutlass.const_expr(0) + + smem = cutlass.utils.SmemAllocator() + + if cutlass.const_expr(self.use_async_copy): + sX = smem.allocate_tensor( + mX.element_type, + cute.make_ordered_layout(tiler_mn, order=(1, 0)), + byte_alignment=16, + ) + + if cutlass.const_expr(cluster_n == 1): + reduction_buffer = smem.allocate_tensor( + Float32, + cute.make_layout((rows_per_block, warps_per_row)), + byte_alignment=4, + ) + mbar_ptr = None + else: + reduction_buffer = smem.allocate_tensor( + Float32, + cute.make_layout((rows_per_block, (warps_per_row, cluster_n))), + byte_alignment=4, + ) + mbar_ptr = smem.allocate_array(cutlass.Int64, num_elems=1) + + if cutlass.const_expr(cluster_n > 1): + if tidx == 0: + cute.arch.mbarrier_init(mbar_ptr, 1) + cute.arch.mbarrier_init_fence() + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + idX = cute.make_identity_tensor(mX.shape) + + gX = cute.local_tile(mX, tiler_mn, (bidx, cluster_y)) + gY = cute.local_tile(mY, tiler_mn, (bidx, cluster_y)) + cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y)) + + mW_expanded_layout = cute.prepend(mW.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + mW_2d = cute.make_tensor(mW.iterator, mW_expanded_layout) + gW = cute.local_tile(mW_2d, tiler_mn, (0, cluster_y)) + + copy_atom_sync = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mX.element_type, + num_bits_per_copy=copy_bits, + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mY.element_type, + num_bits_per_copy=copy_bits, + ) + + if cutlass.const_expr(self.use_async_copy): + copy_atom_async = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + mX.element_type, + num_bits_per_copy=copy_bits, + ) + tiled_copy_load = cute.make_tiled_copy(copy_atom_async, tv_layout, tiler_mn) + else: + tiled_copy_load = cute.make_tiled_copy(copy_atom_sync, tv_layout, tiler_mn) + + tiled_copy_W = cute.make_tiled_copy(copy_atom_sync, tv_layout, tiler_mn) + tiled_copy_store = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn) + + thr_copy_X = tiled_copy_load.get_slice(tidx) + thr_copy_W = tiled_copy_W.get_slice(tidx) + thr_copy_O = tiled_copy_store.get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXcX = thr_copy_X.partition_S(cX) + tXrX = cute.make_fragment_like(tXgX) + + if cutlass.const_expr(self.use_async_copy): + tXsX = thr_copy_X.partition_D(sX) + + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + tXpX = predicate_k(tXcX, limit=H) + tWpW = predicate_k(thr_copy_W.partition_S(cX), limit=H) + row_coord = tXcX[(0, 0), 0, 0] + row_in_bounds = row_coord[0] < M + + if cutlass.const_expr(self.use_async_copy): + if row_in_bounds: + cute.copy(copy_atom_async, tXgX, tXsX, pred=tXpX) + cute.arch.cp_async_commit_group() + + cute.copy(copy_atom_sync, tWgW, tWrW, pred=tWpW) + + cute.arch.cp_async_wait_group(0) + + cute.autovec_copy(tXsX, tXrX) + else: + tXrX.store(cute.zeros_like(tXrX, dtype=mX.element_type)) + if row_in_bounds: + cute.copy(copy_atom_sync, tXgX, tXrX, pred=tXpX) + + cute.copy(copy_atom_sync, tWgW, tWrW, pred=tWpW) + + x = tXrX.load().to(Float32) + x_sq = x * x + sum_sq = row_reduce_sum_multirow( + x_sq, threads_per_row, reduction_buffer, mbar_ptr, cluster_n + ) + + mean_sq = sum_sq / Float32(H) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + + if cutlass.const_expr(cluster_n > 1): + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + else: + cute.arch.barrier() + + if cutlass.const_expr(self.use_async_copy): + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + w = tXrW.load().to(Float32) + y = x * rstd * (w + Float32(weight_bias)) + + tXrO.store(y.to(mY.element_type)) + + if row_in_bounds: + cute.copy(copy_atom_store, tXrO, tXgO, pred=tXpX) + + +@functools.cache +def _compiled(dtype_str: str, H: int, sm_version: int, contiguous: bool): + """Compile like FlashInfer's ``_get_compiled_rmsnorm_kernel`` (weight_bias 0, PDL on).""" + dtype = get_cutlass_dtype(dtype_str) + kernel_obj = _LaunchReleaseRMSNormKernel(dtype, H, 0.0, sm_version=sm_version) + sym_m = cute.sym_int(64) + if contiguous: + elem_bytes = dtype.width // 8 + tensor_align = math.gcd(128, H * elem_bytes) + x_fake = cute.runtime.make_fake_compact_tensor( + dtype, (sym_m, H), stride_order=(1, 0), assumed_align=tensor_align + ) + y_fake = cute.runtime.make_fake_compact_tensor( + dtype, (sym_m, H), stride_order=(1, 0), assumed_align=tensor_align + ) + else: + sym_row_stride_x = cute.sym_int64(divisibility=kernel_obj.vec_size) + sym_row_stride_y = cute.sym_int64(divisibility=kernel_obj.vec_size) + x_fake = cute.runtime.make_fake_tensor( + dtype, (sym_m, H), (sym_row_stride_x, 1), assumed_align=16 + ) + y_fake = cute.runtime.make_fake_tensor( + dtype, (sym_m, H), (sym_row_stride_y, 1), assumed_align=16 + ) + w_fake = cute.runtime.make_fake_compact_tensor(dtype, (H,), assumed_align=16) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel_obj, + x_fake, + w_fake, + y_fake, + Int64(1), + Float32(1e-6), + True, + stream_fake, + options="--enable-tvm-ffi", + ) + + +def can_use(x: torch.Tensor, weight: torch.Tensor) -> bool: + """Whether :func:`rmsnorm_release_at_launch` serves ``x`` (2D BF16, unit column stride).""" + return ( + IS_AVAILABLE + and x.is_cuda + and x.dim() == 2 + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.stride(-1) == 1 + and x.shape[-1] == weight.numel() + and x.data_ptr() % 16 == 0 + and x.stride(0) % 8 == 0 + and x.shape[0] * x.stride(0) < 2**31 + ) + + +def rmsnorm_release_at_launch(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """``flashinfer.norm.rmsnorm(x, weight, eps, enable_pdl=True)``, bit for bit, with the PDL + trigger at launch (see the module docstring). Returns a contiguous BF16 tensor.""" + out = torch.empty(x.shape, dtype=x.dtype, device=x.device) + contiguous = x.is_contiguous() + kernel = _compiled("bfloat16", x.shape[-1], get_sm_version(x.device), contiguous) + kernel(x, weight, out, x.shape[0], eps) + return out + + +__all__ = ["IS_AVAILABLE", "can_use", "rmsnorm_release_at_launch"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/moe/__init__.py new file mode 100644 index 000000000000..e1192a867e74 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Triton kernels of the Kimi K3 MoE layer (the routed-output residual add).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/moe/kimi_k3_moe_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/moe/kimi_k3_moe_residual.py new file mode 100644 index 000000000000..7f9a45e0a876 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/moe/kimi_k3_moe_residual.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Preserve the two BF16 rounding boundaries in a MoE residual update.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _moe_residual_add_kernel( + routed, shared, residual, output, size: tl.constexpr, BLOCK: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < size + routed_value = tl.load(routed + offsets, mask, other=0).to(tl.float32) + shared_value = tl.load(shared + offsets, mask, other=0).to(tl.float32) + residual_value = tl.load(residual + offsets, mask, other=0).to(tl.float32) + moe_output = (routed_value + shared_value).to(tl.bfloat16).to(tl.float32) + tl.store(output + offsets, residual_value + moe_output, mask) + + +def _fused_moe_residual_add( + routed: torch.Tensor, shared: torch.Tensor, residual: torch.Tensor +) -> torch.Tensor: + """Compute residual + BF16(routed + shared), for matching [tokens, hidden] inputs.""" + tensors = (routed, shared, residual) + eligible = ( + routed.ndim == 2 + and 0 < routed.shape[0] <= 96 + and routed.shape[1] == 7168 + and all( + x.shape == routed.shape + and x.device == routed.device + and x.is_cuda + and x.dtype == torch.bfloat16 + and x.is_contiguous() + for x in tensors + ) + and not (torch.is_grad_enabled() and any(x.requires_grad for x in tensors)) + ) + if not eligible: + return residual + (routed + shared) + output = torch.empty_like(routed) + block = 1024 + _moe_residual_add_kernel[(triton.cdiv(routed.numel(), block),)]( + routed, shared, residual, output, routed.numel(), block, num_warps=4 + ) + return output diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/__init__.py new file mode 100644 index 000000000000..a49531debe54 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Triton / CuTe DSL kernels of the Kimi K3 DSpark / DFlash drafter (tail fusion, vocab-parallel +head, fused Markov chain, MLA split plan).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/dspark_markov_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/dspark_markov_fused.py new file mode 100644 index 000000000000..dc6a222424fd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/dspark_markov_fused.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fused DSpark vanilla-Markov greedy block chain. + +Bit-identical to +:func:`tensorrt_llm._torch.models.modeling_speculative.dspark_markov_chain_logits` +driven by a plain full-vocab ``argmax``: step ``i`` adds +``markov_w2 @ markov_w1[prev_i]`` (cast to the logits dtype) to block position +``i``, where ``prev_0`` is the anchor token (an out-of-vocab anchor gives a zero +bias) and ``prev_{i>0}`` is the ``torch.argmax`` of step ``i - 1``'s biased +logits. + +Each step is three launches instead of about a dozen: + +1. ``_markov_prev_embedding_kernel`` -- reduces the previous step's per-tile + argmax partials to ``prev_i`` (or reads the anchor) and gathers the masked + ``markov_w1`` row; +2. the same ``F.linear`` bias GEMV as the reference, so the bias is + bit-identical by construction; +3. ``_markov_add_argmax_kernel`` -- adds the bias to the block logits, stores + the biased row and writes a per-tile argmax. + +Launch configurations depend only on ``(B, vocab)`` and nothing synchronizes +with the host, so the chain is CUDA-graph capturable. +""" + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +_ADD_BLOCK_V = 2048 +_INT32_MAX = tl.constexpr(2**31 - 1) + + +def fused_markov_supported( + base_logits: torch.Tensor, markov_w1: torch.Tensor, markov_w2: torch.Tensor +) -> bool: + """Whether :func:`fused_markov_chain_logits` handles these operands. + + Requires full-vocab fp32 ``[B, K, vocab]`` block logits with a unit vocab + stride and contiguous ``[vocab, rank]`` Markov weights of one dtype. + """ + return ( + base_logits.is_cuda + and base_logits.dim() == 3 + and base_logits.dtype == torch.float32 + and base_logits.stride(-1) == 1 + and markov_w1.dtype == markov_w2.dtype + and markov_w1.is_contiguous() + and markov_w2.is_contiguous() + and markov_w1.shape == markov_w2.shape + and markov_w2.shape[0] == base_logits.shape[-1] + ) + + +@triton.jit +def _argmax_combine(val_a, idx_a, val_b, idx_b): + # torch.argmax order: NaN beats any number, equal values go to the lower index. + nan_a = val_a != val_a + nan_b = val_b != val_b + take_a = tl.where( + nan_a, + (~nan_b) | (idx_a < idx_b), + (val_a > val_b) | ((val_a == val_b) & (idx_a < idx_b)), + ) + return tl.where(take_a, val_a, val_b), tl.where(take_a, idx_a, idx_b) + + +@triton.jit +def _markov_prev_embedding_kernel( + part_val_ptr, + part_idx_ptr, + anchor_ptr, + w1_ptr, + emb_ptr, + vocab, + num_parts, + RANK: tl.constexpr, + FROM_ANCHOR: tl.constexpr, + BLOCK_P: tl.constexpr, + BLOCK_R: tl.constexpr, +): + req = tl.program_id(0) + if FROM_ANCHOR: + prev = tl.load(anchor_ptr + req).to(tl.int64) + else: + offs_p = tl.arange(0, BLOCK_P) + mask_p = offs_p < num_parts + part_val = tl.load( + part_val_ptr + req * num_parts + offs_p, mask=mask_p, other=float("-inf") + ) + part_idx = tl.load(part_idx_ptr + req * num_parts + offs_p, mask=mask_p, other=_INT32_MAX) + _, best_idx = tl.reduce((part_val, part_idx), 0, _argmax_combine) + prev = best_idx.to(tl.int64) + valid = (prev >= 0) & (prev < vocab) + offs_r = tl.arange(0, BLOCK_R) + mask_r = offs_r < RANK + row = tl.load(w1_ptr + tl.where(valid, prev, 0) * RANK + offs_r, mask=mask_r & valid, other=0.0) + tl.store(emb_ptr + req * RANK + offs_r, row, mask=mask_r) + + +@triton.jit +def _markov_add_argmax_kernel( + base_ptr, + out_ptr, + bias_ptr, + stride_base_b, + stride_out_b, + part_val_ptr, + part_idx_ptr, + vocab, + num_parts, + BLOCK_V: tl.constexpr, +): + req = tl.program_id(0) + pid_v = tl.program_id(1) + offs_v = pid_v * BLOCK_V + tl.arange(0, BLOCK_V) + mask_v = offs_v < vocab + row = req.to(tl.int64) + bias = tl.load(bias_ptr + row * vocab + offs_v, mask=mask_v, other=0.0).to(tl.float32) + base = tl.load(base_ptr + row * stride_base_b + offs_v, mask=mask_v, other=0.0) + logits = base + bias + tl.store(out_ptr + row * stride_out_b + offs_v, logits, mask=mask_v) + vals = tl.where(mask_v, logits, float("-inf")) + idx = tl.where(mask_v, offs_v, _INT32_MAX) + best_val, best_idx = tl.reduce((vals, idx), 0, _argmax_combine) + tl.store(part_val_ptr + req * num_parts + pid_v, best_val) + tl.store(part_idx_ptr + req * num_parts + pid_v, best_idx) + + +def fused_markov_chain_logits( + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, +) -> torch.Tensor: + """Greedy-chained vanilla-Markov block logits (fused path). + + Args: + base_logits: ``[B, K, vocab]`` fp32 shared-lm_head block logits (full + vocab, unit vocab stride). + first_prev_tokens: ``[B]`` int32/int64 anchor token ids. + markov_w1: ``[vocab, rank]``, contiguous. + markov_w2: ``[vocab, rank]``, contiguous, same dtype as ``markov_w1``. + Returns: + ``[B, K, vocab]`` fp32 biased logits (a new tensor), bit-identical to + ``dspark_markov_chain_logits`` with a plain argmax. + """ + num_reqs, block_size, vocab = base_logits.shape + rank = markov_w2.shape[1] + device = base_logits.device + out = torch.empty((num_reqs, block_size, vocab), dtype=base_logits.dtype, device=device) + if num_reqs == 0 or block_size == 0: + return out + + num_parts = triton.cdiv(vocab, _ADD_BLOCK_V) + part_val = torch.empty((num_reqs, num_parts), dtype=torch.float32, device=device) + part_idx = torch.empty((num_reqs, num_parts), dtype=torch.int32, device=device) + prev_emb = torch.empty((num_reqs, rank), dtype=markov_w1.dtype, device=device) + for step in range(block_size): + _markov_prev_embedding_kernel[(num_reqs,)]( + part_val, + part_idx, + first_prev_tokens, + markov_w1, + prev_emb, + vocab, + num_parts, + RANK=rank, + FROM_ANCHOR=step == 0, + BLOCK_P=triton.next_power_of_2(num_parts), + BLOCK_R=triton.next_power_of_2(rank), + num_warps=4, + ) + bias = F.linear(prev_emb, markov_w2) + _markov_add_argmax_kernel[(num_reqs, num_parts)]( + base_logits[:, step], + out[:, step], + bias, + base_logits.stride(0), + out.stride(0), + part_val, + part_idx, + vocab, + num_parts, + BLOCK_V=_ADD_BLOCK_V, + num_warps=4, + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/mla_split_plan.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/mla_split_plan.py new file mode 100644 index 000000000000..c5fdb611c5fb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/mla_split_plan.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Split-KV planner of the CuTe DSL MLA decode: its limits and the Triton planning kernel.""" + +import triton +import triton.language as tl + +# Per-request split-KV plan for the CuTe DSL MLA decode kernels. The kernel grid +# carries the tactic's split_kv slots per request; the plan decides how many a +# request uses this step from the batch's KV lengths (the kernel clamps it to +# the grid and the slots beyond exit before their mainloop). +_SPLIT_PLAN_MAX_REQUESTS = 1024 +# CuteDSLNVMlaDecodeBlackwellRunner._MAX_SPLIT_KV: the largest grid any tactic uses. +_SPLIT_PLAN_MAX_SPLIT = 32 +# mma_qk_tiler N of every tactic: the KV tile a split is made of. +_SPLIT_PLAN_KV_TILE = 128 +# Shortest piece worth its partial-output write and reduction read. +_SPLIT_PLAN_MIN_TILES = 4 + + +@triton.jit( + do_not_specialize=["num_requests", "rows_per_request", "wave_clusters", "max_split"], + do_not_specialize_on_alignment=[ + "num_requests", + "rows_per_request", + "wave_clusters", + "max_split", + ], +) +def _mla_split_kv_plan_kernel( + cache_seqs_ptr, + plan_ptr, + num_requests, + rows_per_request, + wave_clusters, + max_split, + KV_TILE: tl.constexpr, + MIN_TILES: tl.constexpr, + PIECE_COST: tl.constexpr, + MAX_PIECES: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.arange(0, BLOCK) + mask = offsets < num_requests + kv_len = tl.load(cache_seqs_ptr + offsets, mask=mask, other=0) + tiles = tl.maximum((kv_len + KV_TILE - 1) // KV_TILE, 1) + tiles = tl.where(mask, tiles, 0) + longest = tl.max(tiles, axis=0) + # Cut the longest request into k equal pieces (every request then takes + # ceil(tiles / piece) pieces) and keep the k whose wave-quantized makespan, + # waves x (piece + fixed piece cost), is smallest; ties go to fewer pieces. + best_piece = longest + best_cost = tl.full([], 2**31 - 1, tl.int32) + for pieces in tl.static_range(1, MAX_PIECES + 1): + piece = tl.maximum((longest + pieces - 1) // pieces, MIN_TILES) + splits = tl.minimum(tl.maximum((tiles + piece - 1) // piece, 1), max_split) + clusters = tl.sum(tl.where(mask, splits, 0), axis=0) * rows_per_request + waves = (clusters + wave_clusters - 1) // wave_clusters + cost = waves * (piece + PIECE_COST) + better = cost < best_cost + best_cost = tl.where(better, cost, best_cost) + best_piece = tl.where(better, piece, best_piece) + plan = tl.minimum(tl.maximum((tiles + best_piece - 1) // best_piece, 1), max_split) + tl.store(plan_ptr + offsets, plan, mask=mask) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/spec_tail_kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/spec_tail_kernels.py new file mode 100644 index 000000000000..09ea8d1bf3ca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/spec_tail_kernels.py @@ -0,0 +1,1585 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Triton kernels for the bookkeeping of a one-model speculative-decoding step. + +Each kernel replaces a chain of small eager PyTorch ops with one launch and +produces the same values: + +* :func:`strict_accept` -- strict (token-equality) acceptance of the previous + draft against the sampled target tokens. +* :func:`decode_draft_block_table` -- the DFlash draft-KV block-table decode + (blocks per request, clamped and divided block offsets). +* :func:`dflash_paged_latent_prepare` -- the per-request bookkeeping of a + DFlash generation step whose drafter context is a paged single-latent cache. +* :func:`store_context_latents` -- the masked store of the context latents of + every drafter layer into its page pool. + +Launch shapes depend only on host batch sizes, so all of them are CUDA-graph +capturable. + +Kimi K3 Rubin (private copy): the kernels below the main ones are Jonas Li's +env-gated tail fusions from ``user/jonasl/k3-rubin-agentic-opt`` (189fe67b81): +``TLLM_SPEC_TAIL_FUSION`` (:func:`fused_strict_accept`) and ``TLLM_DSPARK_TAIL_V2`` +(fused selection + acceptance, the column-parallel block-table decode, the MLA +drafter's context/block preparation and its context latents on a side stream). +Unset, the main kernels above are what runs. +""" + +from typing import NamedTuple, Optional + +import torch +import triton +import triton.language as tl + +from ...features import active as _features + +_VP_INT32_MAX = tl.constexpr(2**31 - 1) + +_BLOCK_TABLE_BLOCK = 2048 + + +@triton.jit +def _strict_accept_kernel( + target_ptr, + draft_ptr, + stride_draft_row, + stride_draft_col, + accepted_ptr, + num_accepted_ptr, + num_ctx, + num_gen, + DRAFT_LEN: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_K: tl.constexpr, +): + width = DRAFT_LEN + 1 + rows = tl.program_id(0) * BLOCK_B + tl.arange(0, BLOCK_B) + cols = tl.arange(0, BLOCK_K) + + # Context requests: accept the sampled token only. + ctx_mask = rows < num_ctx + t_ctx = tl.load(target_ptr + rows, mask=ctx_mask, other=0) + tl.store(accepted_ptr + rows * width, t_ctx.to(tl.int32), mask=ctx_mask) + tl.store(num_accepted_ptr + rows, tl.full([BLOCK_B], 1, tl.int32), mask=ctx_mask) + + # Generation requests: the accepted row is the K+1 target tokens, and the + # count is 1 + the number of leading draft tokens equal to them. + gen_mask = rows < num_gen + tile_mask = gen_mask[:, None] & (cols[None, :] < width) + t_gen = tl.load( + target_ptr + num_ctx + rows[:, None] * width + cols[None, :], mask=tile_mask, other=0 + ).to(tl.int32) + tl.store( + accepted_ptr + (num_ctx + rows[:, None]) * width + cols[None, :], t_gen, mask=tile_mask + ) + draft_mask = gen_mask[:, None] & (cols[None, :] < DRAFT_LEN) + draft = tl.load( + draft_ptr + rows[:, None] * stride_draft_row + cols[None, :] * stride_draft_col, + mask=draft_mask, + other=0, + ).to(tl.int32) + mismatch = draft_mask & (draft != t_gen) + first_mismatch = tl.min(tl.where(mismatch, cols[None, :], DRAFT_LEN), axis=1) + tl.store(num_accepted_ptr + num_ctx + rows, (1 + first_mismatch).to(tl.int32), mask=gen_mask) + + +def strict_accept( + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Strict acceptance of a uniform-width draft in one launch. + + Args: + target_tokens: ``[num_contexts + num_gens * (K + 1)]`` sampled target + tokens (any integer dtype). + draft_tokens: ``[num_gens, K]`` previous draft tokens. + num_contexts, batch_size: request counts. + Returns: + ``accepted_tokens [batch_size, K + 1]`` and ``num_accepted_tokens + [batch_size]``, both int32. Context rows define only column 0. + """ + num_gens = batch_size - num_contexts + draft_len = draft_tokens.shape[-1] + device = target_tokens.device + accepted = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=device) + num_accepted = torch.empty((batch_size,), dtype=torch.int32, device=device) + if batch_size == 0: + return accepted, num_accepted + block_b = min(triton.next_power_of_2(max(num_contexts, num_gens, 1)), 128) + grid = (triton.cdiv(max(num_contexts, num_gens, 1), block_b),) + _strict_accept_kernel[grid]( + target_tokens, + draft_tokens, + draft_tokens.stride(0), + draft_tokens.stride(1), + accepted, + num_accepted, + num_contexts, + num_gens, + DRAFT_LEN=draft_len, + BLOCK_B=block_b, + BLOCK_K=triton.next_power_of_2(draft_len + 1), + num_warps=4, + ) + return accepted, num_accepted + + +@triton.jit +def _decode_block_table_kernel( + src_ptr, + stride_src_row, + stride_src_col, + n_blocks, + divisor, + block_offset, + table_ptr, + stride_table, + counts_ptr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + src_row = src_ptr + row.to(tl.int64) * stride_src_row + table_row = table_ptr + row.to(tl.int64) * stride_table + count = tl.zeros([BLOCK], dtype=tl.int64) + for start in range(0, n_blocks, BLOCK): + cols = start + offs + mask = cols < n_blocks + encoded = tl.load(src_row + cols * stride_src_col, mask=mask, other=-1).to(tl.int64) + encoded -= block_offset + count += (encoded >= 0).to(tl.int64) + tl.store(table_row + cols, (tl.maximum(encoded, 0) // divisor).to(tl.int32), mask=mask) + tl.store(counts_ptr + row, tl.sum(count, axis=0)) + + +def decode_draft_block_table( + src: torch.Tensor, + divisor: int, + table: torch.Tensor, + counts: torch.Tensor, + *, + block_offset: int = 0, +) -> None: + """Decode layer-relative block offsets without modifying ``src``. + + For ``encoded = src - block_offset``, computes + ``counts[i] = (encoded[i] >= 0).sum()`` and + ``table[i] = encoded[i].clamp(min=0) // divisor``. + + Args: + src: ``[rows, blocks]`` encoded block offsets (any strides). + divisor: encoded offset -> pool block index divisor. + table: ``[>= rows, blocks]`` int32 output, unit inner stride. + counts: ``[>= rows]`` int64 output. + block_offset: this layer's K-buffer offset within each physical slot. + """ + rows, n_blocks = src.shape + assert table.shape[-1] == n_blocks and table.stride(-1) == 1 + if rows == 0: + return + _decode_block_table_kernel[(rows,)]( + src, + src.stride(0), + src.stride(1), + n_blocks, + int(divisor), + int(block_offset), + table, + table.stride(0), + counts, + BLOCK=min(_BLOCK_TABLE_BLOCK, triton.next_power_of_2(n_blocks)), + num_warps=4, + ) + + +class DraftBlockPlacement(NamedTuple): + """Where a paged drafter context places each request's own block. + + Block position ``j`` of request ``g`` goes to page ``pages[g, j]`` at + ``offsets[g, j]``, right after the ``num_ctx[g]`` context positions; the + block attends ``seq_lens[g] = num_ctx[g] + block`` keys, repeated per query + token in ``kv_bounds``. + """ + + pages: torch.Tensor # [G, block] int64 + offsets: torch.Tensor # [G, block] int64 + seq_lens: torch.Tensor # [G] int32 + kv_bounds: torch.Tensor # [G * block] int32 + + +class PagedLatentPrepare(NamedTuple): + """What :func:`dflash_paged_latent_prepare` computes for ``G`` generation + requests drafting ``block`` positions from ``width = K + 1`` target tokens.""" + + bonus: torch.Tensor # [G] int64, the last accepted token + query_positions: torch.Tensor # [G, block] int64 + ctx_positions: torch.Tensor # [G * width] int64, RoPE positions of the target latents + store_pages: torch.Tensor # [G * width] int64, destination page of each latent + store_offsets: torch.Tensor # [G * width] int64, offset inside that page + store_mask: torch.Tensor # [G * width] int8, 1 for accepted tokens + num_ctx: torch.Tensor # [G] int64, context length the block attends + blk_pages: torch.Tensor # [G, block] int64, page of each block position + blk_offsets: torch.Tensor # [G, block] int64 + blk_seq_lens: torch.Tensor # [G] int32, num_ctx + block + blk_kv_bounds: torch.Tensor # [G * block] int32, blk_seq_lens per query token + + +@triton.jit +def _paged_latent_prepare_kernel( + num_accepted_ptr, + accepted_ptr, + stride_accepted, + slots_ptr, + ctx_len_ptr, + block_counts_ptr, + table_ptr, + stride_table, + max_blocks, + seq_lens_ptr, + bonus_ptr, + query_pos_ptr, + ctx_pos_ptr, + store_pages_ptr, + store_offsets_ptr, + store_mask_ptr, + num_ctx_ptr, + blk_pages_ptr, + blk_offsets_ptr, + blk_seq_lens_ptr, + blk_kv_bounds_ptr, + num_gens, + first_row, + dummy_slot, + max_ctx, + PAGE_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK_G: tl.constexpr, + BLOCK_W: tl.constexpr, +): + g = tl.arange(0, BLOCK_G) + j = tl.arange(0, BLOCK_W) + gm = g < num_gens + row = first_row + g + nacc = tl.load(num_accepted_ptr + row, mask=gm, other=1).to(tl.int64) + slot = tl.load(slots_ptr + g, mask=gm, other=0).to(tl.int64) + old = tl.load(ctx_len_ptr + slot, mask=gm, other=0).to(tl.int64) + capacity = tl.load(block_counts_ptr + row, mask=gm, other=0).to(tl.int64) * PAGE_SIZE + table_row = table_ptr + row.to(tl.int64) * stride_table + + bonus = tl.load( + accepted_ptr + row.to(tl.int64) * stride_accepted + tl.maximum(nacc - 1, 0), + mask=gm, + other=0, + ) + tl.store(bonus_ptr + g, bonus.to(tl.int64), mask=gm) + + # The block continues from the context length advanced by this step. + now = tl.minimum(old + nacc, max_ctx) + qmask = gm[:, None] & (j[None, :] < BLOCK_SIZE) + blk = g[:, None] * BLOCK_SIZE + j[None, :] + tl.store(query_pos_ptr + blk, now[:, None] + j[None, :], mask=qmask) + + # Context latents of this step's target tokens: RoPE position, destination + # (column clamped into the request's allocation) and write mask. + wmask = gm[:, None] & (j[None, :] < WIDTH) + flat = g[:, None] * WIDTH + j[None, :] + pos = old[:, None] + j[None, :] + tl.store(ctx_pos_ptr + flat, pos, mask=wmask) + col = tl.maximum(tl.minimum(pos, capacity[:, None] - 1), 0) + page = tl.load(table_row[:, None] + tl.minimum(col // PAGE_SIZE, max_blocks - 1), mask=wmask) + tl.store(store_pages_ptr + flat, page.to(tl.int64), mask=wmask) + tl.store(store_offsets_ptr + flat, col % PAGE_SIZE, mask=wmask) + tl.store(store_mask_ptr + flat, (j[None, :] < nacc[:, None]).to(tl.int8), mask=wmask) + + # Advance the context length; the shared dummy slot stays empty. + new_len = tl.where(slot == dummy_slot, 0, now) + # Every lane has read its old length before any lane writes a new one. + tl.debug_barrier() + tl.store(ctx_len_ptr + slot, new_len, mask=gm) + tl.store(ctx_len_ptr + dummy_slot, tl.zeros([], dtype=tl.int64)) + + # Advertised context, capped so the block still fits in the allocation. + num_ctx = tl.minimum(new_len, tl.maximum(capacity - BLOCK_SIZE, 0)) + tl.store(num_ctx_ptr + g, num_ctx, mask=gm) + tl.store( + seq_lens_ptr + row, + tl.full([BLOCK_G], WIDTH, tl.int32).to(seq_lens_ptr.dtype.element_ty), + mask=gm, + ) + + # Block-decode operands: where the block's own latents go, and its bounds. + bpos = num_ctx[:, None] + j[None, :] + bpage = tl.load(table_row[:, None] + tl.minimum(bpos // PAGE_SIZE, max_blocks - 1), mask=qmask) + tl.store(blk_pages_ptr + blk, bpage.to(tl.int64), mask=qmask) + tl.store(blk_offsets_ptr + blk, bpos % PAGE_SIZE, mask=qmask) + seq_len = (num_ctx + BLOCK_SIZE).to(tl.int32) + tl.store(blk_seq_lens_ptr + g, seq_len, mask=gm) + tl.store( + blk_kv_bounds_ptr + blk, + tl.zeros([BLOCK_G, BLOCK_W], tl.int32) + seq_len[:, None], + mask=qmask, + ) + + +def dflash_paged_latent_prepare( + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + first_row: int, + dummy_slot: int, + max_ctx: int, + page_size: int, + block_size: int, + width: int, +) -> PagedLatentPrepare: + """The DFlash generation-step bookkeeping for a paged single-latent context. + + For generation request ``g`` (batch row ``first_row + g``, context slot + ``slots[g]``) computes what ``K3DFlashWorker.prepare_1st_drafter_inputs`` and + ``_build_mla_block_fixup`` compute eagerly, and applies the same side + effects: ``ctx_len`` advanced in place (the dummy slot zeroed) and + ``seq_lens_cuda`` of the generation rows set to ``width``. + + Args: + num_accepted_tokens: ``[batch]`` accepted counts. + accepted_tokens: ``[batch, width]`` accepted tokens, unit inner stride. + slots: ``[G]`` context slots. + ctx_len: ``[slots]`` int64 context lengths. + block_counts: ``[batch]`` pages held per batch row. + block_tables: ``[batch, max_blocks]`` int32 page tables, unit inner stride. + seq_lens_cuda: ``[batch]`` attention sequence lengths. + first_row: batch row of generation request 0. + dummy_slot: the shared padding slot. + max_ctx: context length ceiling. + page_size, block_size, width: page size, drafted block positions and + target tokens per request. + """ + num_gens = slots.shape[0] + device = ctx_len.device + i64 = dict(dtype=torch.int64, device=device) + out = PagedLatentPrepare( + bonus=torch.empty((num_gens,), **i64), + query_positions=torch.empty((num_gens, block_size), **i64), + ctx_positions=torch.empty((num_gens * width,), **i64), + store_pages=torch.empty((num_gens * width,), **i64), + store_offsets=torch.empty((num_gens * width,), **i64), + store_mask=torch.empty((num_gens * width,), dtype=torch.int8, device=device), + num_ctx=torch.empty((num_gens,), **i64), + blk_pages=torch.empty((num_gens, block_size), **i64), + blk_offsets=torch.empty((num_gens, block_size), **i64), + blk_seq_lens=torch.empty((num_gens,), dtype=torch.int32, device=device), + blk_kv_bounds=torch.empty((num_gens * block_size,), dtype=torch.int32, device=device), + ) + if num_gens == 0: + return out + assert accepted_tokens.stride(-1) == 1 and block_tables.stride(-1) == 1 + _paged_latent_prepare_kernel[(1,)]( + num_accepted_tokens, + accepted_tokens, + accepted_tokens.stride(0), + slots, + ctx_len, + block_counts, + block_tables, + block_tables.stride(0), + block_tables.shape[1], + seq_lens_cuda, + *out, + num_gens, + first_row, + dummy_slot, + max_ctx, + PAGE_SIZE=page_size, + BLOCK_SIZE=block_size, + WIDTH=width, + BLOCK_G=triton.next_power_of_2(num_gens), + BLOCK_W=triton.next_power_of_2(max(width, block_size)), + num_warps=4, + ) + return out + + +@triton.jit +def _store_context_latents_kernel( + latent_ptr, + stride_latent_tok, + stride_latent_layer, + pool_ptrs_ptr, + stride_pool_page, + stride_pool_off, + pages_ptr, + offsets_ptr, + mask_ptr, + head_dim, + BLOCK_D: tl.constexpr, +): + tok = tl.program_id(0) + layer = tl.program_id(1) + offs = tl.arange(0, BLOCK_D) + dmask = offs < head_dim + page = tl.load(pages_ptr + tok) + off = tl.load(offsets_ptr + tok) + keep = tl.load(mask_ptr + tok).to(tl.float32) + src = latent_ptr + tok.to(tl.int64) * stride_latent_tok + layer * stride_latent_layer + vals = tl.load(src + offs, mask=dmask, other=0.0) + # The eager path multiplies by the mask in the latent dtype: x * 1 is x and + # x * 0 a signed zero either way. + vals = (vals.to(tl.float32) * keep).to(vals.dtype) + pool = tl.load(pool_ptrs_ptr + layer).to(tl.pointer_type(vals.dtype)) + tl.store(pool + page * stride_pool_page + off * stride_pool_off + offs, vals, mask=dmask) + + +def store_context_latents( + latent: torch.Tensor, + pools: list[torch.Tensor], + pool_ptrs: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + mask: torch.Tensor, +) -> None: + """``pools[l][pages[t], 0, 0, offsets[t]] = latent[t, l] * mask[t]`` for every + token ``t`` and layer ``l``. + + Args: + latent: ``[N, L, D]`` per-token, per-layer latents (unit inner stride). + pools: the ``L`` per-layer ``[pages, 1, 1, page_size, D]`` pool views, + all with the same strides. + pool_ptrs: ``[L]`` int64 base addresses of ``pools``. + pages, offsets: ``[N]`` destination page and offset in the page. + mask: ``[N]`` int8 write mask; 0 stores a zeroed latent. + """ + num_tokens, num_layers, head_dim = latent.shape + if num_tokens == 0: + return + pool0 = pools[0] + assert latent.stride(-1) == 1 and pool0.stride(-1) == 1 + assert all(p.stride() == pool0.stride() and p.dtype == latent.dtype for p in pools) + _store_context_latents_kernel[(num_tokens, num_layers)]( + latent, + latent.stride(0), + latent.stride(1), + pool_ptrs, + pool0.stride(0), + pool0.stride(3), + pages, + offsets, + mask, + head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + num_warps=4, + ) + + +# -------------------------------------------------------------------------------------------------- +# Jonas Li's env-gated tail fusions (k3-rubin-agentic-opt). +# -------------------------------------------------------------------------------------------------- + +# Largest number of generation requests the one-program fused tail kernels +# (selection + acceptance, the Markov steps) take; larger batches keep the +# unfused, row-parallel kernels. Measured in the decode graph at 1-4 requests. +TAIL_V2_MAX_GENS = 16 + +# Columns per program of the parallel draft block-table decode. +BLOCK_TABLE_CHUNK = _BLOCK_TABLE_BLOCK + + +def spec_tail_fusion_enabled() -> bool: + """Whether ``TLLM_SPEC_TAIL_FUSION=1`` selects the fused tail kernels.""" + return _features().spec_tail_fusion + + +def tail_v2_enabled() -> bool: + """Whether ``TLLM_DSPARK_TAIL_V2=1`` (with ``TLLM_SPEC_TAIL_FUSION=1``) selects + the second set of fused tail kernels. They write bitwise the values of the + paths they replace: the MLA drafter's context and block RoPE kernels with + in-kernel table lookups and the block latents stored straight into the + paged pool, the column-parallel draft block-table decode, shape-only + operands built once instead of every step, and a TP-sharded drafter's + context latents computed on a side stream, concurrently with the + target-token selection and acceptance.""" + return spec_tail_fusion_enabled() and _features().dspark_tail_v2 + + +@triton.jit +def _fused_strict_accept_kernel( + target_ptr, + draft_ptr, + stride_draft, + accepted_ptr, + num_accepted_ptr, + num_ctx, + num_gen, + pool_ptr, + counter_ptr, + base_total, + frac, + max_total, + DRAFT_LEN: tl.constexpr, + FORCE: tl.constexpr, + COUNTER_STRIDE: tl.constexpr, + SLOT_STRIDE: tl.constexpr, + POOL_MASK: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_K: tl.constexpr, +): + width = DRAFT_LEN + 1 + rows = tl.arange(0, BLOCK_B) + cols = tl.arange(0, BLOCK_K) + + # Context requests: accept the sampled token only. + ctx_mask = rows < num_ctx + t_ctx = tl.load(target_ptr + rows, mask=ctx_mask, other=0) + tl.store(accepted_ptr + rows * width, t_ctx.to(tl.int32), mask=ctx_mask) + tl.store(num_accepted_ptr + rows, tl.full([BLOCK_B], 1, tl.int32), mask=ctx_mask) + + # Generation requests: accepted row = the K+1 target tokens; the accepted + # length is 1 + the number of leading draft tokens equal to them. + gen_mask = rows < num_gen + tile_mask = gen_mask[:, None] & (cols[None, :] < width) + t_gen = tl.load( + target_ptr + num_ctx + rows[:, None] * width + cols[None, :], mask=tile_mask, other=0 + ).to(tl.int32) + tl.store( + accepted_ptr + (num_ctx + rows[:, None]) * width + cols[None, :], t_gen, mask=tile_mask + ) + draft_mask = gen_mask[:, None] & (cols[None, :] < DRAFT_LEN) + d = tl.load( + draft_ptr + rows[:, None] * stride_draft + cols[None, :], mask=draft_mask, other=0 + ).to(tl.int32) + mismatch = draft_mask & (d != t_gen) + first_mismatch = tl.min(tl.where(mismatch, cols[None, :], DRAFT_LEN), axis=1) + n = 1 + first_mismatch + if FORCE == 1: + n = tl.full([BLOCK_B], 0, tl.int32) + base_total + elif FORCE == 2: + # Same draw as SpecWorkerBase._apply_force_accepted_tokens: bump the + # device counter, hash (counter, slot) into the fixed uniform pool. + counter = tl.load(counter_ptr) + 1 + # Every warp reads the old value before any warp writes the new one. + tl.debug_barrier() + tl.store(counter_ptr, counter) + idx = (counter * COUNTER_STRIDE + rows.to(tl.int64) * SLOT_STRIDE) & POOL_MASK + rand = tl.load(pool_ptr + idx, mask=gen_mask, other=1.0) + extra = (rand < frac).to(tl.int32) + n = tl.minimum(base_total + extra, max_total) + tl.store(num_accepted_ptr + num_ctx + rows, n.to(tl.int32), mask=gen_mask) + + +def fused_strict_accept( + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Strict acceptance in one launch. + + Args: + target_tokens: ``[num_contexts + num_gens * (K + 1)]`` sampled target + tokens (any integer dtype). + draft_tokens: ``[num_gens, K]`` previous draft tokens (unit inner stride). + num_contexts, batch_size: request counts. + force: ``None`` for no override, ``(base_total,)`` for a constant + accepted length, or ``(base_total, frac, max_total, pool, counter, + counter_stride, slot_stride)`` for the fractional override. + Returns: + ``accepted_tokens [batch_size, K + 1]`` int32 and + ``num_accepted_tokens [batch_size]`` int32, matching + ``SpecWorkerBase._sample_and_accept_draft_tokens_base`` (context rows + only define column 0, as there). + """ + num_gens = batch_size - num_contexts + draft_len = draft_tokens.shape[-1] + device = target_tokens.device + accepted = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=device) + num_accepted = torch.empty((batch_size,), dtype=torch.int32, device=device) + if batch_size == 0: + return accepted, num_accepted + assert draft_tokens.dim() == 2 and (num_gens == 0 or draft_tokens.stride(-1) == 1) + force_mode, base_total, frac, max_total = 0, 0, 0.0, 0 + pool, counter, counter_stride, slot_stride, pool_mask = target_tokens, target_tokens, 1, 1, 0 + if force is not None and len(force) == 1: + force_mode, base_total = 1, int(force[0]) + elif force is not None: + base_total, frac, max_total, pool, counter, counter_stride, slot_stride = force + force_mode = 2 + pool_mask = pool.numel() - 1 + _fused_strict_accept_kernel[(1,)]( + target_tokens, + draft_tokens, + draft_tokens.stride(0) if num_gens > 0 else draft_len, + accepted, + num_accepted, + num_contexts, + num_gens, + pool, + counter, + base_total, + float(frac), + max_total, + DRAFT_LEN=draft_len, + FORCE=force_mode, + COUNTER_STRIDE=counter_stride, + SLOT_STRIDE=slot_stride, + POOL_MASK=pool_mask, + BLOCK_B=triton.next_power_of_2(max(num_contexts, num_gens, 1)), + BLOCK_K=triton.next_power_of_2(draft_len + 1), + num_warps=4, + ) + return accepted, num_accepted + + +@triton.jit +def _pair_argmax(vals, idx, axis: tl.constexpr): + # The index the vocab-parallel _argmax_combine reduction selects, from + # native reductions: NaN beats every number and the lowest index wins among + # NaNs; otherwise the largest value wins and the lowest index among equal + # values (+0.0 == -0.0). The selected index is unique, so it does not depend + # on the reduction order. + is_nan = vals != vals + any_nan = tl.max(is_nan.to(tl.int32), axis=axis) > 0 + vmax = tl.max(tl.where(is_nan, float("-inf"), vals), axis=axis) + cand = tl.where( + tl.expand_dims(any_nan, axis), is_nan, (vals == tl.expand_dims(vmax, axis)) & ~is_nan + ) + return tl.min(tl.where(cand, idx, _VP_INT32_MAX), axis=axis) + + +@triton.jit +def _select_accept_kernel( + pairs_ptr, + stride_row, + n_slots, + draft_ptr, + stride_draft, + accepted_ptr, + num_accepted_ptr, + num_ctx, + num_gen, + pool_ptr, + counter_ptr, + base_total, + frac, + max_total, + prev_ptr, + stride_prev, + state_idx_ptr, + dummy_ptr, + kv_lens_ptr, + rewind_ptr, + DRAFT_LEN: tl.constexpr, + FORCE: tl.constexpr, + COUNTER_STRIDE: tl.constexpr, + SLOT_STRIDE: tl.constexpr, + POOL_MASK: tl.constexpr, + KDA_REPLAY: tl.constexpr, + KV_REWIND: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_SLOTS: tl.constexpr, +): + # Pairs are (value, index) float32 couples, slot s of a row at columns + # [2 s, 2 s + 1]; a row is loaded as one contiguous run and split. + width = DRAFT_LEN + 1 + rows = tl.arange(0, BLOCK_B) + cols = tl.arange(0, BLOCK_K) + gen_mask = rows < num_gen + tile_mask = gen_mask[:, None] & (cols[None, :] < width) + draft_mask = gen_mask[:, None] & (cols[None, :] < DRAFT_LEN) + + # Everything that does not depend on the selection is read first. + d = tl.load( + draft_ptr + rows[:, None] * stride_draft + cols[None, :], mask=draft_mask, other=0 + ).to(tl.int32) + if KDA_REPLAY: + slot = tl.load(state_idx_ptr + num_ctx + rows, mask=gen_mask, other=0).to(tl.int64) + dummy = tl.load(dummy_ptr + num_ctx + rows, mask=gen_mask, other=0) != 0 + current = tl.load(prev_ptr + slot * stride_prev, mask=gen_mask, other=0) + if KV_REWIND: + kv = tl.load(kv_lens_ptr + num_ctx + rows, mask=gen_mask) + if FORCE == 2: + counter = tl.load(counter_ptr) + 1 + idx = (counter * COUNTER_STRIDE + rows.to(tl.int64) * SLOT_STRIDE) & POOL_MASK + rand = tl.load(pool_ptr + idx, mask=gen_mask, other=1.0) + # Every warp reads the old value before any warp writes the new one. + tl.debug_barrier() + tl.store(counter_ptr, counter) + + lanes = tl.arange(0, 2 * BLOCK_SLOTS) + lane_ok = lanes < 2 * n_slots + + # Context requests: accept the selected token only. + c = tl.arange(0, BLOCK_C) + cmask = c < num_ctx + cm2 = cmask[:, None] & lane_ok[None, :] + craw = tl.load( + pairs_ptr + c[:, None].to(tl.int64) * stride_row + lanes[None, :], mask=cm2, other=0.0 + ) + cvals, cidxf = tl.split(tl.reshape(craw, [BLOCK_C, BLOCK_SLOTS, 2])) + cok = cmask[:, None] & (tl.arange(0, BLOCK_SLOTS)[None, :] < n_slots) + t_ctx = _pair_argmax( + tl.where(cok, cvals, float("-inf")), tl.where(cok, cidxf.to(tl.int32), _VP_INT32_MAX), 1 + ) + tl.store(accepted_ptr + c * width, t_ctx, mask=cmask) + tl.store(num_accepted_ptr + c, tl.full([BLOCK_C], 1, tl.int32), mask=cmask) + + # Generation requests: the K+1 selected tokens of each request (the global + # argmax / draw of every row's exchanged pairs, as the vocab-parallel + # selection reduces them), then strict acceptance as _strict_accept_kernel. + prow = num_ctx + rows[:, None] * width + cols[None, :] + m3 = tile_mask[:, :, None] & lane_ok[None, None, :] + graw = tl.load( + pairs_ptr + prow[:, :, None].to(tl.int64) * stride_row + lanes[None, None, :], + mask=m3, + other=0.0, + ) + gvals, gidxf = tl.split(tl.reshape(graw, [BLOCK_B, BLOCK_K, BLOCK_SLOTS, 2])) + gok = tile_mask[:, :, None] & (tl.arange(0, BLOCK_SLOTS)[None, None, :] < n_slots) + t_gen = _pair_argmax( + tl.where(gok, gvals, float("-inf")), tl.where(gok, gidxf.to(tl.int32), _VP_INT32_MAX), 2 + ) + tl.store( + accepted_ptr + (num_ctx + rows[:, None]) * width + cols[None, :], t_gen, mask=tile_mask + ) + mismatch = draft_mask & (d != t_gen) + first_mismatch = tl.min(tl.where(mismatch, cols[None, :], DRAFT_LEN), axis=1) + n = 1 + first_mismatch + if FORCE == 1: + n = tl.full([BLOCK_B], 0, tl.int32) + base_total + elif FORCE == 2: + extra = (rand < frac).to(tl.int32) + n = tl.minimum(base_total + extra, max_total) + n = n.to(tl.int32) + tl.store(num_accepted_ptr + num_ctx + rows, n, mask=gen_mask) + + if KDA_REPLAY: + # KDAReplayState.record_acceptance: record the accepted draft prefix per + # state slot; dummy requests keep the slot's current value. All slots + # were read before any is written, as gather-then-put does. + accepted_drafts = tl.maximum(n - 1, 0).to(current.dtype) + tl.debug_barrier() + tl.store( + prev_ptr + slot * stride_prev, + tl.where(dummy, current, accepted_drafts), + mask=gen_mask, + ) + if KV_REWIND: + # K3DFlashWorker._prepare_kv_for_draft_forward. + tl.store(rewind_ptr + rows, 1 - n, mask=gen_mask) + tl.store(kv_lens_ptr + num_ctx + rows, kv + 1, mask=gen_mask) + + +def fused_select_accept( + pairs: tuple, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, + kda_replay: Optional[tuple] = None, + kv_lens: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Vocab-parallel target selection, strict acceptance and the per-step + bookkeeping that follows it, in one launch. + + Args: + pairs: ``(pairs, row stride, slot stride, slots)`` from + ``K3VocabShardedLogits.select_pairs``: every rank's ``(value, global + index)`` candidates of the ``num_contexts + num_gens * (K + 1)`` rows. + draft_tokens, num_contexts, batch_size, force: as :func:`fused_strict_accept`. + kda_replay: ``(prev_num_accepted_tokens, state_indices, dummy_mask)`` to + also apply ``KDAReplayState.record_acceptance``, or None. + ``prev_num_accepted_tokens`` is a 1-D int32 view with any slot stride + (the replay count is the first int32 of each slot's record); + ``state_indices`` and ``dummy_mask`` cover the whole batch, contexts first. + kv_lens: ``attn_metadata.kv_lens_cuda`` to also apply + ``K3DFlashWorker._prepare_kv_for_draft_forward`` (its gen rows += 1), or None. + Returns: + ``accepted_tokens [batch_size, K + 1]`` int32, ``num_accepted_tokens + [batch_size]`` int32 (the values of the reduce + :func:`fused_strict_accept` + path) and, with ``kv_lens``, the ``[num_gens]`` int32 KV rewind amounts + ``1 - num_accepted``. + """ + pair_buf, stride_row, stride_slot, n_slots = pairs + num_gens = batch_size - num_contexts + draft_len = draft_tokens.shape[-1] + device = pair_buf.device + accepted = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=device) + num_accepted = torch.empty((batch_size,), dtype=torch.int32, device=device) + rewind = ( + torch.empty((num_gens,), dtype=torch.int32, device=device) if kv_lens is not None else None + ) + if batch_size == 0: + return accepted, num_accepted, rewind + assert draft_tokens.dim() == 2 and (num_gens == 0 or draft_tokens.stride(-1) == 1) + force_mode, base_total, frac, max_total = 0, 0, 0.0, 0 + pool, counter, counter_stride, slot_stride, pool_mask = pair_buf, pair_buf, 1, 1, 0 + if force is not None and len(force) == 1: + force_mode, base_total = 1, int(force[0]) + elif force is not None: + base_total, frac, max_total, pool, counter, counter_stride, slot_stride = force + force_mode = 2 + pool_mask = pool.numel() - 1 + prev, state_idx, dummy = kda_replay if kda_replay is not None else (accepted,) * 3 + if kda_replay is not None: + assert prev.dim() == 1 and prev.dtype == torch.int32 + assert state_idx.stride(-1) == 1 and dummy.stride(-1) == 1 + # Interleaved (value, index) slots, as the MNNVL / single-rank exchange lays them out. + assert stride_slot == 2 and pair_buf.stride(-1) == 1 + _select_accept_kernel[(1,)]( + pair_buf, + stride_row, + n_slots, + draft_tokens, + draft_tokens.stride(0) if num_gens > 0 else draft_len, + accepted, + num_accepted, + num_contexts, + num_gens, + pool, + counter, + base_total, + float(frac), + max_total, + prev, + prev.stride(0), + state_idx, + dummy, + kv_lens if kv_lens is not None else accepted, + rewind if rewind is not None else accepted, + DRAFT_LEN=draft_len, + FORCE=force_mode, + COUNTER_STRIDE=counter_stride, + SLOT_STRIDE=slot_stride, + POOL_MASK=pool_mask, + KDA_REPLAY=kda_replay is not None, + KV_REWIND=kv_lens is not None, + BLOCK_C=triton.next_power_of_2(max(num_contexts, 1)), + BLOCK_B=triton.next_power_of_2(max(num_gens, 1)), + BLOCK_K=triton.next_power_of_2(draft_len + 1), + BLOCK_SLOTS=triton.next_power_of_2(n_slots), + num_warps=4, + ) + return accepted, num_accepted, rewind + + +@triton.jit +def _decode_block_table_chunk_kernel( + src_ptr, + stride_src, + stride_src_col, + n_blocks, + divisor, + block_offset, + table_ptr, + stride_table, + partial_ptr, + BLOCK: tl.constexpr, +): + # Program (row, chunk) decodes columns [chunk * BLOCK, (chunk + 1) * BLOCK) + # and records how many of them are allocated (encoded >= 0). + row = tl.program_id(0) + chunk = tl.program_id(1) + cols = chunk * BLOCK + tl.arange(0, BLOCK) + mask = cols < n_blocks + src_row = src_ptr + row.to(tl.int64) * stride_src + encoded = tl.load(src_row + cols * stride_src_col, mask=mask, other=-1).to(tl.int64) + encoded -= block_offset + decoded = tl.maximum(encoded, 0) // divisor + tl.store(table_ptr + row.to(tl.int64) * stride_table + cols, decoded.to(tl.int32), mask=mask) + count = tl.sum((encoded >= 0).to(tl.int32), axis=0) + tl.store(partial_ptr + row * tl.num_programs(1) + chunk, count) + + +@triton.jit +def _sum_block_counts_kernel( + partial_ptr, n_chunks, counts_ptr, rows, BLOCK_R: tl.constexpr, BLOCK_C: tl.constexpr +): + r = tl.arange(0, BLOCK_R) + c = tl.arange(0, BLOCK_C) + mask = (r[:, None] < rows) & (c[None, :] < n_chunks) + partial = tl.load(partial_ptr + r[:, None] * n_chunks + c[None, :], mask=mask, other=0) + tl.store(counts_ptr + r, tl.sum(partial.to(tl.int64), axis=1), mask=r < rows) + + +def decode_draft_block_table_parallel( + src: torch.Tensor, + divisor: int, + table: torch.Tensor, + counts: torch.Tensor, + partial: torch.Tensor, + *, + block_offset: int = 0, +) -> None: + """:func:`decode_draft_block_table` with the columns of a row spread over programs. + + Same outputs (integer arithmetic, order-free). ``partial`` is an int32 + scratch of at least ``rows * cdiv(blocks, _BLOCK_TABLE_BLOCK)`` elements. + """ + rows, n_blocks = src.shape + assert table.shape[-1] == n_blocks and table.stride(-1) == 1 + if rows == 0: + return + n_chunks = triton.cdiv(n_blocks, _BLOCK_TABLE_BLOCK) + assert partial.dtype == torch.int32 and partial.numel() >= rows * n_chunks + _decode_block_table_chunk_kernel[(rows, n_chunks)]( + src, + src.stride(0), + src.stride(1), + n_blocks, + int(divisor), + int(block_offset), + table, + table.stride(0), + partial, + BLOCK=_BLOCK_TABLE_BLOCK, + num_warps=8, + ) + _sum_block_counts_kernel[(1,)]( + partial, + n_chunks, + counts, + rows, + BLOCK_R=triton.next_power_of_2(rows), + BLOCK_C=triton.next_power_of_2(n_chunks), + num_warps=4, + ) + + +@triton.jit +def _ctx_positions_kernel( + ctx_len_ptr, slots_ptr, out_ptr, num_gens, WIDTH: tl.constexpr, BLOCK: tl.constexpr +): + i = tl.arange(0, BLOCK) + mask = i < num_gens * WIDTH + g = i // WIDTH + slot = tl.load(slots_ptr + g, mask=mask, other=0).to(tl.int64) + old = tl.load(ctx_len_ptr + slot, mask=mask, other=0).to(tl.int64) + tl.store(out_ptr + i, old + (i - g * WIDTH), mask=mask) + + +def dflash_ctx_positions(ctx_len: torch.Tensor, slots: torch.Tensor, width: int) -> torch.Tensor: + """``ctx_len[slots[g]] + j`` for the ``width`` tokens ``j`` of each request + ``g``: the RoPE positions :func:`dflash_gen_prepare` assigns the step's + context latents (it reads ``ctx_len`` before advancing it).""" + num_gens = slots.shape[0] + out = torch.empty((num_gens * width,), dtype=torch.long, device=ctx_len.device) + if num_gens > 0: + _ctx_positions_kernel[(1,)]( + ctx_len, + slots, + out, + num_gens, + WIDTH=width, + BLOCK=triton.next_power_of_2(num_gens * width), + num_warps=1, + ) + return out + + +@triton.jit +def _dflash_gen_prepare_kernel( + num_accepted_ptr, + accepted_ptr, + stride_accepted, + slots_ptr, + ctx_len_ptr, + block_counts_ptr, + table_ptr, + stride_table, + max_blocks, + seq_lens_ptr, + bonus_ptr, + embed_ids_ptr, + query_pos_ptr, + ctx_pos_ptr, + store_pages_ptr, + store_offsets_ptr, + store_mask_ptr, + num_ctx_ptr, + blk_pages_ptr, + blk_offsets_ptr, + blk_seq_lens_ptr, + blk_kv_bounds_ptr, + blk_owned_ptr, + num_gens, + num_ctx_reqs, + dummy_slot, + max_ctx, + mask_token_id, + cp_rank, + PAGE_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + CP_SIZE: tl.constexpr, + BLOCK_G: tl.constexpr, + BLOCK_W: tl.constexpr, +): + g = tl.arange(0, BLOCK_G) + j = tl.arange(0, BLOCK_W) + gm = g < num_gens + row = num_ctx_reqs + g + nacc = tl.load(num_accepted_ptr + row, mask=gm, other=1).to(tl.int64) + slot = tl.load(slots_ptr + g, mask=gm, other=0).to(tl.int64) + old = tl.load(ctx_len_ptr + slot, mask=gm, other=0).to(tl.int64) + counts = tl.load(block_counts_ptr + row, mask=gm, other=0).to(tl.int64) + table_row = table_ptr + row.to(tl.int64) * stride_table + + # Bonus = last accepted token; the embedding input appends the mask token. + bidx = tl.maximum(nacc - 1, 0) + bonus = tl.load(accepted_ptr + row.to(tl.int64) * stride_accepted + bidx, mask=gm, other=0) + bonus = bonus.to(tl.int64) + tl.store(bonus_ptr + g, bonus, mask=gm) + tl.store(embed_ids_ptr + g, bonus, mask=gm) + tl.store(embed_ids_ptr + num_gens, mask_token_id.to(tl.int64)) + + # Block query positions continue from the advanced context length. + now = tl.minimum(old + nacc, max_ctx) + qmask = gm[:, None] & (j[None, :] < BLOCK_SIZE) + tl.store( + query_pos_ptr + g[:, None] * BLOCK_SIZE + j[None, :], now[:, None] + j[None, :], mask=qmask + ) + + # Context-latent store of this step's K+1 target tokens: RoPE positions, + # destination page / offset (column clamped into the allocation) and mask. + # Under helix (CP_SIZE > 1) global page p lives on rank p % CP_SIZE as its + # local page p // CP_SIZE; positions on other ranks' pages are skipped (-1). + wmask = gm[:, None] & (j[None, :] < WIDTH) + flat = g[:, None] * WIDTH + j[None, :] + pos = old[:, None] + j[None, :] + tl.store(ctx_pos_ptr + flat, pos, mask=wmask) + capacity = counts * PAGE_SIZE + gpage = pos // PAGE_SIZE + col = (gpage // CP_SIZE) * PAGE_SIZE + pos % PAGE_SIZE + col = tl.maximum(tl.minimum(col, capacity[:, None] - 1), 0) + page_idx = tl.minimum(col // PAGE_SIZE, max_blocks - 1) + page = tl.load(table_row[:, None] + page_idx, mask=wmask, other=0).to(tl.int64) + tl.store(store_pages_ptr + flat, page, mask=wmask) + tl.store(store_offsets_ptr + flat, col % PAGE_SIZE, mask=wmask) + store = (j[None, :] < nacc[:, None]).to(tl.int8) + store = tl.where(gpage % CP_SIZE == cp_rank, store, -1) + tl.store(store_mask_ptr + flat, store, mask=wmask) + + # Advance the context length; the shared padding / dummy slot stays empty. + new_len = tl.where(slot == dummy_slot, 0, now) + tl.debug_barrier() + tl.store(ctx_len_ptr + slot, new_len, mask=gm) + tl.store(ctx_len_ptr + dummy_slot, tl.zeros([], dtype=tl.int64)) + + # Advertised context, capped so the block still fits in the allocation + # (CP_SIZE local pages cover CP_SIZE times as many global positions). + allocated = tl.maximum(capacity * CP_SIZE - BLOCK_SIZE, 0) + num_ctx = tl.minimum(new_len, allocated) + tl.store(num_ctx_ptr + g, num_ctx, mask=gm) + seq_lens = tl.full([BLOCK_G], WIDTH, tl.int32).to(seq_lens_ptr.dtype.element_ty) + tl.store(seq_lens_ptr + row, seq_lens, mask=gm) + + # Block-decode operands: where the block's own latents go, and its bounds + # (rank-local slots and lengths under helix, as HelixDraftShard computes). + bpos = num_ctx[:, None] + j[None, :] + bgpage = bpos // PAGE_SIZE + bslot = (bgpage // CP_SIZE) * PAGE_SIZE + bpos % PAGE_SIZE + bpage = tl.load( + table_row[:, None] + tl.minimum(bslot // PAGE_SIZE, max_blocks - 1), mask=qmask, other=0 + ) + tl.store(blk_pages_ptr + g[:, None] * BLOCK_SIZE + j[None, :], bpage.to(tl.int64), mask=qmask) + tl.store(blk_offsets_ptr + g[:, None] * BLOCK_SIZE + j[None, :], bslot % PAGE_SIZE, mask=qmask) + total = num_ctx + BLOCK_SIZE + if CP_SIZE > 1: + tl.store( + blk_owned_ptr + g[:, None] * BLOCK_SIZE + j[None, :], + (bgpage % CP_SIZE == cp_rank).to(tl.int8), + mask=qmask, + ) + span = PAGE_SIZE * CP_SIZE + rem = total % span - cp_rank * PAGE_SIZE + total = (total // span) * PAGE_SIZE + tl.minimum(tl.maximum(rem, 0), PAGE_SIZE) + seq_len = total.to(tl.int32) + tl.store(blk_seq_lens_ptr + g, seq_len, mask=gm) + tl.store( + blk_kv_bounds_ptr + g[:, None] * BLOCK_SIZE + j[None, :], + tl.zeros([BLOCK_G, BLOCK_W], dtype=tl.int32) + seq_len[:, None], + mask=qmask, + ) + + +def dflash_gen_prepare( + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + num_contexts: int, + num_gens: int, + dummy_slot: int, + max_ctx: int, + mask_token_id: int, + page_size: int, + block_size: int, + width: int, + cp_size: int = 1, + cp_rank: int = 0, +) -> dict: + """The DFlash paged-context bookkeeping of one generation step in one launch. + + Computes, for each generation request ``g`` (batch row ``num_contexts + g``), + exactly what ``K3DFlashWorker.prepare_1st_drafter_inputs`` and + ``_build_mla_block_fixup`` compute eagerly when the drafter context lives in + the manager's paged pool: the bonus token, the block query positions, the + RoPE positions / destination pages / write mask of the ``width`` target + latents, the advanced ``ctx_len`` (written in place; the dummy slot is + zeroed), the advertised context length, ``seq_lens_cuda = width`` and the + block-decode page / offset / bound operands. + + With ``cp_size > 1`` the context pages are helix-sharded (``HelixDraftShard``): + the store mask is -1 on other ranks' positions, the block operands are + rank-local and ``blk_owned`` marks the block positions this rank holds. + """ + device = ctx_len.device + long = dict(dtype=torch.long, device=device) + out = dict( + bonus=torch.empty((num_gens,), **long), + embed_ids=torch.empty((num_gens + 1,), **long), + query_positions=torch.empty((num_gens, block_size), **long), + ctx_positions=torch.empty((num_gens * width,), **long), + store_pages=torch.empty((num_gens * width,), **long), + store_offsets=torch.empty((num_gens * width,), **long), + store_mask=torch.empty((num_gens * width,), dtype=torch.int8, device=device), + num_ctx=torch.empty((num_gens,), **long), + blk_pages=torch.empty((num_gens, block_size), **long), + blk_offsets=torch.empty((num_gens, block_size), **long), + blk_seq_lens=torch.empty((num_gens,), dtype=torch.int32, device=device), + blk_kv_bounds=torch.empty((num_gens * block_size,), dtype=torch.int32, device=device), + ) + if cp_size > 1: + out["blk_owned"] = torch.empty((num_gens, block_size), dtype=torch.int8, device=device) + assert accepted_tokens.stride(-1) == 1 and block_tables.stride(-1) == 1 + _dflash_gen_prepare_kernel[(1,)]( + num_accepted_tokens, + accepted_tokens, + accepted_tokens.stride(0), + slots, + ctx_len, + block_counts, + block_tables, + block_tables.stride(0), + block_tables.shape[1], + seq_lens_cuda, + out["bonus"], + out["embed_ids"], + out["query_positions"], + out["ctx_positions"], + out["store_pages"], + out["store_offsets"], + out["store_mask"], + out["num_ctx"], + out["blk_pages"], + out["blk_offsets"], + out["blk_seq_lens"], + out["blk_kv_bounds"], + out.get("blk_owned", out["store_mask"]), + num_gens, + num_contexts, + dummy_slot, + max_ctx, + mask_token_id, + cp_rank, + PAGE_SIZE=page_size, + BLOCK_SIZE=block_size, + WIDTH=width, + CP_SIZE=cp_size, + BLOCK_G=triton.next_power_of_2(max(num_gens, 1)), + BLOCK_W=triton.next_power_of_2(max(width, block_size)), + num_warps=4, + ) + return out + + +@triton.jit +def _noise_embedding_kernel( + embed_ptr, out_ptr, num_gens, hidden, BLOCK_SIZE: tl.constexpr, BLOCK_H: tl.constexpr +): + g = tl.program_id(0) + j = tl.program_id(1) + offs = tl.program_id(2) * BLOCK_H + tl.arange(0, BLOCK_H) + mask = offs < hidden + src_row = tl.where(j == 0, g, num_gens).to(tl.int64) + vals = tl.load(embed_ptr + src_row * hidden + offs, mask=mask) + tl.store(out_ptr + (g.to(tl.int64) * BLOCK_SIZE + j) * hidden + offs, vals, mask=mask) + + +def dflash_noise_embedding(combined_embed: torch.Tensor, block_size: int) -> torch.Tensor: + """``[G, block, H]``: the mask-token embedding everywhere and the bonus + embedding at block slot 0, from ``combined_embed = embed([bonus..., mask])``.""" + num_gens = combined_embed.shape[0] - 1 + hidden = combined_embed.shape[-1] + combined_embed = combined_embed.contiguous() + out = torch.empty( + (num_gens, block_size, hidden), dtype=combined_embed.dtype, device=combined_embed.device + ) + if num_gens > 0: + block_h = 1024 + _noise_embedding_kernel[(num_gens, block_size, triton.cdiv(hidden, block_h))]( + combined_embed, out, num_gens, hidden, BLOCK_SIZE=block_size, BLOCK_H=block_h + ) + return out + + +@triton.jit +def _store_mla_context_kernel( + latent_ptr, + stride_latent_tok, + stride_latent_layer, + pool_ptrs_ptr, + stride_pool_page, + stride_pool_off, + pages_ptr, + offsets_ptr, + mask_ptr, + head_dim, + BLOCK_D: tl.constexpr, + POOL_FP8: tl.constexpr, +): + tok = tl.program_id(0) + layer = tl.program_id(1) + offs = tl.arange(0, BLOCK_D) + dmask = offs < head_dim + page = tl.load(pages_ptr + tok) + off = tl.load(offsets_ptr + tok) + flag = tl.load(mask_ptr + tok) + keep = (flag > 0).to(tl.float32) + src = latent_ptr + tok.to(tl.int64) * stride_latent_tok + layer * stride_latent_layer + vals = tl.load(src + offs, mask=dmask, other=0.0) + # Same bf16 multiply as the eager k.mul_(mask): exact for x * 1, signed zero for x * 0. + vals = (vals.to(tl.float32) * keep).to(vals.dtype) + if POOL_FP8: + # to_draft_fp8: saturated E4M3 at unit scale, round to nearest even. + vals = tl.minimum(tl.maximum(vals.to(tl.float32), -448.0), 448.0).to(tl.float8e4nv) + pool = tl.load(pool_ptrs_ptr + layer).to(tl.pointer_type(vals.dtype)) + dst = pool + page * stride_pool_page + off * stride_pool_off + tl.store(dst + offs, vals, mask=dmask & (flag >= 0)) + + +def store_mla_context_latents( + latent: torch.Tensor, + pools: list, + pool_ptrs: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + mask: torch.Tensor, +) -> None: + """``pools[l][pages[t], 0, 0, offsets[t]] = latent[t, l] * mask[t]`` for all t, l. + + Args: + latent: ``[N, L, D]`` per-token, per-layer MLA latents (any token and + layer strides, unit inner stride). + pools: the L per-layer ``[pages, 1, 1, page_size, D]`` pool views + (identical strides). + pool_ptrs: ``[L]`` int64 base addresses of ``pools``. + pages, offsets: ``[N]`` destination page and in-page offset. + mask: ``[N]`` int8 write mask (0 writes a zeroed latent, as eagerly; + -1 leaves the slot untouched). + + An FP8 E4M3 pool stores ``to_draft_fp8`` of the masked latent. + """ + num_tokens, num_layers, head_dim = latent.shape + if num_tokens == 0: + return + pool0 = pools[0] + assert latent.stride(-1) == 1 and pool0.stride(-1) == 1 + assert all(p.stride() == pool0.stride() and p.dtype == pool0.dtype for p in pools) + pool_fp8 = pool0.dtype == torch.float8_e4m3fn + assert pool_fp8 or pool0.dtype == latent.dtype + _store_mla_context_kernel[(num_tokens, num_layers)]( + latent, + latent.stride(0), + latent.stride(1), + pool_ptrs, + pool0.stride(0), + pool0.stride(3), + pages, + offsets, + mask, + head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + POOL_FP8=pool_fp8, + num_warps=4, + ) + + +@triton.jit +def _scatter_rows_kernel( + src_ptr, + stride_src, + dst_ptr, + stride_dst_page, + stride_dst_off, + pages_ptr, + offsets_ptr, + rows_mask_ptr, + width, + BLOCK_D: tl.constexpr, + POOL_FP8: tl.constexpr, + HAS_ROWS_MASK: tl.constexpr, +): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK_D) + mask = offs < width + if HAS_ROWS_MASK: + mask = mask & (tl.load(rows_mask_ptr + row) != 0) + page = tl.load(pages_ptr + row) + off = tl.load(offsets_ptr + row) + vals = tl.load(src_ptr + row.to(tl.int64) * stride_src + offs, mask=mask) + if POOL_FP8: + # to_draft_fp8: saturated E4M3 at unit scale, round to nearest even. + vals = tl.minimum(tl.maximum(vals.to(tl.float32), -448.0), 448.0).to(tl.float8e4nv) + tl.store(dst_ptr + page * stride_dst_page + off * stride_dst_off + offs, vals, mask=mask) + + +def scatter_rows_to_pool( + src: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + rows_mask: Optional[torch.Tensor] = None, +) -> None: + """``pool[pages[i], offsets[i]] = src[i]`` for a ``[pages, page_size, D]`` pool view. + + The same writes as the eager advanced-index assignment (a pure copy); an + FP8 E4M3 pool stores ``to_draft_fp8(src[i])``. With ``rows_mask`` (bool or + integer ``[rows]``), rows where it is 0 / False write nothing: the slot + keeps its value, as ``torch.where(rows_mask, src, pool[pages, offsets])`` + written back would leave it. + """ + rows, width = src.shape + if rows == 0: + return + assert src.stride(-1) == 1 and pool.stride(-1) == 1 and pool.shape[-1] == width + pool_fp8 = pool.dtype == torch.float8_e4m3fn + assert pool_fp8 or pool.dtype == src.dtype + if rows_mask is not None: + rows_mask = rows_mask.reshape(-1) + assert rows_mask.numel() == rows and rows_mask.is_contiguous() + if rows_mask.dtype == torch.bool: + rows_mask = rows_mask.view(torch.uint8) + _scatter_rows_kernel[(rows,)]( + src, + src.stride(0), + pool, + pool.stride(0), + pool.stride(1), + pages, + offsets, + pages if rows_mask is None else rows_mask, + width, + BLOCK_D=triton.next_power_of_2(width), + POOL_FP8=pool_fp8, + HAS_ROWS_MASK=rows_mask is not None, + num_warps=4, + ) + + +@triton.jit +def _next_new_tokens_kernel( + accepted_ptr, + stride_accepted, + num_accepted_ptr, + draft_ptr, + stride_draft, + out_ptr, + batch_size, + DRAFT_LEN: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_K: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_B + tl.arange(0, BLOCK_B) + cols = tl.arange(0, BLOCK_K) + rmask = rows < batch_size + nacc = tl.load(num_accepted_ptr + rows, mask=rmask, other=1) + last = tl.load(accepted_ptr + rows.to(tl.int64) * stride_accepted + (nacc - 1), mask=rmask) + width = DRAFT_LEN + 1 + tl.store(out_ptr + rows * width, last.to(out_ptr.dtype.element_ty), mask=rmask) + dmask = rmask[:, None] & (cols[None, :] < DRAFT_LEN) + d = tl.load(draft_ptr + rows[:, None].to(tl.int64) * stride_draft + cols[None, :], mask=dmask) + tl.store( + out_ptr + rows[:, None] * width + 1 + cols[None, :], + d.to(out_ptr.dtype.element_ty), + mask=dmask, + ) + + +@triton.jit +def _next_new_tokens_rewind_kernel( + accepted_ptr, + stride_accepted, + num_accepted_ptr, + draft_ptr, + stride_draft, + out_ptr, + batch_size, + kv_lens_ptr, + rewind_ptr, + num_ctx, + num_gen, + DRAFT_LEN: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_K: tl.constexpr, +): + # _next_new_tokens_kernel for a batch in one program, plus + # K3DFlashWorker._apply_kv_rewind_after_draft on the gen rows. + rows = tl.arange(0, BLOCK_B) + cols = tl.arange(0, BLOCK_K) + rmask = rows < batch_size + nacc = tl.load(num_accepted_ptr + rows, mask=rmask, other=1) + last = tl.load(accepted_ptr + rows.to(tl.int64) * stride_accepted + (nacc - 1), mask=rmask) + width = DRAFT_LEN + 1 + tl.store(out_ptr + rows * width, last.to(out_ptr.dtype.element_ty), mask=rmask) + dmask = rmask[:, None] & (cols[None, :] < DRAFT_LEN) + d = tl.load(draft_ptr + rows[:, None].to(tl.int64) * stride_draft + cols[None, :], mask=dmask) + tl.store( + out_ptr + rows[:, None] * width + 1 + cols[None, :], + d.to(out_ptr.dtype.element_ty), + mask=dmask, + ) + gmask = rows < num_gen + kv = tl.load(kv_lens_ptr + num_ctx + rows, mask=gmask) + amount = tl.load(rewind_ptr + rows, mask=gmask) + tl.store(kv_lens_ptr + num_ctx + rows, tl.maximum(kv - amount, 0), mask=gmask) + + +def fused_next_new_tokens_rewind( + accepted_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + kv_lens: torch.Tensor, + rewind: torch.Tensor, + num_contexts: int, +) -> torch.Tensor: + """:func:`fused_next_new_tokens` and ``kv_lens[num_contexts:][:num_gens] = + (kv_lens - rewind).clamp(min=0)`` in one launch (one program).""" + batch_size, draft_len = next_draft_tokens.shape + num_gens = rewind.shape[0] + dtype = torch.promote_types(accepted_tokens.dtype, next_draft_tokens.dtype) + out = torch.empty((batch_size, draft_len + 1), dtype=dtype, device=accepted_tokens.device) + if batch_size == 0 and num_gens == 0: + return out + assert accepted_tokens.stride(-1) == 1 and next_draft_tokens.stride(-1) == 1 + assert kv_lens.stride(-1) == 1 and rewind.is_contiguous() + _next_new_tokens_rewind_kernel[(1,)]( + accepted_tokens, + accepted_tokens.stride(0), + num_accepted_tokens, + next_draft_tokens, + next_draft_tokens.stride(0), + out, + batch_size, + kv_lens, + rewind, + num_contexts, + num_gens, + DRAFT_LEN=draft_len, + BLOCK_B=triton.next_power_of_2(max(batch_size, num_gens, 1)), + BLOCK_K=triton.next_power_of_2(max(draft_len, 1)), + num_warps=4, + ) + return out + + +def fused_next_new_tokens( + accepted_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> torch.Tensor: + """``cat([accepted[b, nacc[b] - 1], next_draft[b]])`` per request, in one launch + (``SpecWorkerBase._prepare_next_new_tokens`` with identity batch indices).""" + batch_size, draft_len = next_draft_tokens.shape + dtype = torch.promote_types(accepted_tokens.dtype, next_draft_tokens.dtype) + out = torch.empty((batch_size, draft_len + 1), dtype=dtype, device=accepted_tokens.device) + if batch_size == 0: + return out + assert accepted_tokens.stride(-1) == 1 and next_draft_tokens.stride(-1) == 1 + block_b = 128 + _next_new_tokens_kernel[(triton.cdiv(batch_size, block_b),)]( + accepted_tokens, + accepted_tokens.stride(0), + num_accepted_tokens, + next_draft_tokens, + next_draft_tokens.stride(0), + out, + batch_size, + DRAFT_LEN=draft_len, + BLOCK_B=block_b, + BLOCK_K=triton.next_power_of_2(max(draft_len, 1)), + num_warps=4, + ) + return out + + +@triton.jit +def _vocab_shard_embedding_kernel( + ids_ptr, weight_ptr, out_ptr, vocab_start, vocab_end, hidden, BLOCK_H: tl.constexpr +): + row = tl.program_id(0) + offs = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + mask = offs < hidden + token = tl.load(ids_ptr + row).to(tl.int64) + local = (token >= vocab_start) & (token < vocab_end) + src = tl.where(local, token - vocab_start, 0) + vals = tl.load(weight_ptr + src * hidden + offs, mask=mask & local, other=0.0) + tl.store(out_ptr + row.to(tl.int64) * hidden + offs, vals, mask=mask) + + +def vocab_shard_embedding_applies(embedding, ids: torch.Tensor) -> bool: + """Whether ``embedding(ids)`` of a vocab-parallel ``layers.EmbeddingTable`` can run + as :func:`vocab_shard_lookup` of this rank's rows followed by the table's all-reduce.""" + weight = getattr(embedding, "weight", None) + return ( + getattr(embedding, "vocab_parallel", False) + and getattr(embedding, "all_reduce", None) is not None + and weight is not None + and weight.is_contiguous() + and ids.dim() == 1 + ) + + +def vocab_shard_lookup( + ids: torch.Tensor, weight: torch.Tensor, vocab_start: int, vocab_end: int +) -> torch.Tensor: + """This rank's rows ``weight[ids - vocab_start]`` of a vocab shard ``[vocab_start, + vocab_end)``, zero for ids outside it: one lookup kernel instead of the eager mask / lookup / + fill chain, same values. The table's all-reduce of the result is ``embedding(ids)``.""" + hidden = weight.shape[1] + out = torch.empty((ids.shape[0], hidden), dtype=weight.dtype, device=weight.device) + if ids.shape[0] > 0: + block_h = 1024 + _vocab_shard_embedding_kernel[(ids.shape[0], triton.cdiv(hidden, block_h))]( + ids, + weight, + out, + vocab_start, + vocab_end, + hidden, + BLOCK_H=block_h, + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/vocab_parallel_kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/vocab_parallel_kernels.py new file mode 100644 index 000000000000..441b2b8b1ea2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec/vocab_parallel_kernels.py @@ -0,0 +1,717 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Triton kernels of the vocab-parallel token selection and their launchers. + +The selection reduces each rank's vocab shard of the logits to ``(value, global index)`` pairs +(:func:`local_select`, :func:`markov_add_local_argmax`), exchanges them, and reduces the exchanged +pairs to tokens (:func:`reduce_pairs`, :func:`markov_reduce_embed`); :func:`markov_chain` runs a +whole DSpark Markov chain in one launch. ``dflash.K3VocabParallelHead`` owns the exchange +buffers and calls these through the ``ops/spec`` entries. +""" + +from typing import NamedTuple, Optional + +import torch +import triton +import triton.language as tl + +_INT32_MAX = tl.constexpr(2**31 - 1) +# Index reported by a split without valid columns; exact in float32 and above +# every supported vocab id. +_EMPTY_INDEX = tl.constexpr(2**24) +_LOCAL_BLOCK_V = 2048 +# Shard columns per program of the one-launch Markov chain. +_MARKOV_BLOCK_V = 128 +# Capacity of the one-launch Markov chain's exchange sets. +_CHAIN_MAX_STEPS = 16 +_CHAIN_MAX_REQS = 16 +_REDUCE_BLOCK_ROWS = 32 +# Programs the local reductions aim to launch (rows x column splits). +_TARGET_PROGRAMS = 256 + + +@triton.jit +def _argmax_combine(val_a, idx_a, val_b, idx_b): + # torch.argmax order: NaN beats any number, equal values go to the lower index. + nan_a = val_a != val_a + nan_b = val_b != val_b + take_a = tl.where( + nan_a, + (~nan_b) | (idx_a < idx_b), + (~nan_b) & ((val_a > val_b) | ((val_a == val_b) & (idx_a < idx_b))), + ) + return tl.where(take_a, val_a, val_b), tl.where(take_a, idx_a, idx_b) + + +@triton.jit +def _gumbel(r): + # Standard Gumbel sample -log(-log(u)) from 32 random bits. 23 bits map to + # the open interval (0, 1); with 23 bits the midpoint (k + 0.5) / 2^23 is + # exact in float32, so u never rounds to 0 or 1. + u = ((r >> 9).to(tl.float32) + 0.5) * (1.0 / 8388608.0) + return -tl.log(-tl.log(u)) + + +@triton.jit +def _store_pair( + out_ptr, + row, + stride_out, + slot, + best_val, + gidx, + zero_lo, + zero_hi, + fill, + ZERO_WIDTH: tl.constexpr, +): + # Publishes (value, global index) at columns [2 * slot, 2 * slot + 1] of the + # row. With ZERO_WIDTH > 0 the program flagged `fill` also zeroes the row + # outside [zero_lo, zero_hi) -- the other ranks' columns -- so that an + # all-reduce sum acts as an all-gather. + base = out_ptr + row.to(tl.int64) * stride_out + tl.store(base + 2 * slot, best_val) + tl.store(base + 2 * slot + 1, gidx.to(tl.float32)) + if ZERO_WIDTH > 0: + if fill: + offs = tl.arange(0, ZERO_WIDTH) + tl.store( + base + offs, + tl.zeros([ZERO_WIDTH], tl.float32), + mask=(offs < zero_lo) | (offs >= zero_hi), + ) + + +@triton.jit +def _global_index(best_idx, col_base): + # A split without valid columns reports an index above any vocab id, so it + # never wins a tie. + return tl.where(best_idx == _INT32_MAX, _EMPTY_INDEX, col_base + best_idx) + + +@triton.jit +def _local_argmax_kernel( + x_ptr, + stride_x, + n_cols, + col_base, + chunk, + out_ptr, + stride_out, + slot0, + zero_lo, + zero_hi, + ZERO_WIDTH: tl.constexpr, + BLOCK_V: tl.constexpr, +): + # Program (row, split) reduces columns [split * chunk, (split + 1) * chunk). + row = tl.program_id(0) + split = tl.program_id(1) + lo = split * chunk + hi = tl.minimum(lo + chunk, n_cols) + offs = tl.arange(0, BLOCK_V) + x_row = x_ptr + row.to(tl.int64) * stride_x + best_val = tl.full([BLOCK_V], float("-inf"), tl.float32) + best_idx = tl.full([BLOCK_V], _INT32_MAX, tl.int32) + for start in range(lo, hi, BLOCK_V): + cols = start + offs + mask = cols < hi + vals = tl.load(x_row + cols, mask=mask, other=float("-inf")).to(tl.float32) + idx = tl.where(mask, cols, _INT32_MAX) + best_val, best_idx = _argmax_combine(best_val, best_idx, vals, idx) + val, idx = tl.reduce((best_val, best_idx), 0, _argmax_combine) + _store_pair( + out_ptr, + row, + stride_out, + slot0 + split, + val, + _global_index(idx, col_base), + zero_lo, + zero_hi, + split == 0, + ZERO_WIDTH, + ) + + +@triton.jit +def _local_sample_kernel( + x_ptr, + stride_x, + n_cols, + col_base, + chunk, + temp_ptr, + seed_ptr, + offset_ptr, + greedy_threshold, + out_ptr, + stride_out, + slot0, + zero_lo, + zero_hi, + ZERO_WIDTH: tl.constexpr, + BLOCK_V: tl.constexpr, +): + # Gumbel-max over one split of this rank's shard. The noise of global vocab + # index v is output (v % 4) of Philox(seed, counter=(v // 4, row, offset + # lo, offset hi)): disjoint streams per row and per step, identical on + # every rank. col_base and chunk are multiples of 4, so 4-column groups + # never straddle splits or shards. + GROUPS: tl.constexpr = BLOCK_V // 4 + row = tl.program_id(0) + split = tl.program_id(1) + lo = split * chunk + hi = tl.minimum(lo + chunk, n_cols) + offs_g = tl.arange(0, GROUPS) + offs_l = tl.arange(0, 4) + x_row = x_ptr + row.to(tl.int64) * stride_x + temperature = tl.load(temp_ptr + row).to(tl.float32) + is_greedy = temperature <= greedy_threshold + inv_t = 1.0 / tl.where(is_greedy, 1.0, temperature) + seed = tl.load(seed_ptr + row).to(tl.int64) + offset = tl.load(offset_ptr + row).to(tl.int64) + c1 = tl.full([GROUPS], 0, tl.uint32) + row.to(tl.uint32) + c2 = tl.full([GROUPS], 0, tl.uint32) + (offset & 0xFFFFFFFF).to(tl.uint32) + c3 = tl.full([GROUPS], 0, tl.uint32) + ((offset >> 32) & 0xFFFFFFFF).to(tl.uint32) + best_val = tl.full([GROUPS, 4], float("-inf"), tl.float32) + best_idx = tl.full([GROUPS, 4], _INT32_MAX, tl.int32) + for start in range(lo, hi, BLOCK_V): + cols = start + offs_g[:, None] * 4 + offs_l[None, :] + mask = cols < hi + vals = tl.load(x_row + cols, mask=mask, other=float("-inf")).to(tl.float32) + if not is_greedy: + counter = ((col_base + start) // 4 + offs_g).to(tl.uint32) + r0, r1, r2, r3 = tl.philox(seed, counter, c1, c2, c3) + # [GROUPS, 2, 2] -> [GROUPS, 4] in lane order r0, r1, r2, r3. + noise = tl.join( + tl.join(_gumbel(r0), _gumbel(r2)), tl.join(_gumbel(r1), _gumbel(r3)) + ).reshape([GROUPS, 4]) + vals = vals * inv_t + noise + idx = tl.where(mask, cols, _INT32_MAX) + best_val, best_idx = _argmax_combine(best_val, best_idx, vals, idx) + best_val, best_idx = tl.reduce((best_val, best_idx), 1, _argmax_combine) + val, idx = tl.reduce((best_val, best_idx), 0, _argmax_combine) + _store_pair( + out_ptr, + row, + stride_out, + slot0 + split, + val, + _global_index(idx, col_base), + zero_lo, + zero_hi, + split == 0, + ZERO_WIDTH, + ) + + +@triton.jit +def _reduce_pairs_kernel( + pairs_ptr, + stride_row, + stride_slot, + n_rows, + n_slots, + out_ptr, + BLOCK_ROWS: tl.constexpr, + BLOCK_SLOTS: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + slots = tl.arange(0, BLOCK_SLOTS) + mask = (rows[:, None] < n_rows) & (slots[None, :] < n_slots) + ptrs = pairs_ptr + rows[:, None].to(tl.int64) * stride_row + slots[None, :] * stride_slot + vals = tl.load(ptrs, mask=mask, other=float("-inf")) + gidx = tl.load(ptrs + 1, mask=mask, other=0.0).to(tl.int32) + gidx = tl.where(mask, gidx, _INT32_MAX) + _, best = tl.reduce((vals, gidx), 1, _argmax_combine) + tl.store(out_ptr + rows, best.to(out_ptr.dtype.element_ty), mask=rows < n_rows) + + +@triton.jit +def _markov_reduce_embed_kernel( + pairs_ptr, + stride_row, + stride_slot, + n_slots, + anchor_ptr, + w1_ptr, + emb_ptr, + tokens_ptr, + stride_tokens, + vocab, + step, + RANK: tl.constexpr, + FROM_ANCHOR: tl.constexpr, + BLOCK_SLOTS: tl.constexpr, + BLOCK_R: tl.constexpr, +): + # Resolves prev_step for one request (the anchor at step 0, else the global + # argmax of the previous step, which is also draft token step - 1) and + # gathers its markov_w1 row; an out-of-vocab prev gives a zero row. + req = tl.program_id(0) + if FROM_ANCHOR: + prev = tl.load(anchor_ptr + req).to(tl.int64) + else: + slots = tl.arange(0, BLOCK_SLOTS) + smask = slots < n_slots + ptrs = pairs_ptr + req.to(tl.int64) * stride_row + slots * stride_slot + vals = tl.load(ptrs, mask=smask, other=float("-inf")) + gidx = tl.load(ptrs + 1, mask=smask, other=0.0).to(tl.int32) + gidx = tl.where(smask, gidx, _INT32_MAX) + _, best = tl.reduce((vals, gidx), 0, _argmax_combine) + prev = best.to(tl.int64) + tl.store(tokens_ptr + req.to(tl.int64) * stride_tokens + step - 1, prev.to(tl.int32)) + valid = (prev >= 0) & (prev < vocab) + offs_r = tl.arange(0, BLOCK_R) + mask_r = offs_r < RANK + w1_row = tl.load( + w1_ptr + tl.where(valid, prev, 0) * RANK + offs_r, mask=mask_r & valid, other=0.0 + ) + tl.store(emb_ptr + req * RANK + offs_r, w1_row, mask=mask_r) + + +@triton.jit +def _markov_add_local_argmax_kernel( + base_ptr, + stride_base, + bias_ptr, + stride_bias, + n_cols, + col_base, + chunk, + out_ptr, + stride_out, + slot0, + zero_lo, + zero_hi, + ZERO_WIDTH: tl.constexpr, + BLOCK_V: tl.constexpr, +): + # base (shard logits of block position `step`) + bias, upcast to fp32 like + # the gathered reference, then the argmax of one split of the shard. + req = tl.program_id(0) + split = tl.program_id(1) + lo = split * chunk + hi = tl.minimum(lo + chunk, n_cols) + offs = tl.arange(0, BLOCK_V) + base_row = base_ptr + req.to(tl.int64) * stride_base + bias_row = bias_ptr + req.to(tl.int64) * stride_bias + best_val = tl.full([BLOCK_V], float("-inf"), tl.float32) + best_idx = tl.full([BLOCK_V], _INT32_MAX, tl.int32) + for start in range(lo, hi, BLOCK_V): + cols = start + offs + mask = cols < hi + base = tl.load(base_row + cols, mask=mask, other=0.0).to(tl.float32) + bias = tl.load(bias_row + cols, mask=mask, other=0.0).to(tl.float32) + vals = tl.where(mask, base + bias, float("-inf")) + idx = tl.where(mask, cols, _INT32_MAX) + best_val, best_idx = _argmax_combine(best_val, best_idx, vals, idx) + val, idx = tl.reduce((best_val, best_idx), 0, _argmax_combine) + _store_pair( + out_ptr, + req, + stride_out, + slot0 + split, + val, + _global_index(idx, col_base), + zero_lo, + zero_hi, + split == 0, + ZERO_WIDTH, + ) + + +@triton.jit +def _gdc_wait(): + # griddepcontrol.wait with a memory clobber, so that no load is moved above + # it (tl.extra.cuda.gdc_wait has none). + tl.inline_asm_elementwise( + "griddepcontrol.wait; // dummy $0", + "=r,~{memory}", + [], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +@triton.jit +def _argmax_pairs(vals, idx, axis: tl.constexpr): + # The (value, index) that a _argmax_combine reduction along ``axis`` returns, + # from native reductions: NaN beats every number and the lowest index wins + # among NaNs; otherwise the largest value, lowest index among equal values + # (+0.0 == -0.0). The selected index is unique, so it does not depend on the + # reduction order; the value is the selected entry's (NaN, or the maximum). + is_nan = vals != vals + any_nan = tl.max(is_nan.to(tl.int32), axis=axis) > 0 + vmax = tl.max(tl.where(is_nan, float("-inf"), vals), axis=axis) + cand = tl.where( + tl.expand_dims(any_nan, axis), is_nan, (vals == tl.expand_dims(vmax, axis)) & ~is_nan + ) + best = tl.min(tl.where(cand, idx, _INT32_MAX), axis=axis) + return tl.where(any_nan, float("nan"), vmax), best + + +_NEG_ZERO_BITS = tl.constexpr(-(2**31)) + + +@triton.jit +def _wait_pairs_written(xptr, xm): + # Spin until no pair of the tile holds the -0.0 sentinel. Only a scalar is + # loop-carried: Triton 3.8 fails to compile (TritonGPURemoveLayoutConversions) + # a while loop that carries the [16, TP] tiles once TP >= 8. + pending = tl.full((), 1, tl.int32) + while pending > 0: + xv = tl.load(xptr, mask=xm, other=0.0, volatile=True) + xi = tl.load(xptr + 1, mask=xm, other=0.0, volatile=True) + dirty = xm & ( + (xv.to(tl.int32, bitcast=True) == _NEG_ZERO_BITS) + | (xi.to(tl.int32, bitcast=True) == _NEG_ZERO_BITS) + ) + pending = tl.max(tl.max(dirty.to(tl.int32), axis=1), axis=0) + + +@triton.jit +def _markov_chain_kernel( + anchor_ptr, + w1_ptr, + w2_ptr, + stride_w2, + base_ptr, + stride_base_req, + stride_base_step, + tokens_ptr, + stride_tokens, + partial_ptr, + ticket_ptr, + launch_ptr, + mc_ptr, + uc_ptr, + num_reqs, + vocab, + n_cols, + col_base, + tp_rank, + STEPS: tl.constexpr, + RANK: tl.constexpr, + TP: tl.constexpr, + BLOCK_B: tl.constexpr, + BLOCK_V: tl.constexpr, + BLOCK_P: tl.constexpr, + SET_FLOATS: tl.constexpr, +): + # The whole DSpark Markov chain of one draft block in one launch. Every + # step: resolve prev, gather its markov_w1 row, compute this program's + # markov_w2 bias columns, add the step's shard logits and reduce to a + # (value, global index) candidate; the last program of the rank to finish + # reduces the rank's candidates and multicasts the rank's pair per request + # into every rank's exchange buffer; every program then polls its local copy + # until all ranks' pairs of the step are present (Lamport: the buffer holds + # -0.0 where nothing was written yet; published values never are -0.0). + # + # Exchange sets: launch L uses set L % 3 (STEPS regions of [request, rank, + # 2]) and clears set (L - 1) % 3, which every program of this rank finished + # reading in launch L - 1; another rank writes that set again only in launch + # L + 2, which needs this rank's pairs of launch L + 1, published after this + # launch completed -- so the clear needs no fence. All programs read L + # before their first ticket; program 0 advances it after the last step. + # Every program must be resident at once (the caller launches at most one + # per SM), since each one waits for all programs' pairs every step. + pid = tl.program_id(0) + nprog = tl.num_programs(0) + reqs = tl.arange(0, BLOCK_B) + rmask = reqs < num_reqs + cols = pid * BLOCK_V + tl.arange(0, BLOCK_V) + cmask = cols < n_cols + r = tl.arange(0, RANK) + ranks = tl.arange(0, TP) + launch = tl.load(launch_ptr, volatile=True) + cur_set = launch % 3 + if pid == 0: + # The -0.0 sentinel is written as its bit pattern: a -0.0 float constant + # is not guaranteed to keep its sign through compilation. + z = tl.arange(0, SET_FLOATS) + tl.store( + uc_ptr.to(tl.pointer_type(tl.int32)) + ((launch + 2) % 3) * SET_FLOATS + z, + tl.full([SET_FLOATS], _NEG_ZERO_BITS, tl.int32), + ) + # The step logits come from the previous grid; nothing is read before this. + _gdc_wait() + prev = tl.load(anchor_ptr + reqs, mask=rmask, other=-1).to(tl.int64) + xm = rmask[:, None] & (ranks[None, :] < TP) + for step in tl.static_range(STEPS + 1): + if step < STEPS: + # Issued before the wait for the previous step's pairs: neither + # depends on it. + w2 = tl.load( + w2_ptr + cols[:, None].to(tl.int64) * stride_w2 + r[None, :], mask=cmask[:, None] + ) + base = tl.load( + base_ptr + + reqs[:, None].to(tl.int64) * stride_base_req + + step * stride_base_step + + cols[None, :], + mask=rmask[:, None] & cmask[None, :], + other=0.0, + ) + if step > 0: + # Wait for every rank's pair of the previous step, then take the winner. + region = cur_set * SET_FLOATS + (step - 1) * (BLOCK_B * TP * 2) + xptr = uc_ptr + region + (reqs[:, None] * TP + ranks[None, :]) * 2 + _wait_pairs_written(xptr, xm) + # Read again once all are written: a slot is written once per launch + # and cleared only in a later one, so this read sees the polled values. + xv = tl.load(xptr, mask=xm, other=0.0, volatile=True) + xi = tl.load(xptr + 1, mask=xm, other=0.0, volatile=True) + _, best = _argmax_pairs( + tl.where(xm, xv, float("-inf")), tl.where(xm, xi.to(tl.int32), _INT32_MAX), 1 + ) + prev = best.to(tl.int64) + if pid == 0: + tl.store( + tokens_ptr + reqs.to(tl.int64) * stride_tokens + step - 1, + prev.to(tl.int32), + mask=rmask, + ) + if step < STEPS: + valid = rmask & (prev >= 0) & (prev < vocab) + w1 = tl.load( + w1_ptr + tl.where(valid, prev, 0)[:, None] * RANK + r[None, :], + mask=valid[:, None], + other=0.0, + ) + # F.linear(markov_w1[prev], w2) rounds the FP32 dot to the weights' dtype. + bias = tl.dot(w1, tl.trans(w2)).to(w2.dtype).to(tl.float32) + vals = tl.where( + rmask[:, None] & cmask[None, :], base.to(tl.float32) + bias, float("-inf") + ) + idx = tl.where(cmask[None, :], cols[None, :], _INT32_MAX) + tl.zeros( + [BLOCK_B, BLOCK_V], tl.int32 + ) + cand_val, cand_idx = _argmax_pairs(vals, idx, 1) + cand_idx = _global_index(cand_idx, col_base) + pbase = partial_ptr + ((step * nprog + pid) * BLOCK_B + reqs) * 2 + tl.store(pbase, cand_val) + tl.store(pbase + 1, cand_idx.to(tl.float32)) + tl.debug_barrier() + ticket = tl.atomic_add(ticket_ptr + step, 1, sem="acq_rel", scope="gpu") + if ticket == nprog - 1: + p = tl.arange(0, BLOCK_P) + pm = (p[:, None] < nprog) & rmask[None, :] + qptrs = partial_ptr + ((step * nprog + p[:, None]) * BLOCK_B + reqs[None, :]) * 2 + qv = tl.load(qptrs, mask=pm, other=float("-inf"), volatile=True) + qi = tl.where( + pm, + tl.load(qptrs + 1, mask=pm, other=0.0, volatile=True).to(tl.int32), + _INT32_MAX, + ) + win_val, win_idx = _argmax_pairs(qv, qi, 0) + # -0.0 marks an empty slot; publish +0.0 instead (argmax-equal). + win_val = tl.where( + win_val.to(tl.int32, bitcast=True) == _NEG_ZERO_BITS, 0.0, win_val + ) + tl.atomic_xchg(ticket_ptr + step, 0, sem="relaxed", scope="gpu") + out = ( + mc_ptr + + cur_set * SET_FLOATS + + step * (BLOCK_B * TP * 2) + + (reqs * TP + tp_rank) * 2 + ) + tl.store(out, win_val, mask=rmask) + tl.store(out + 1, win_idx.to(tl.float32), mask=rmask) + if pid == 0: + tl.store(launch_ptr, launch + 1) + + +class _PairLayout(NamedTuple): + """Where the local kernels publish pairs, for one exchange.""" + + buffer: torch.Tensor + stride_out: int + slot0: int # first slot owned by this rank + zero_lo: int # columns [zero_lo, zero_hi) are this rank's; the rest is zero-filled + zero_hi: int + zero_width: int # 0: no zero fill (NCCL all-gather) + splits: int # column splits per rank (slots per rank) + chunk: int # columns per split + + +def local_select( + local: torch.Tensor, + col_base: int, + layout: _PairLayout, + sampling: Optional[tuple] = None, +) -> None: + """Publish each split's ``(value, global index)`` pair of the shard logits ``local`` + ``[rows, shard]`` into ``layout``: the argmax, or with ``sampling = (temperatures, seeds, + offsets, greedy_threshold)`` the Gumbel-max draw of temperature-only sampling.""" + rows, n_cols = local.shape + if rows == 0: + return + kernel = _local_argmax_kernel if sampling is None else _local_sample_kernel + extra = () if sampling is None else tuple(sampling) + kernel[(rows, layout.splits)]( + local, + local.stride(0), + n_cols, + col_base, + layout.chunk, + *extra, + layout.buffer, + layout.stride_out, + layout.slot0, + layout.zero_lo, + layout.zero_hi, + ZERO_WIDTH=layout.zero_width, + BLOCK_V=_LOCAL_BLOCK_V, + num_warps=8, + ) + + +def reduce_pairs( + pairs: torch.Tensor, stride_row: int, stride_slot: int, n_slots: int, out: torch.Tensor +) -> torch.Tensor: + """``out[r]`` = the global index of the winning exchanged pair of row ``r``.""" + rows = out.shape[0] + if rows > 0: + _reduce_pairs_kernel[(triton.cdiv(rows, _REDUCE_BLOCK_ROWS),)]( + pairs, + stride_row, + stride_slot, + rows, + n_slots, + out, + BLOCK_ROWS=_REDUCE_BLOCK_ROWS, + BLOCK_SLOTS=triton.next_power_of_2(n_slots), + num_warps=4, + ) + return out + + +def markov_reduce_embed( + pairs: torch.Tensor, + stride_row: int, + stride_slot: int, + n_slots: int, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + emb: torch.Tensor, + tokens: torch.Tensor, + vocab: int, + step: int, +) -> None: + """Resolve ``prev_step`` per request (the anchor at step 0, else the winner of the exchanged + pairs, also stored as draft token ``step - 1``) and gather its ``markov_w1`` row into + ``emb``.""" + num_reqs = emb.shape[0] + rank = markov_w1.shape[1] + _markov_reduce_embed_kernel[(num_reqs,)]( + pairs, + stride_row, + stride_slot, + n_slots, + anchors, + markov_w1, + emb, + tokens, + tokens.stride(0), + vocab, + step, + RANK=rank, + FROM_ANCHOR=step == 0, + BLOCK_SLOTS=triton.next_power_of_2(n_slots), + BLOCK_R=triton.next_power_of_2(rank), + num_warps=4, + ) + + +def markov_add_local_argmax( + step_logits: torch.Tensor, bias: torch.Tensor, col_base: int, layout: _PairLayout +) -> None: + """Publish each split's argmax pair of ``step_logits + bias`` (``[requests, shard]``, added + in FP32) into ``layout``.""" + num_reqs, shard = step_logits.shape + _markov_add_local_argmax_kernel[(num_reqs, layout.splits)]( + step_logits, + step_logits.stride(0), + bias, + bias.stride(0), + shard, + col_base, + layout.chunk, + layout.buffer, + layout.stride_out, + layout.slot0, + layout.zero_lo, + layout.zero_hi, + ZERO_WIDTH=layout.zero_width, + BLOCK_V=_LOCAL_BLOCK_V, + num_warps=8, + ) + + +def markov_chain( + local_block: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: torch.Tensor, + tokens: torch.Tensor, + partial: torch.Tensor, + tickets: torch.Tensor, + launch: torch.Tensor, + mcast: torch.Tensor, + local: torch.Tensor, + set_floats: int, + vocab: int, + col_base: int, + tp: int, + tp_rank: int, +) -> torch.Tensor: + """The DSpark greedy Markov chain of one draft block in one launch (``_markov_chain_kernel``): + ``tokens`` ``[B, K]`` int32 from the shard block logits ``[B, K, shard]``.""" + num_reqs, block_size, shard = local_block.shape + programs = triton.cdiv(shard, _MARKOV_BLOCK_V) + _markov_chain_kernel[(programs,)]( + first_prev_tokens, + markov_w1, + w2_shard, + w2_shard.stride(0), + local_block, + local_block.stride(0), + local_block.stride(1), + tokens, + tokens.stride(0), + partial, + tickets, + launch, + mcast, + local, + num_reqs, + vocab, + shard, + col_base, + tp_rank, + STEPS=block_size, + RANK=markov_w1.shape[1], + TP=tp, + BLOCK_B=_CHAIN_MAX_REQS, + BLOCK_V=_MARKOV_BLOCK_V, + BLOCK_P=triton.next_power_of_2(programs), + SET_FLOATS=set_floats, + num_warps=4, + launch_pdl=True, + ) + return tokens diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec_decode_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec_decode_head.py new file mode 100644 index 000000000000..6d6c9dc40137 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/spec_decode_head.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The head of a speculative decode step in one launch. + +The head is the device work a captured decode step runs before its first layer: +the overlap-scheduler correction of position ids and KV lengths and the Helix +verify-group buffers derived from the corrected positions +(``TrtllmAttentionMetadata.recompute_helix_spec_buffers``). ``SpecDecodeHead`` +is a ``RuntimeHooks.decode_prologue`` that writes the same values with one +kernel (``ops/spec/k3_spec_decode_head``) instead of one launch per buffer. + +It reads main's ``DecoderInputPreparer`` and ``TrtllmAttentionMetadata`` directly +(``_config.disable_overlap_scheduler``, ``_buffers.previous_*_offsets_cuda``, +``_helix_spec_tokens_valid``); ``test_k3_engine_probes`` pins those names. +""" + +from typing import TYPE_CHECKING, Any, Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + +# The module, not the instance: the entry imports this module for its launcher. +from ..ops.spec import k3_spec_decode_head as _entry + +if TYPE_CHECKING: + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer + + +@triton.jit +def _helix_local_len(g, tokens_per_block, ledger, rank_base): + # Rank-local prefix length of the first g global tokens (populate_helix_spec_buffers). + full = g // ledger + return full * tokens_per_block + tl.minimum( + tl.maximum(g - full * ledger - rank_base, 0), tokens_per_block + ) + + +_SCALARS = [ + "num_pos", + "kv_start", + "num_kv", + "num_helix", + "tokens_per_block", + "cp_size", + "cp_rank", + "kv_bound_start", + "tokens_per_seq", +] + + +@triton.jit(do_not_specialize=_SCALARS, do_not_specialize_on_alignment=_SCALARS) +def _spec_head_kernel( + pos_ids_ptr, + pos_offset_ptr, + num_pos, + kv_lens_ptr, + kv_offset_ptr, + kv_start, + num_kv, + helix_pos_ptr, + helix_slots_ptr, + helix_bounds_ptr, + helix_zero_ptr, + num_helix, + tokens_per_block, + cp_size, + cp_rank, + kv_bound_start, + tokens_per_seq, + HAS_OVERLAP: tl.constexpr, + HAS_HELIX: tl.constexpr, + BLOCK: tl.constexpr, +): + # One program: the per-sequence KV lengths read back the bounds it stores. + lanes = tl.arange(0, BLOCK) + if HAS_OVERLAP: + for start in range(0, num_pos, BLOCK): + t = start + lanes + mask = t < num_pos + pos = tl.load(pos_ids_ptr + t, mask=mask) + tl.store(pos_ids_ptr + t, pos + tl.load(pos_offset_ptr + t, mask=mask), mask=mask) + for start in range(0, num_kv, BLOCK): + s = start + lanes + mask = s < num_kv + kv = tl.load(kv_lens_ptr + kv_start + s, mask=mask) + tl.store( + kv_lens_ptr + kv_start + s, kv + tl.load(kv_offset_ptr + s, mask=mask), mask=mask + ) + if HAS_HELIX: + ledger = tokens_per_block * cp_size + rank_base = cp_rank * tokens_per_block + for start in range(0, num_helix, BLOCK): + t = start + lanes + mask = t < num_helix + position = tl.load(helix_pos_ptr + t, mask=mask, other=0) + if HAS_OVERLAP: + position += tl.load(pos_offset_ptr + t, mask=mask, other=0) + tl.store(helix_pos_ptr + t, position, mask=mask) + local_before = _helix_local_len(position, tokens_per_block, ledger, rank_base) + bound = _helix_local_len(position + 1, tokens_per_block, ledger, rank_base) + owned = (position // tokens_per_block) % cp_size == cp_rank + tl.store(helix_slots_ptr + t, tl.where(owned, local_before, -1), mask=mask) + tl.store(helix_bounds_ptr + t, bound, mask=mask) + tl.store(helix_zero_ptr + t, bound == 0, mask=mask) + # A sequence's KV length is the bound of its last token. The lane that + # stores it is the one that stored its overlap-corrected value above. + tl.debug_barrier() + num_seqs = num_helix // tokens_per_seq + for start in range(0, num_seqs, BLOCK): + s = start + lanes + mask = s < num_seqs + last = tl.load(helix_bounds_ptr + s * tokens_per_seq + tokens_per_seq - 1, mask=mask) + tl.store(kv_lens_ptr + kv_bound_start + s, last, mask=mask) + + +def fused_spec_head_prep( + position_ids: torch.Tensor, + pos_offsets: torch.Tensor, + num_pos: int, + kv_lens: torch.Tensor, + kv_offsets: torch.Tensor, + kv_start: int, + num_kv: int, + *, + has_overlap: bool, + helix_positions: Optional[torch.Tensor] = None, + helix_slots: Optional[torch.Tensor] = None, + helix_bounds: Optional[torch.Tensor] = None, + helix_zero_kv: Optional[torch.Tensor] = None, + num_helix: int = 0, + tokens_per_block: int = 1, + cp_size: int = 1, + cp_rank: int = 0, + kv_bound_start: int = 0, + tokens_per_seq: int = 1, +) -> None: + """The decode-step head bookkeeping in one launch. + + With ``has_overlap``: ``position_ids[:num_pos] += pos_offsets[:num_pos]`` and + ``kv_lens[kv_start:kv_start + num_kv] += kv_offsets[:num_kv]``. With Helix + buffers (``num_helix > 0``), after ``helix_positions[:num_helix] += + pos_offsets[:num_helix]`` (only with ``has_overlap``), what + ``TrtllmAttentionMetadata.recompute_helix_spec_buffers(num_helix, tokens_per_seq)`` + writes: rank-local write slots, rank-local KV bounds, the zero-bound mask and + ``kv_lens[kv_bound_start + i]`` = the bound of sequence ``i``'s last token. + + All integer tensors are int32 with unit stride; ``helix_zero_kv`` is bool. + """ + for tensor in (position_ids, pos_offsets, kv_lens, kv_offsets): + assert tensor.dtype == torch.int32 and tensor.stride(-1) == 1 + has_helix = num_helix > 0 + if has_helix: + for tensor in (helix_positions, helix_slots, helix_bounds): + assert tensor.dtype == torch.int32 and tensor.is_contiguous() + assert helix_zero_kv.dtype == torch.bool and helix_zero_kv.is_contiguous() + assert num_helix % tokens_per_seq == 0 + else: + helix_positions = helix_slots = helix_bounds = helix_zero_kv = kv_lens + _spec_head_kernel[(1,)]( + position_ids, + pos_offsets, + num_pos, + kv_lens, + kv_offsets, + kv_start, + num_kv, + helix_positions, + helix_slots, + helix_bounds, + helix_zero_kv, + num_helix, + tokens_per_block, + cp_size, + cp_rank, + kv_bound_start, + tokens_per_seq, + HAS_OVERLAP=has_overlap, + HAS_HELIX=has_helix, + BLOCK=512, + num_warps=4, + ) + + +def spec_decode_kv_range(attn_metadata: TrtllmAttentionMetadata) -> tuple[int, int]: + """``(first sequence, count)`` whose KV lengths a speculative step corrects: the generation + requests, or the context requests they run as when they extend a context + (``num_chunked_ctx_requests``); the two branches of ``DecoderInputPreparer.preprocess_inputs``.""" + num_ctx = attn_metadata.num_contexts + num_chunked = attn_metadata.num_chunked_ctx_requests + if num_ctx >= num_chunked and num_chunked > 0: + return num_ctx - num_chunked, num_chunked + return num_ctx, attn_metadata.num_generations + + +class SpecDecodeHead: + """``RuntimeHooks.decode_prologue`` running a speculative decode step's head in one kernel. + + Covers int32 two-dimensional position ids and int32 KV lengths of + TRT-LLM attention metadata; any other step keeps the native prologue. + The kernel also writes the per-token zero-KV flags (``helix_kv_bounds == 0``); + main's metadata keeps no buffer for them, so they go to a scratch buffer here. + """ + + def __init__(self) -> None: + self._zero_kv: Optional[torch.Tensor] = None + + def _zero_kv_scratch(self, like: torch.Tensor) -> torch.Tensor: + scratch = self._zero_kv + if scratch is None or scratch.device != like.device or scratch.numel() < like.numel(): + scratch = torch.empty(like.numel(), dtype=torch.bool, device=like.device) + self._zero_kv = scratch + return scratch + + def __call__( + self, + preparer: "DecoderInputPreparer", + inputs: dict[str, Any], + *, + enable_spec_decode: bool, + runtime_draft_len: int, + ) -> bool: + md = inputs.get("attn_metadata") + if ( + not enable_spec_decode + or not isinstance(md, TrtllmAttentionMetadata) + or md.kv_cache_manager is None + ): + return False + kv_lens = md.kv_lens_cuda + overlap = not preparer._config.disable_overlap_scheduler + helix = preparer.mapping.has_cp_helix() and md._helix_spec_tokens_valid + position_ids = inputs["position_ids"] + if ( + not (overlap or helix) + or kv_lens.dtype != torch.int32 + or position_ids.ndim != 2 + or position_ids.dtype != torch.int32 + ): + return False + num_ctx_tokens = md.num_ctx_tokens + gen_tokens = inputs["input_ids"].shape[0] - num_ctx_tokens + num_ctx_requests = md.num_contexts + kv_start, num_kv = spec_decode_kv_range(md) + helix_kwargs = {} + if helix: + tokens_per_seq = preparer.get_runtime_tokens_per_gen_step(runtime_draft_len) + assert gen_tokens % tokens_per_seq == 0 and ( + gen_tokens // tokens_per_seq == md.num_generations + ), ( + f"helix spec expects uniform verify groups: {gen_tokens} gen tokens over " + f"{md.num_generations} generation rows of width {tokens_per_seq}" + ) + helix_kwargs = dict( + helix_positions=md.helix_position_offsets, + helix_slots=md.helix_local_slots, + helix_bounds=md.helix_kv_bounds, + helix_zero_kv=self._zero_kv_scratch(md.helix_kv_bounds), + num_helix=gen_tokens, + tokens_per_block=md.kv_cache_manager.tokens_per_block, + cp_size=md.mapping.cp_size, + cp_rank=md.mapping.cp_rank, + kv_bound_start=num_ctx_requests, + tokens_per_seq=tokens_per_seq, + ) + buffers = preparer._buffers + _entry.k3_spec_decode_head( + position_ids[0, num_ctx_tokens:], + buffers.previous_pos_id_offsets_cuda, + gen_tokens, + kv_lens, + buffers.previous_kv_lens_offsets_cuda, + kv_start, + num_kv, + has_overlap=overlap, + **helix_kwargs, + ) + md.on_update_kv_lens() + return True + + +spec_decode_head = SpecDecodeHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/layers.py new file mode 100644 index 000000000000..452d239dcc00 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/layers.py @@ -0,0 +1,1289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kernel-call primitives of the Kimi K3 Rubin model: K3-owned weights and direct kernel calls. + +Each container is an ``nn.Module`` that registers the parameter names the checkpoint and the +kernels read (``weight``, ``weight_scale``, ...), loads its rank's shard with an explicit +``load`` (and ``prepare`` where the kernel wants another layout), and runs one fixed kernel chain: + +* ``linear``: BF16 ``F.linear``; FP8 128 x 128 block scales on the low-M GEMM, the SM107 MXFP8 + GEMM (``mxfp8_gemm``) or the SM100 / SM103 swap-AB GEMM. +* ``ops/gemm/_mxfp8_gemm``: the SM107 MXFP8 GEMM over K3's tactics and its autotuner runner. +* ``norm``: FlashInfer RMSNorm (plain and fused residual add) on a K3 weight. +* ``mlp``: gate-up GEMM, SiTU / SiLU-and-mul, down GEMM and its all-reduce. +* ``embedding``: masked (vocab-parallel) lookup and its all-reduce. +* ``conv``: the short-convolution weight the KDA kernels read. +* ``allreduce``: the TP all-reduce on the MNNVL / trtllm all-reduce op. +* ``streams``: two branches on the current and an auxiliary stream. + +Tensor-parallel all-reduce of K3's row-parallel projections, as one fixed kernel call. + +:class:`TpAllReduce` decides its transport once, at construction, from the mapping and the requested +strategy (``ModelConfig.allreduce_strategy``, ``AUTO`` in every K3 recipe): + +* **MNNVL**: ``AUTO`` / ``MNNVL`` on a group MNNVL serves (BF16 / FP16 / FP32, no CP, aarch64 with + every NVLink up, and the group spans nodes, an explicit ``MNNVL``, or ``TLLM_TEST_MNNVL=1``). Each + call is ``trtllm::mnnvl_fusion_allreduce`` on the Lamport workspace of + ``get_or_scale_allreduce_mnnvl_workspace`` (one per mapping, shared with every other MNNVL + all-reduce and the pushed all-reduces of ``layers.py``; allocated here, a TP-wide collective, + and grown on the first eager call that needs more). The kernel is one-shot up to + ``TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES`` of ``tokens * hidden * ranks * element size``, two-shot above. + Every K3 all-reduce runs here on Rubin / GB300 NVL72: helix GEN groups span nodes, and the dense + MLP's TP-4 subgroup is declared over the whole world. +* **The trtllm all-reduce op** (the catalog's ``allreduce`` entry: ``trtllm::allreduce`` on the IPC + workspace of ``get_allreduce_workspace``) for any other group, with the strategy fixed: an + explicit one as given, ``AUTO`` as the op's own size heuristic + (``TLLM_DISABLE_ALLREDUCE_AUTOTUNE=1`` behaviour of main's ``AllReduce``; main otherwise + autotunes the strategy there). +* **Identity** with one rank or under attention DP (no TP group to reduce over). + +Fusions: ``NONE`` and the residual / plain RMSNorm epilogues the MNNVL kernel takes +(``AllReduceParams``), with main's output contract (a tensor, or a tuple for fusions with several +outputs). + +For ``layers.py`` an MNNVL instance also answers ``mnnvl_allreduce`` (itself), ``mapping``, +``dtype``, ``symm_mem_allreduce`` (None) and ``get_required_workspace_size``. + +Projections on K3-owned weights: tensor-parallel shards, BF16 GEMM and FP8 block-scale GEMM. + +Weight containers register the parameters under the names the checkpoint loader looks up and keep +this rank's shard. Their forward is one explicit kernel chain. A row-parallel BF16 projection given +an all-reduce (``layers.TpAllReduce``, as ``all_reduce``) reduces its partial with it, with +the caller's ``all_reduce_params`` (``enable_allreduce=False``: the partial; an RMSNorm fusion: +the fused outputs), as main's ``Linear(ROW, reduce_output=True)``. + +**Sharding** (:func:`shard_range` / :func:`load_shard`): rank ``r`` of ``tp`` takes a contiguous +range of the split dim (output features for ``COLUMN``, input features for ``ROW``), the first +``width % tp`` ranks one unit more; for FP8 block-scale weights the unit is a 128-wide block and +the block-scale grid follows the weight's range (``scale_span=128``). :func:`load_weight_shard` +is main's ceil-divided split (``ceil(width / tp)`` per rank), which the NVFP4 experts load with. + +**BF16** (:class:`Bf16Linear`): ``F.linear`` (cuBLASLt). + +**FP8 block scales** (:class:`Fp8BlockScaleLinear`): E4M3 ``weight [N, K]`` with one FP32 scale per +128 x 128 block (``weight_scale [ceil(N/128), ceil(K/128)]``). :meth:`Fp8BlockScaleLinear.prepare` +rounds the scales to UE8M0 (``resmooth_to_fp8_e8m0`` rescales the codes to match) and lays them +out for the GEMM of this GPU; :func:`fp8_block_scale_gemm` then runs, on BF16 rows: + +* SM107 (CuTe DSL Rubin): ``fp8_quantize_1x128_cutedsl_ue8m0`` (activation to E4M3 with K32 + R128c4 UE8M0 scales) and ``ops.gemm.k3_mxfp8_gemm`` (``ops/gemm/_mxfp8_gemm.py``), the weight + scale in the K32 R128c4 layout; +* SM100 / SM103: ``fp8_swap_ab_gemm(disable_ue8m0_cast=True)`` (DeepGEMM; activation quantized in + the kernel), the weight scale packed for DeepGEMM; +* either, where ``low_m_max`` covers the rows: ``trtllm::fp8_block_scale_low_m_gemm`` on the + DeepGEMM-packed scale (on SM107 a copy built from the K32 R128c4 one). + +RMSNorm on a K3-owned weight: the FlashInfer RMSNorm kernels through their catalog entries. + +``rms_norm(x, w, eps)`` is the catalog's ``flashinfer_rmsnorm`` (``w * x / rms(x)``, FP32 +statistics, the weight applied to the BF16-rounded normalized value); ``fused_add_rms_norm(x, r, w, +eps)`` is the catalog's ``flashinfer_fused_add_rmsnorm``, which writes ``r += x`` and +``x = rms_norm(r)`` in place. Both launch with PDL where ``TRTLLM_ENABLE_PDL`` is set (the op reads +it). + +:class:`RMSNormWeight` holds the ``[hidden]`` weight under the parameter name ``weight`` and the +epsilon as ``variance_epsilon`` (also readable as ``eps``), the names the checkpoint loader and the +kernels that take a norm's weight read. The KDA output norm is one too: its kernel is +``rms_norm_gated_token_major`` (``modules/mamba/layernorm_gated.py``), called by the mixer on +``weight`` / ``eps``. + +Depthwise causal short-convolution weight of a KDA q / k / v stream. + +:class:`ShortConvWeight` registers ``weight`` ``[channels, 1, kernel_size]`` and no bias, the +parameter a depthwise ``nn.Conv1d`` (``groups=channels``) registers, so the checkpoint's +``*_conv1d.weight`` loads by name. The convolution itself is the caller's kernel call on that +weight (``causal_conv1d_fn``, the KDA prefill / decode / verify kernels), with ``activation`` +(SiLU) applied inside it. + +Two launches on two CUDA streams, joined on the calling stream. + +``fork_join(fn0, fn1, fork, join, aux)`` runs ``fn0`` on the calling stream and ``fn1`` on ``aux``: + + fork.record() # calling stream + r0 = fn0() # calling stream + with torch.cuda.stream(aux): + fork.wait(); r1 = fn1(); join.record() + join.wait() # calling stream, before anything reads r1 + +Both run on the calling stream, in that order, unless ``do_multi_stream()`` (the engine's +per-thread flag, set while CUDA graphs are captured) holds and ``aux`` is given. Outside graph +capture the CUDA tensors ``fn1`` returns (directly or in nested tuples / lists / dict values) are +recorded on the calling stream, so the caching allocator does not hand their memory to later +``aux`` work while the calling stream still reads them. + +Token embedding on a K3-owned table. + +:class:`EmbeddingTable` holds ``weight`` ``[rows, hidden]``: the whole vocabulary, or with a +vocab-parallel mapping (``tp_size > 1``, no attention DP) rank ``r``'s rows +``[r * ceil(V / tp), min((r + 1) * ceil(V / tp), V))``, zero-padded to ``ceil(V / tp)`` rows. Its +forward is + + hit = (ids >= start) & (ids < end) # start, end: the rows this rank holds + out = F.embedding(hit * (ids - start * hit), weight) + out.masked_fill_(~hit, 0) + out = all_reduce(out) # vocab-parallel, unless the caller reduces + +so an id outside the table (the rejection sampler pads with -1) embeds to zeros and every vocab +row comes from exactly one rank. :func:`gather_embedding` is the one-launch lookup of a whole table +(``ops/embedding/k3_embedding_gather.py``, bitwise the masked lookup). + +Gated MLP on K3-owned BF16 weights: gate/up GEMM, gated activation, down GEMM, all-reduce. + + h = x @ gate_up_proj.weight.T # F.linear, [M, 2 * I_local] = [gate | up] + a = act(h) # family situ_and_mul or trtllm::silu_and_mul + y = a @ down_proj.weight.T # F.linear, [M, H] (row-parallel partial) + y = down_proj.all_reduce(y, params) # layers.TpAllReduce, tp > 1 only + +``gate_up_proj.weight`` holds this rank's rows ``[start, end)`` of the gate weight followed by the +same rows of the up weight (``start, end`` from :func:`layers.shard_range` over the +intermediate size); ``down_proj.weight`` the same columns of the down weight. ``params`` is the +caller's ``final_all_reduce_params`` (``enable_allreduce=False`` returns the partial). Activations: + +* ``situ``: the family ``situ_and_mul(h, beta, linear_beta)`` entry, the Triton SiTU kernel + (``beta * tanh(gate / beta) * sigmoid(gate) * (linear_beta * tanh(up / linear_beta) or up)``, + FP32 math); +* ``silu``: ``trtllm::silu_and_mul(h)``, the Triton SwiGLU kernel. + +:func:`subgroup_mapping` is the mapping of a TP subgroup of ``tp_size`` ranks: consecutive ranks +reduce together and the remaining factor is folded into PP, the subgroup the all-reduce runs over +when the MLP's TP size is smaller than the model's. + +Pushed MNNVL all-reduce, private to the Kimi K3 Rubin package. + +A producer kernel may store its BF16 output straight into the Lamport buffers of the MNNVL +all-reduce that follows it (the K3 low-M FP8 GEMM's push epilogue, ``kernels/k3_lamport_ar.py``). +:func:`lamport_push_target` hands the buffers to the producer and :func:`finish_lamport_push` +completes the all-reduce, bitwise like the C++ kernel the call replaces. + +This is the push API of the Rubin agentic branch (``189fe67b81``, ``distributed/ops.py``: +``LamportPushTarget``/``LamportPending`` and the ``AllReduce`` / +``MNNVLAllReduce`` methods ``lamport_push_target``, ``lamport_push_tensor`` and +``finish_lamport_push``), as free functions that take the all-reduce: ``layers.TpAllReduce`` +(or main's ``AllReduce``), read through ``mnnvl_allreduce`` (the MNNVL transport: its ``mapping``, +``dtype`` and ``get_required_workspace_size``) and ``symm_mem_allreduce``. They work on the shared +MNNVL workspace of ``get_or_scale_allreduce_mnnvl_workspace`` (same dict keys and the same 9-word +``buffer_flags`` layout). + +Differences from the branch, all forced by main: + +* The C++ one-/two-shot switch is ``distributed.ops._MNNVL_ONE_SHOT_THRESHOLD_BYTES`` + (``TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES``): the pushed call must reproduce the kernel main's + ``MNNVLAllReduce`` would have run. +* The device view of a workspace is built on the first eager call; during CUDA graph capture a + workspace without a view, or one that would have to grow, gives no target (the caller falls + back to the plain all-reduce). The eager warm-up forward that precedes every capture builds it. +""" + +from __future__ import annotations + +import dataclasses +import math +import os +import platform +from typing import Any, Callable, Dict, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allreduce import allreduce +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.activation.situ_and_mul import ( + situ_and_mul, +) +from tensorrt_llm._torch.distributed.mnnvl_memory import MnnvlMemory +from tensorrt_llm._torch.distributed.ops import ( + _MNNVL_ONE_SHOT_THRESHOLD_BYTES, + get_allreduce_workspace, + get_or_scale_allreduce_mnnvl_workspace, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._utils import get_sm_version, is_device_integrated +from tensorrt_llm.functional import AllReduceFusionOp, AllReduceParams, AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +MNNVL_DTYPES = (torch.float16, torch.bfloat16, torch.float32) +#: Fusions ``trtllm::mnnvl_fusion_allreduce`` runs (besides ``NONE``). +MNNVL_FUSION_OPS = frozenset( + { + AllReduceFusionOp.RMS_NORM, + AllReduceFusionOp.RESIDUAL_RMS_NORM, + AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8, + AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + AllReduceFusionOp.RESIDUAL_RMS_NORM_OUT_QUANT_FP8, + AllReduceFusionOp.RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4, + } +) + + +def mnnvl_serves(mapping: Mapping, dtype: torch.dtype, strategy: AllReduceStrategy) -> bool: + """Whether the all-reduce of ``mapping`` runs on MNNVL for ``strategy`` (``AUTO`` / ``MNNVL``).""" + if strategy not in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL): + return False + if dtype not in MNNVL_DTYPES or mapping.has_cp(): + return False + if os.environ.get("TLLM_TEST_MNNVL", "0") == "1": + return True + supported = MnnvlMemory.supports_mnnvl() and "aarch64" in platform.machine().lower() + return supported and (strategy == AllReduceStrategy.MNNVL or mapping.is_multi_node()) + + +def mnnvl_workspace_bytes(num_tokens: int, hidden: int, world: int, dtype: torch.dtype) -> int: + """Lamport buffer bytes one MNNVL call of ``[num_tokens, hidden]`` needs.""" + elem = torch.tensor([], dtype=dtype).element_size() + if num_tokens * hidden * world * elem <= _MNNVL_ONE_SHOT_THRESHOLD_BYTES: + return num_tokens * hidden * world * elem + return 2 * math.ceil(num_tokens / world) * world * hidden * elem + + +class TpAllReduce: + """The all-reduce over ``mapping``'s TP group; ``self(x)`` returns the reduced tensor.""" + + def __init__( + self, + mapping: Mapping, + dtype: torch.dtype, + strategy: AllReduceStrategy = AllReduceStrategy.AUTO, + ) -> None: + self.mapping = mapping + self.dtype = dtype + self.strategy = strategy + self.symm_mem_allreduce = None + self.transport = "identity" + self.workspace = None + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return + if mnnvl_serves(mapping, dtype, strategy): + self.transport = "mnnvl" + get_or_scale_allreduce_mnnvl_workspace(mapping, dtype) + return + if strategy in ( + AllReduceStrategy.MNNVL, + AllReduceStrategy.SYMM_MEM, + AllReduceStrategy.UB, + AllReduceStrategy.LOWPRECISION, + ): + raise ValueError(f"K3 all-reduce: strategy {strategy.name} is not served here") + self.transport = "op" + if strategy not in (AllReduceStrategy.NCCL, AllReduceStrategy.NCCL_SYMMETRIC): + self.workspace = get_allreduce_workspace(mapping) + + get_required_workspace_size = staticmethod(mnnvl_workspace_bytes) + + @property + def mnnvl_allreduce(self) -> Optional["TpAllReduce"]: + """Self on the MNNVL transport (what ``layers.py`` reads), else None.""" + return self if self.transport == "mnnvl" else None + + def __call__( + self, x: torch.Tensor, *, all_reduce_params: Optional[AllReduceParams] = None + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + params = all_reduce_params + if self.transport == "identity" or (params is not None and not params.enable_allreduce): + return x + x = x.contiguous() + params = params or AllReduceParams() + if self.transport == "mnnvl": + return self._mnnvl(x, params) + return allreduce( + x, + params.residual, + params.norm_weight, + params.scale, + params.bias, + self.workspace, + self.mapping.tp_group, + self.strategy, + params.fusion_op, + params.eps, + params.trigger_completion_at_end, + ) + + def _mnnvl( + self, x: torch.Tensor, params: AllReduceParams + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + fusion_op = params.fusion_op + is_fusion = fusion_op != AllReduceFusionOp.NONE + if is_fusion and fusion_op not in MNNVL_FUSION_OPS: + raise ValueError(f"the MNNVL all-reduce does not fuse {fusion_op!r}") + hidden = x.shape[-1] + num_tokens = x.numel() // hidden + need = mnnvl_workspace_bytes(num_tokens, hidden, self.mapping.tp_size, self.dtype) + if need >= 2**32 - 1: + raise ValueError( + f"MNNVL all-reduce of ({num_tokens}, {hidden}) over {self.mapping.tp_size} ranks " + f"needs {need} workspace bytes, past the kernel's uint32 limit" + ) + workspace = get_or_scale_allreduce_mnnvl_workspace( + self.mapping, self.dtype, buffer_size_bytes=need + ) + outputs = torch.ops.trtllm.mnnvl_fusion_allreduce( + x, + params.norm_weight, + params.residual, + params.eps, + workspace["uc_buffer"].view(self.dtype).view(3, -1), + workspace["buffer_flags"], + is_fusion, + params.scale, + int(fusion_op), + fp8_payload=False, + ) + if is_fusion and len(outputs) > 1: + return tuple(outputs) + return outputs[0] + + +FP8_BLOCK = 128 + + +def _rubin_mxfp8() -> bool: + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + return get_sm_version() == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE + + +def _split(width: int, parts: int, index: int) -> int: + return (width // parts) * index + min(width % parts, index) + + +def shard_range(width: int, tp_size: int, tp_rank: int, alignment: int = 1) -> Tuple[int, int]: + """``[start, end)`` of rank ``tp_rank`` over a split dim of ``width`` elements.""" + if alignment <= 1: + return _split(width, tp_size, tp_rank), _split(width, tp_size, tp_rank + 1) + if width % alignment: + raise ValueError(f"split dim {width} is not a multiple of {alignment}") + blocks = width // alignment + return ( + _split(blocks, tp_size, tp_rank) * alignment, + _split(blocks, tp_size, tp_rank + 1) * alignment, + ) + + +def load_shard( + tensor: torch.Tensor, + tp_size: int, + tp_rank: int, + mode: Optional[TensorParallelMode], + *, + features: Optional[int] = None, + alignment: int = 1, + scale_span: Optional[int] = None, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """This rank's slice of a checkpoint tensor (a weight, or with ``scale_span`` its block-scale + grid) on ``device``. ``features``: the full split-dim width of the weight (default: the + tensor's). A 1-D tensor of a ``ROW`` projection, or a split dim of width 1, is not split.""" + device = device if device is not None else tensor.device + if mode is None or tp_size <= 1: + return tensor.to(device) + dim = TensorParallelMode.split_dim(mode) + if (tensor.dim() == 1 and dim == 1) or tensor.shape[dim] == 1: + return tensor.to(device) + start, end = shard_range(features or tensor.shape[dim], tp_size, tp_rank, alignment) + if scale_span: + if start % scale_span or end % scale_span: + raise ValueError(f"shard [{start}, {end}) does not align with {scale_span}-wide scales") + start, end = start // scale_span, end // scale_span + return tensor.narrow(dim, start, end - start).to(device) + + +def load_weight_shard( + weight, + tensor_parallel_size: int = 1, + tensor_parallel_rank: int = 0, + tensor_parallel_mode: Optional[TensorParallelMode] = None, + device: torch.device = torch.device("cpu"), + return_slice_indices: bool = False, +) -> torch.Tensor: + """This rank's ceil-divided slice of ``weight`` (a tensor or a safetensors slice) on ``device``: + ``ceil(width / tp)`` elements of the split dim per rank, the last rank the remainder. A 1-D + tensor of a ``ROW`` projection, or a split dim of width 1, is not split. On an integrated GPU + the tensor stays on its own device. ``return_slice_indices`` (tensors only): also the slices.""" + # Skip device transfers on integrated GPUs to conserve shared memory + if weight.device.type != device.type and is_device_integrated(): + # For integrated GPU systems (e.g., DGX Spark), CPU and GPU share limited physical memory. + # Avoiding device transfers reduces memory consumption and unnecessary data copies, + # enabling support for larger models on memory-constrained systems. + logger.warning_once( + f"[load_weight_shard] Skipping device transfer from {weight.device} to {device} " + "on integrated GPU to conserve shared memory.", + key="load_weight_shard_skip_device_transfer_with_integrated_gpu", + ) + device = weight.device + if isinstance(weight, torch.Tensor): + tensor_shape = weight.shape + + def maybe_convert_to_torch_tensor(tensor: torch.Tensor, indices: list[slice] | None = None): + if indices is None: + # Avoid unnecessary copy + result = (tensor.to(device), [slice(d) for d in tensor.shape]) + else: + result = (tensor[indices].to(device), indices) + return result if return_slice_indices else result[0] + + # WAR to check whether it is a safetensor slice since safetensor didn't register the type to + # the module; safetensors slice, supports lazy loading, type(weight) is `builtin.PySafeSlice` + elif hasattr(weight, "get_shape"): + tensor_shape = weight.get_shape() + + def maybe_convert_to_torch_tensor( + tensor, indices: Union[slice, tuple[slice]] = slice(None) + ): + return tensor[indices].to(device) + else: + raise ValueError(f"unsupported weight type: {type(weight)}") + if tensor_parallel_mode is None or tensor_parallel_size <= 1: + return maybe_convert_to_torch_tensor(weight) + + split_dim = TensorParallelMode.split_dim(tensor_parallel_mode) + + if len(tensor_shape) == 1 and split_dim == 1: + return maybe_convert_to_torch_tensor(weight) + + width = tensor_shape[split_dim] + if width == 1: + return maybe_convert_to_torch_tensor(weight) + + slice_width = math.ceil(width / tensor_parallel_size) + slice_start = tensor_parallel_rank * slice_width + slice_end = min((tensor_parallel_rank + 1) * slice_width, width) + slice_obj = [slice(d) for d in tensor_shape] + slice_obj[split_dim] = slice(slice_start, slice_end) + return maybe_convert_to_torch_tensor(weight, tuple(slice_obj)) + + +class _ShardedLinear(nn.Module): + alignment = 1 + + def __init__( + self, + in_features: int, + out_features: int, + tp_size: int, + tp_rank: int, + tp_mode: Optional[TensorParallelMode], + ) -> None: + super().__init__() + self.tp_size, self.tp_rank, self.tp_mode = tp_size, tp_rank, tp_mode + self.full_in_features, self.full_out_features = in_features, out_features + self.in_features, self.out_features = in_features, out_features + if tp_size > 1 and tp_mode is not None: + dim_width = in_features if tp_mode == TensorParallelMode.ROW else out_features + start, end = shard_range(dim_width, tp_size, tp_rank, self.alignment) + if tp_mode == TensorParallelMode.ROW: + self.in_features = end - start + else: + self.out_features = end - start + + def split_features(self) -> int: + return ( + self.full_in_features + if self.tp_mode == TensorParallelMode.ROW + else self.full_out_features + ) + + def shard(self, tensor: torch.Tensor, scale_span: Optional[int] = None) -> torch.Tensor: + """This rank's slice of a full checkpoint tensor of the projection, on its device.""" + return load_shard( + tensor, + self.tp_size, + self.tp_rank, + self.tp_mode, + features=self.split_features(), + alignment=self.alignment, + scale_span=scale_span, + device=self.weight.device, + ) + + +def bf16_linear( + x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None +) -> torch.Tensor: + """``x @ weight.T (+ bias)``: cuBLASLt via ``F.linear``.""" + return F.linear(x, weight, bias) + + +class Bf16Linear(_ShardedLinear): + """BF16 (or FP16 / FP32) ``weight [out, in]`` (and ``bias [out]``) of one rank.""" + + def __init__( + self, + in_features: int, + out_features: int, + *, + bias: bool = False, + dtype: torch.dtype = torch.bfloat16, + tp_size: int = 1, + tp_rank: int = 0, + tp_mode: Optional[TensorParallelMode] = None, + all_reduce: Optional[Callable[..., torch.Tensor]] = None, + device: Optional[torch.device] = None, + ) -> None: + super().__init__(in_features, out_features, tp_size, tp_rank, tp_mode) + if all_reduce is not None and tp_mode != TensorParallelMode.ROW: + raise ValueError("only a row-parallel projection reduces its output") + self.all_reduce = all_reduce + self.reduce_output = all_reduce is not None + self.weight = nn.Parameter( + torch.empty(self.out_features, self.in_features, dtype=dtype, device=device), + requires_grad=False, + ) + if bias: + self.bias = nn.Parameter( + torch.empty(self.out_features, dtype=dtype, device=device), requires_grad=False + ) + else: + self.register_parameter("bias", None) + + def load(self, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> None: + """Copy this rank's shard of the full checkpoint ``weight`` (and ``bias``).""" + self.weight.data.copy_(self.shard(weight).to(self.weight.dtype)) + if bias is not None: + self.bias.data.copy_(self.shard(bias).to(self.bias.dtype)) + + def forward( + self, x: torch.Tensor, *, all_reduce_params: Optional[AllReduceParams] = None + ) -> torch.Tensor: + # Bias on rank 0 only for a row-parallel projection (its partial sums are all-reduced). + bias = self.bias if self.tp_mode != TensorParallelMode.ROW or self.tp_rank == 0 else None + out = bf16_linear(x, self.weight, bias) + if self.all_reduce is not None: + out = self.all_reduce(out, all_reduce_params=all_reduce_params) + return out + + +def prepare_fp8_block_weight( + weight: torch.Tensor, weight_scale: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """``(codes, scale)`` for :func:`fp8_block_scale_gemm` from checkpoint E4M3 codes and their FP32 + 128 x 128 grid: scales rounded to UE8M0 with the codes rescaled to match (in place), then the + scale laid out for this GPU's GEMM (K32 R128c4 UE8M0 bytes on SM107, DeepGEMM-packed UE8M0 + otherwise).""" + from tensorrt_llm.quantization.utils.fp8_utils import ( + resmooth_to_fp8_e8m0, + transform_k128_scales_to_cutedsl_mxfp8_layout, + transform_sf_into_required_layout, + ) + + weight, weight_scale = resmooth_to_fp8_e8m0( + weight.contiguous(), weight_scale.contiguous().float() + ) + n, k = weight.shape + if _rubin_mxfp8(): + return weight, transform_k128_scales_to_cutedsl_mxfp8_layout(weight_scale, mn=n, k=k) + return weight, transform_sf_into_required_layout( + weight_scale, mn=n, k=k, recipe=(1, 128, 128), is_sfa=False + ) + + +def low_m_gemm_rows(n: int, k: int) -> int: + """Largest M ``trtllm::fp8_block_scale_low_m_gemm`` serves for an ``[n, k]`` weight here (0: + none).""" + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + if not IS_CUTLASS_DSL_AVAILABLE: + return 0 + from tensorrt_llm._torch.kernels.blackwell import fp8_block_scale_low_m_gemm + + return fp8_block_scale_low_m_gemm.max_m_for(n, k) + + +def low_m_gemm_scale(weight_scale: torch.Tensor, n: int, k: int) -> torch.Tensor: + """The DeepGEMM-packed scale the low-M GEMM reads, from a :func:`prepare_fp8_block_weight` + scale (itself on SM100 / SM103, repacked from K32 R128c4 on SM107).""" + if not _rubin_mxfp8(): + return weight_scale + from tensorrt_llm._torch.kernels.blackwell import fp8_block_scale_low_m_gemm + + return fp8_block_scale_low_m_gemm.pack_block_scale_from_mxfp8_layout(weight_scale, n, k) + + +def fp8_block_scale_gemm( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + *, + low_m_max: int = 0, + low_m_scale: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """BF16 ``x [M, K] @ weight.T`` on :func:`prepare_fp8_block_weight` operands (see the module + docstring for the chain per GPU).""" + if ( + low_m_max + and 1 <= x.shape[0] <= low_m_max + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ): + return torch.ops.trtllm.fp8_block_scale_low_m_gemm(x, weight, low_m_scale) + if _rubin_mxfp8(): + from .ops.gemm.k3_mxfp8_gemm import k3_mxfp8_gemm + + act, act_scale = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(x) + return k3_mxfp8_gemm(act, weight, act_scale, weight_scale) + return torch.ops.trtllm.fp8_swap_ab_gemm( + x, weight, weight_scale, output_dtype=x.dtype, disable_ue8m0_cast=True + ) + + +class Fp8BlockScaleLinear(_ShardedLinear): + """FP8 block-scale ``weight`` / ``weight_scale`` of one rank (plus the unit ``input_scale`` / + ``inv_input_scale`` checkpoint FP8 projections also register), bias-free.""" + + alignment = FP8_BLOCK + + def __init__( + self, + in_features: int, + out_features: int, + *, + tp_size: int = 1, + tp_rank: int = 0, + tp_mode: Optional[TensorParallelMode] = None, + device: Optional[torch.device] = None, + ) -> None: + super().__init__(in_features, out_features, tp_size, tp_rank, tp_mode) + n, k = self.out_features, self.in_features + self.weight = nn.Parameter( + torch.empty(n, k, dtype=torch.float8_e4m3fn, device=device), requires_grad=False + ) + self.weight_scale = nn.Parameter( + torch.empty(-(-n // FP8_BLOCK), -(-k // FP8_BLOCK), dtype=torch.float32, device=device), + requires_grad=False, + ) + self.input_scale = nn.Parameter( + torch.tensor(1.0, dtype=torch.float32, device=device), requires_grad=False + ) + self.inv_input_scale = nn.Parameter( + torch.tensor(1.0, dtype=torch.float32, device=device), requires_grad=False + ) + self.register_parameter("bias", None) + self.prepared = False + self.low_m_max = 0 + self.low_m_scale: Optional[torch.Tensor] = None + + def load(self, weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + """Copy this rank's shard of the full checkpoint codes and their FP32 block-scale grid + (a modelopt ``[N/128, 1, K/128, 1]`` grid is squeezed first).""" + if weight_scale.dim() == 4: + weight_scale = weight_scale.squeeze(1).squeeze(-1) + self.weight.data.copy_(self.shard(weight)) + self.weight_scale.data.copy_(self.shard(weight_scale, scale_span=FP8_BLOCK).float()) + self.prepared = False + + def prepare(self, low_m: bool = True) -> None: + """Once after loading: :func:`prepare_fp8_block_weight` in place of the parameters, and + with ``low_m`` the rows the low-M GEMM serves and its scale.""" + if self.prepared: + return + weight, scale = prepare_fp8_block_weight(self.weight.data, self.weight_scale.data) + self.weight = nn.Parameter(weight, requires_grad=False) + self.weight_scale = nn.Parameter(scale, requires_grad=False) + n, k = self.out_features, self.in_features + self.low_m_max = low_m_gemm_rows(n, k) if low_m else 0 + self.low_m_scale = low_m_gemm_scale(scale, n, k) if self.low_m_max else None + self.prepared = True + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if not self.prepared: + raise RuntimeError("Fp8BlockScaleLinear.prepare() must run after loading") + x_flat = x.reshape(-1, x.shape[-1]) + out = fp8_block_scale_gemm( + x_flat, + self.weight, + self.weight_scale, + low_m_max=self.low_m_max, + low_m_scale=self.low_m_scale, + ) + return out.reshape(*x.shape[:-1], self.out_features) + + +_FLASHINFER_DTYPES = (torch.bfloat16, torch.float16) + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """``weight * x / sqrt(mean(x^2) + eps)`` over the last dim of a BF16 / FP16 ``x``.""" + if x.dtype not in _FLASHINFER_DTYPES: + raise TypeError(f"rms_norm takes BF16 / FP16 activations, got {x.dtype}") + return flashinfer_rmsnorm(x, weight, eps) + + +def fused_add_rms_norm( + x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float +) -> Tuple[torch.Tensor, torch.Tensor]: + """In place: ``residual += x``, then ``x = rms_norm(residual)``. Returns ``(x, residual)``.""" + if x.dtype not in _FLASHINFER_DTYPES: + raise TypeError(f"fused_add_rms_norm takes BF16 / FP16 activations, got {x.dtype}") + return flashinfer_fused_add_rmsnorm(x, residual, weight, eps) + + +class RMSNormWeight(nn.Module): + """The ``[hidden]`` RMSNorm weight (``weight``, initialized to ones) and its epsilon.""" + + def __init__( + self, + hidden_size: int, + eps: float, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.variance_epsilon = eps + + @property + def eps(self) -> float: + """``variance_epsilon``, under the name the gated (KDA ``o_norm``) and attention-residual + norm readers use.""" + return self.variance_epsilon + + def forward( + self, x: torch.Tensor, residual: Optional[torch.Tensor] = None + ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]: + if residual is None: + return rms_norm(x, self.weight, self.variance_epsilon) + return fused_add_rms_norm(x, residual, self.weight, self.variance_epsilon) + + +class ShortConvWeight(nn.Module): + """``weight [channels, 1, kernel_size]`` of one rank's channels.""" + + def __init__( + self, + channels: int, + kernel_size: int, + *, + activation: Optional[str] = "silu", + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.hidden_size, self.kernel_size, self.activation = channels, kernel_size, activation + self.weight = nn.Parameter( + torch.empty(channels, 1, kernel_size, dtype=dtype, device=device) + ) + self.register_parameter("bias", None) + + +def _record_on(obj: Any, stream: torch.cuda.Stream) -> None: + if isinstance(obj, torch.Tensor): + if obj.is_cuda: + try: + obj.record_stream(stream) + except RuntimeError as error: + # cudaMallocAsync asserts on memory it did not allocate (an external buffer wrapped + # as a tensor), which it never frees: nothing to record. + if ( + torch.cuda.memory.get_allocator_backend() != "cudaMallocAsync" + or "ptr not found in ptr_info" not in str(error) + ): + raise + elif isinstance(obj, (tuple, list)): + for item in obj: + _record_on(item, stream) + elif isinstance(obj, dict): + for item in obj.values(): + _record_on(item, stream) + + +def fork( + fn0: Callable[[], Any], + fn1: Callable[[], Any], + fork_event: torch.cuda.Event, + join_event: torch.cuda.Event, + aux_stream: Optional[torch.cuda.Stream], +) -> Tuple[Any, Any, Optional[torch.cuda.Event]]: + """``fn0()`` and ``fn1()`` and the event the calling stream must wait on before it reads + ``fn1()``'s outputs (None when both ran on the calling stream).""" + if not (do_multi_stream() and aux_stream is not None): + return fn0(), fn1(), None + fork_event.record() + result0 = fn0() + with torch.cuda.stream(aux_stream): + fork_event.wait() + result1 = fn1() + join_event.record() + if not torch.cuda.is_current_stream_capturing(): + _record_on(result1, torch.cuda.current_stream()) + return result0, result1, join_event + + +def fork_join( + fn0: Callable[[], Any], + fn1: Callable[[], Any], + fork_event: torch.cuda.Event, + join_event: torch.cuda.Event, + aux_stream: Optional[torch.cuda.Stream], +) -> Tuple[Any, Any]: + """``(fn0(), fn1())``, ``fn1`` on ``aux_stream`` where :func:`fork` forks; joined.""" + result0, result1, join = fork(fn0, fn1, fork_event, join_event, aux_stream) + if join is not None: + join.wait() + return result0, result1 + + +def masked_embedding(ids: torch.Tensor, weight: torch.Tensor, start: int, end: int) -> torch.Tensor: + """Rows ``ids - start`` of ``weight`` where ``start <= ids < end``, zeros elsewhere.""" + hit = (ids >= start) & (ids < end) + out = F.embedding(hit * (ids - start * hit), weight) + out.masked_fill_(~hit.unsqueeze(-1), 0) + return out + + +def gather_embedding(ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """:func:`masked_embedding` of a whole table in one Triton launch (PDL where enabled).""" + from .ops.embedding.k3_embedding_gather import k3_embedding_gather + + return k3_embedding_gather(ids, weight) + + +class EmbeddingTable(nn.Module): + """``weight`` of one rank and the lookup above.""" + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + *, + dtype: Optional[torch.dtype] = None, + mapping: Optional[Mapping] = None, + vocab_parallel: bool = False, + reduce_output: bool = True, + allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + mapping = mapping or Mapping() + self.num_embeddings, self.embedding_dim = num_embeddings, embedding_dim + self.vocab_parallel = ( + vocab_parallel and mapping.tp_size > 1 and not mapping.enable_attention_dp + ) + if self.vocab_parallel: + rows = math.ceil(num_embeddings / mapping.tp_size) + self.vocab_start = mapping.tp_rank * rows + self.vocab_end = min(self.vocab_start + rows, num_embeddings) + self.all_reduce = ( + TpAllReduce(mapping, dtype or torch.bfloat16, allreduce_strategy) + if reduce_output + else None + ) + else: + rows = num_embeddings + self.vocab_start, self.vocab_end = 0, num_embeddings + self.all_reduce = None + self.weight = nn.Parameter( + torch.empty(rows, embedding_dim, dtype=dtype, device=device), requires_grad=False + ) + + def load(self, weight: torch.Tensor) -> None: + """Copy this rank's rows of the full checkpoint table (zero padding rows past the vocab).""" + rows = self.vocab_end - self.vocab_start + self.weight.data[:rows].copy_(weight[self.vocab_start : self.vocab_end]) + self.weight.data[rows:].zero_() + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + out = masked_embedding( + ids, + self.weight, + self.vocab_start, + self.vocab_end if self.vocab_parallel else self.weight.shape[0], + ) + if self.all_reduce is not None: + out = self.all_reduce(out) + return out + + +def subgroup_mapping(mapping: Mapping, tp_size: int) -> Mapping: + """``mapping`` regrouped into TP groups of ``tp_size`` consecutive ranks.""" + if mapping.tp_size % tp_size: + raise ValueError(f"TP {mapping.tp_size} does not split into groups of {tp_size}") + pp_size = mapping.pp_size * mapping.tp_size // tp_size + return Mapping( + world_size=tp_size * pp_size, + rank=mapping.rank, + gpus_per_node=mapping.gpus_per_node, + tp_size=tp_size, + pp_size=pp_size, + enable_attention_dp=mapping.enable_attention_dp, + ) + + +def gated_activation( + h: torch.Tensor, activation: str, beta: float = 1.0, linear_beta: Optional[float] = None +) -> torch.Tensor: + """``act(gate) * up`` of a packed ``[..., 2 * I]`` = ``[gate | up]`` tensor.""" + if activation == "situ": + out = situ_and_mul(h.reshape(-1, h.shape[-1]), beta, linear_beta) + return out.reshape(*h.shape[:-1], h.shape[-1] // 2) + if activation == "silu": + return torch.ops.trtllm.silu_and_mul( + h, swiglu_limit=None, swiglu_alpha=None, swiglu_beta=None + ) + raise ValueError(f"unknown gated activation {activation!r}") + + +class GatedMlp(nn.Module): + """``gate_up_proj`` / ``down_proj`` BF16 weights of one rank and the chain above.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + *, + activation: str, + beta: float = 1.0, + linear_beta: Optional[float] = None, + dtype: torch.dtype = torch.bfloat16, + mapping: Optional[Mapping] = None, + reduce_output: bool = True, + allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + mapping = mapping or Mapping() + self.mapping = mapping + self.activation, self.beta, self.linear_beta = activation, beta, linear_beta + self.hidden_size, self.intermediate_size = hidden_size, intermediate_size + tp, rank = mapping.tp_size, mapping.tp_rank + self.inter_start, self.inter_end = shard_range(intermediate_size, tp, rank) + local = self.inter_end - self.inter_start + self.gate_up_proj = Bf16Linear(hidden_size, 2 * local, dtype=dtype, device=device) + self.down_proj = Bf16Linear( + intermediate_size, + hidden_size, + dtype=dtype, + tp_size=tp, + tp_rank=rank, + tp_mode=TensorParallelMode.ROW, + all_reduce=( + TpAllReduce(mapping, dtype, allreduce_strategy) + if reduce_output and tp > 1 + else None + ), + device=device, + ) + + def load(self, gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor) -> None: + """Copy this rank's rows of the full checkpoint gate / up ``[I, H]`` and columns of down + ``[H, I]``.""" + start, end = self.inter_start, self.inter_end + local = end - start + weight = self.gate_up_proj.weight + weight.data[:local].copy_(gate[start:end].to(weight.device, weight.dtype)) + weight.data[local:].copy_(up[start:end].to(weight.device, weight.dtype)) + self.down_proj.load(down) + + def forward( + self, x: torch.Tensor, final_all_reduce_params: Optional[AllReduceParams] = None + ) -> torch.Tensor: + h = bf16_linear(x, self.gate_up_proj.weight) + a = gated_activation(h, self.activation, self.beta, self.linear_beta) + return self.down_proj(a, all_reduce_params=final_all_reduce_params) + + +# Pushed all-reduces of at most this many bytes (num_tokens * hidden_dim * group_size * 2) use the +# one-shot transport, larger ones the two-shot transport: the transport of the C++ kernel the +# call replaces. The result always reproduces the C++ kernel main selects. +_K3_AR_PUSH_ONE_SHOT_MAX_BYTES = int(_MNNVL_ONE_SHOT_THRESHOLD_BYTES) +# Pushed all-reduces whose consumer finishes them (``deferred``: the K3 attn_res kernel reduces +# the slots itself, after o_proj and after the row-parallel MoE tail) use the one-shot transport up +# to this many bytes, the two-shot transport above. +_K3_FUSED_FINISH_ONE_SHOT_MAX_BYTES = _K3_AR_PUSH_ONE_SHOT_MAX_BYTES + + +@dataclasses.dataclass(frozen=True) +class LamportWorkspaceView: + """Device addresses of one MNNVL all-reduce workspace for pushed calls, as the routed experts' + output push takes them. + + ``buffer``: this rank's unicast mapping of the workspace; ``mc_base``: its multicast mapping; + ``peer_ptrs``: int64 CUDA tensor with every rank's unicast mapping (this rank's own at index + ``rank``); ``flags``: the uint32 rotation state shared with the C++ kernels; + ``bytes_per_buffer``: size of one of the three Lamport buffers. + """ + + world: int + rank: int + buffer: torch.Tensor + mc_base: int + peer_ptrs: torch.Tensor + flags: torch.Tensor + bytes_per_buffer: int + + @property + def uc_local(self) -> int: + """Device address of ``buffer``.""" + return self.buffer.data_ptr() + + +# Device views of the workspaces for pushed all-reduces, by the address of the workspace's buffer +# (a scaled workspace is a new buffer). +_LAMPORT_VIEWS: Dict[int, object] = {} +# The latest shared MNNVL workspace an eager call here got, by all-reduce mapping (workspaces only +# grow: a later one is at least as large). +_SEEN_WORKSPACES: Dict[Mapping, Dict[str, Any]] = {} + + +class LamportPushTarget: + """One MNNVL all-reduce whose input the producer pushes. + + Made by :func:`lamport_push_target` for a single call. The producer stores its BF16 + ``[num_tokens, hidden]`` output into the Lamport slots of ``workspace`` (a + ``k3_lamport_ar.LamportWorkspace``) with the ``one_shot`` or two-shot transport and sets + ``pushed``; :func:`finish_lamport_push` then writes the all-reduce result into ``output``, + reproducing the C++ one-shot kernel if ``cpp_one_shot``, else the two-shot kernel. + """ + + def __init__(self, workspace, one_shot: bool, cpp_one_shot: bool, output: torch.Tensor): + self.workspace = workspace + self.one_shot = one_shot + self.cpp_one_shot = cpp_one_shot + self.output = output + self.pushed = False + # Set by ``lamport_push_target(deferred=True)``: the one-shot transport, finished by the + # consumer (``LamportPending``). + self.deferred = False + # The call carries all-gather rows (``k3_lamport_ar`` "ag", one-shot transport): its + # result is ``bf16(y + bf16(sum of the slots))``. + self.ag = False + + +class LamportPending: + """A pushed all-reduce whose finish is left to its consumer + (``finish_lamport_push(..., deferred=True)``). + + Every rank has stored its partial into the slots of the call (one-shot or two-shot + transport), but no kernel has reduced them or advanced the workspace's rotation state: the + consumer's kernel does both (``k3_attn_res_lamport``), or ``materialize`` runs + the finish kernel into a tensor. Exactly one of the two must run, before the next all-reduce on + the workspace. + """ + + def __init__(self, target: LamportPushTarget): + self.target = target + self.workspace = target.workspace + self.shape = tuple(target.output.shape) + + @property + def finish_args(self) -> Tuple[int, int, bool, int, int]: + """``(finish_world, finish_chunk, finish_two_shot, mc_base, rank)`` of + ``k3_attn_res_lamport``: the ranks per partial sum of the C++ kernel the result + reproduces, the transport and the workspace addresses.""" + from .kernels import k3_lamport_ar + + ws = self.workspace + return ( + ws.world, + k3_lamport_ar.reduce_chunk(ws.world, self.target.cpp_one_shot), + not self.target.one_shot, + ws.mc_base, + ws.rank, + ) + + @property + def finish_ag(self) -> bool: + """Whether the call carries all-gather rows (``finish_ag`` of ``k3_attn_res_lamport``).""" + return getattr(self.target, "ag", False) + + def materialize(self) -> torch.Tensor: + """Finish the all-reduce into a BF16 tensor (the finish kernel on the current stream).""" + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + from .ops.comm.k3_lamport_finish import k3_lamport_finish + + return k3_lamport_finish( + self.workspace, + self.target.output, + self.target.one_shot, + use_pdl=get_env_enable_pdl(), + cpp_one_shot=self.target.cpp_one_shot, + ag=self.finish_ag, + ) + + +def _lamport_view(mnnvl: Any, workspace: dict): + """The ``k3_lamport_ar.LamportWorkspace`` of an MNNVL workspace (built once per buffer).""" + key = workspace["uc_buffer"].data_ptr() + view = _LAMPORT_VIEWS.get(key) + if view is None: + from .kernels import k3_lamport_ar + + view = k3_lamport_ar.LamportWorkspace.from_mnnvl( + workspace, mnnvl.mapping.tp_size, mnnvl.mapping.tp_rank + ) + _LAMPORT_VIEWS[key] = view + logger.info_once( + f"[MNNVL AllReduce] pushed all-reduce: Lamport workspace view " + f"for tp_size {mnnvl.mapping.tp_size} ({view.bytes_per_buffer} bytes per buffer)", + key=("k3_ar_push_view", mnnvl.mapping.tp_size), + ) + return view + + +def mnnvl_lamport_push_target( + mnnvl: Any, + num_tokens: int, + hidden_dim: int, + one_shot: Optional[bool] = None, +) -> Optional[LamportPushTarget]: + """``MNNVLAllReduce.lamport_push_target`` of the branch: a target for a plain all-reduce of + BF16 ``[num_tokens, hidden_dim]`` on the workspace of ``mnnvl`` (an all-reduce's MNNVL + transport, ``mnnvl_allreduce``), or None. ``one_shot`` forces the transport instead of the + size rule.""" + from .kernels import k3_lamport_ar + + world = mnnvl.mapping.tp_size + nbytes = num_tokens * hidden_dim * world * 2 + cpp_one_shot = nbytes <= _MNNVL_ONE_SHOT_THRESHOLD_BYTES + if one_shot is None: + one_shot = nbytes <= _K3_AR_PUSH_ONE_SHOT_MAX_BYTES + need = max( + mnnvl.get_required_workspace_size(num_tokens, hidden_dim, world, mnnvl.dtype), + k3_lamport_ar.required_bytes(num_tokens, hidden_dim, world, one_shot), + ) + if need >= 2**32 - 1: + return None + capturing = torch.cuda.is_current_stream_capturing() + if capturing: + # Main raises when a workspace is created or grows during capture, and building the view + # is a host-to-device copy: both must have happened in the eager warm-up of this shape. + seen = _SEEN_WORKSPACES.get(mnnvl.mapping) + if ( + seen is None + or seen["buffer_size_bytes"] < need + or seen["uc_buffer"].data_ptr() not in _LAMPORT_VIEWS + ): + return None + workspace = get_or_scale_allreduce_mnnvl_workspace( + mnnvl.mapping, mnnvl.dtype, buffer_size_bytes=need + ) + if capturing: + if workspace["uc_buffer"].data_ptr() not in _LAMPORT_VIEWS: + # Another all-reduce grew the workspace after the warm-up's push here. + return None + else: + _SEEN_WORKSPACES[mnnvl.mapping] = workspace + view = _lamport_view(mnnvl, workspace) + output = torch.empty( + num_tokens, hidden_dim, dtype=mnnvl.dtype, device=workspace["uc_buffer"].device + ) + return LamportPushTarget(view, one_shot, cpp_one_shot, output) + + +def mnnvl_workspace_mpi_comm(workspace: Dict[str, object]): + """The MPI communicator of a shared MNNVL all-reduce workspace + (``get_or_scale_allreduce_mnnvl_workspace``), for the package's private Lamport workspaces: + ``McastGPUBuffer`` takes its Fortran handle. The workspace keys it ``"comm"``; under a non-MPI + orchestrator that entry is the TP ProcessGroup, which these workspaces do not support.""" + comm = workspace.get("comm") + if comm is None or not hasattr(comm, "py2f"): + raise RuntimeError( + "Kimi K3 private MNNVL workspaces need the MPI communicator of the shared MNNVL " + "workspace (MPI orchestrator, workspace attached)" + ) + return comm + + +def _mnnvl_finish( + target: LamportPushTarget, deferred: bool = False +) -> Union[torch.Tensor, LamportPending]: + """``MNNVLAllReduce.finish_lamport_push`` of the branch.""" + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + from .ops.comm.k3_lamport_finish import k3_lamport_finish + + if deferred: + return LamportPending(target) + return k3_lamport_finish( + target.workspace, + target.output, + target.one_shot, + use_pdl=get_env_enable_pdl(), + cpp_one_shot=target.cpp_one_shot, + ) + + +def lamport_push_target( + allreduce: Any, + num_tokens: int, + hidden_dim: int, + all_reduce_params: Optional[AllReduceParams] = None, + deferred: bool = False, +) -> Optional[LamportPushTarget]: + """A ``LamportPushTarget`` if ``allreduce`` would run the MNNVL all-reduce on a BF16 + ``[num_tokens, hidden_dim]`` input with these parameters (plain sum, no fusion); else None. + + A producer that stores its output into the target (and sets ``pushed``) replaces the input of + this all-reduce; :func:`finish_lamport_push` then returns the same result the all-reduce would. + Targets must be finished in the order they were pushed, by every rank. ``deferred`` sets + ``target.deferred`` for ``finish_lamport_push(deferred=True)``, whose consumer finishes the + call; it takes the one-shot transport where the call is at most + ``_K3_FUSED_FINISH_ONE_SHOT_MAX_BYTES`` (``num_tokens * hidden_dim * world * 2``), the + two-shot one above. + """ + mnnvl = getattr(allreduce, "mnnvl_allreduce", None) + if ( + mnnvl is None + or allreduce.mapping.tp_size == 1 + or getattr(allreduce, "symm_mem_allreduce", None) is not None + ): + return None + if all_reduce_params is not None and ( + not all_reduce_params.enable_allreduce + or all_reduce_params.fusion_op != AllReduceFusionOp.NONE + ): + return None + if mnnvl.dtype != torch.bfloat16 or num_tokens <= 0 or hidden_dim % 8: + return None + if not deferred: + return mnnvl_lamport_push_target(mnnvl, num_tokens, hidden_dim) + one_shot = ( + num_tokens * hidden_dim * allreduce.mapping.tp_size * 2 + <= _K3_FUSED_FINISH_ONE_SHOT_MAX_BYTES + ) + target = mnnvl_lamport_push_target(mnnvl, num_tokens, hidden_dim, one_shot=one_shot) + if target is not None: + target.deferred = True + return target + + +def lamport_push_tensor(allreduce: Any, target: LamportPushTarget, x: torch.Tensor) -> None: + """Push the contiguous BF16 ``[num_tokens, hidden_dim]`` tensor ``x`` into ``target`` (a copy + kernel), for a producer without the push epilogue.""" + del allreduce # the target carries its workspace; kept for the branch's call shape + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + from .ops.comm.k3_lamport_push import k3_lamport_push + + k3_lamport_push(target.workspace, x.contiguous(), target.one_shot, use_pdl=get_env_enable_pdl()) + target.pushed = True + + +def finish_lamport_push( + allreduce: Any, + target: LamportPushTarget, + deferred: bool = False, +) -> Union[torch.Tensor, LamportPending]: + """The all-reduce result of a target the producer has pushed: a tensor, or with ``deferred`` + (where ``target.deferred``) a ``LamportPending`` that the next kernel finishes.""" + del allreduce # the target carries its workspace; kept for the branch's call shape + if not target.pushed: + raise RuntimeError("the producer did not push this all-reduce") + return _mnnvl_finish(target, deferred and target.deferred) + + +__all__ = [ + "LamportPending", + "LamportPushTarget", + "finish_lamport_push", + "lamport_push_target", + "lamport_push_tensor", + "mnnvl_lamport_push_target", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla.py new file mode 100644 index 000000000000..af6d2dfa70c4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla.py @@ -0,0 +1,2938 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MLA on Rubin: K3-owned weights and explicit kernel chains. + +K3 MLA is DeepSeek-V3-style multi-latent attention with K3's deltas: + +* **NoPE.** The MLA kernels' RoPE table is an identity table, so their rotation is a copy. +* **Output gate before ``o_proj``.** ``sigmoid(g_proj(hidden_states)) * attn_output``. +* **Softmax scale** ``(qk_nope + qk_rope) ** -0.5 = 192 ** -0.5`` (``q_scaling`` 1). + +This module holds the weights and the forward of every phase (``K3MLAAttention``), the context +and generation attention chains, K3's split-K context FMHA, the fused Helix tail and the layer's +runtime (``K3MLARuntime``); ``ops/attention/_fp4_mla_decode`` holds K3's FP4 decode and +``kernels/mla/q_norm_pdl`` the PDL RMSNorm. +``TrtllmAttention`` and the KV cache managers are the runtime. + +Kimi K3 MLA: K3-owned weights and the explicit kernel chain of every phase. + +Weights (``layers`` containers, the checkpoint loader's names): the fused q_a / kv_a projection +``kv_a_proj_with_mqa``, ``q_a_layernorm``, ``q_b_proj`` (head-sharded), ``kv_a_layernorm``, +``kv_b_proj`` (head-sharded, ``[all K | all V]`` rows), the output gate ``g_proj`` and ``o_proj`` +(sharded over TP x CP), and the absorption weights the loader derives from kv_b +(``k_b_proj_trans`` / ``v_b_proj``, their FP8 scales and BF16 dequantized copies, and under Helix +``helix_v_b_proj_full``). FP8 projections are ``Fp8BlockScaleLinear`` until the loader swaps them +for its FP8 weight-read GEMM; BF16 projections are ``Bf16Linear``. The context and generation +attention runtime is ``TrtllmAttention`` (``mha``: 192-dim context heads, ``mqa``: the 576-dim +absorbed latent head), whose RoPE table is an identity table owned here (K3 is NoPE). + +One forward of ``T`` tokens (the first ``T_ctx`` context rows): + +1. **Front**: the fused q_a / kv_a projection (on the MXFP8 layer input the decoder layer hands + over, where its GEMM takes it; a BF16 one is ``trtllm::cublas_mm``), the q_a RMSNorm + (FlashInfer, or with ``features.mla_q_norm_pdl`` its PDL-at-launch variant) beside the kv_a + RMSNorm on the attention aux stream, then q_b (the context rows of a Helix rank: its head range + only) beside the latent row ``concat(kv_a_norm, k_pe)``. +2. **Context** rows: ``forward_context``. +3. **Generation** rows: ``forward_absorption_generation``. +4. **Tail**: ``o_proj(attn * sigmoid(g_proj(x)))``; the TP / CP sum of o_proj's partials is the + caller's (``K3MLARuntime``). + +Generation-only Helix batches of at most 128 tokens take the side-gate forward (the zero-KV mask +and the output gate on a side stream, o_proj gated in its GEMM); FP4-KV layers of at most 16 tokens +the fused query (``forward_gen_q_fused``). + +The attention part runs as :func:`mla_custom_op_inplace` under ``eager_on_graph`` (a breakable CUDA +graph breaks around it) where the layer is registered in the model's ``mla_layers``. + +Kimi K3 MLA generation: the absorbed decode as one kernel chain per profile. + +Every function takes the layer (:class:`~.kimi_k3_mla_attention.K3MLAAttention`) for its +weights and attention backends. The chain of one generation batch of ``T`` tokens: + +* **W_UK and the cache update** (:func:`forward_absorption_generation`): the absorption BMM of the + query's NoPE half into the latent query (BF16 ``trtllm::bmm_out`` on ``k_b_proj_trans``, or + ``torch.bmm`` on its dequantized copy for FP8 KV-B) beside + ``TrtllmAttention.mla_rope_generation`` (identity RoPE, the cache append and, on an FP4 cache, + the FP4 query quantization), on the attention aux stream for a BF16 cache. +* **Attention** (:func:`gen_attention`): K3's packed FP4 decode (``ops/attention/_fp4_mla_decode``) on an FP4 + cache with the CuTe DSL FP4 MLA backend, else ``TrtllmAttention.forward``. +* **Helix CP** (:func:`attn_forward_gen`): the V up-projection of this rank's partials + (``trtllm::bmm_out`` on ``helix_v_b_proj_full``), then the fifo-v2 MNNVL all-to-all + (``HelixAllToAllNative.alltoall_native``, which zeroes the rows this rank holds no KV for) and + ``trtllm::helix_post_process_native``; or, where ``helix_tail`` serves the batch, the decode + with its split reduction deferred and one ``k3_helix_tail`` kernel (fold, V up-projection, + exchange, combine). +* **W_UV** without Helix: ``trtllm::bmm_out`` on ``v_b_proj`` (BF16) or ``torch.bmm`` on + ``v_b_proj_dequant`` (FP8) into the layer output. + +The fused query (:func:`forward_gen_q_fused`, generation-only Helix batches of FP4-KV layers): the +fused q_a / kv_a projection, then ``k3_mla_latent_row`` and the KV-only FP4 cache append +(``fp4_mla_kv_append``) on the aux stream beside ``k3_mla_q_fused`` (q_a RMSNorm, q_b, W_UK, RoPE +and the FP4 query quantization in one kernel), then the decode and the Helix tail as above. + +The Kimi-K3 Helix generation tail as one kernel. + +In a Helix generation batch whose V up-projection runs before the exchange, +:class:`K3HelixTail` replaces the packed split-KV reduction, +the V up-projection BMM, the FIFO all-to-all and the combine by one launch of +``cute_dsl_kernels/blackwell/k3_helix_tail.py``: the decode's FP32 split partials are folded by +the senders (the reduction is deferred with ``DeferredPackedReduce``), every head's rows are +projected and pushed into the head owner's Lamport receive buffer (MNNVL fabric memory), and the +receivers combine them into the ``[T, H / P * 128]`` o_proj input. Decode backends that store +their rows directly (no split partials, the FP8 layers) take the same kernel without the fold. + +Sizes and features: + +* ``_TAIL_MB``: receive buffer per rank (two halves), 16 MiB; a batch whose records do not fit a + half takes the unfused path. The kernel runs one CTA per SM: every CTA folds, projects and + sends, then combines tokens. +* o_proj's operand: where the side-stream output gate is joined before the tail, the receivers + also store the MXFP8 form of ``bf16(out * gate)`` (``mx_out``), and o_proj runs the MXFP8 + low-M GEMM with the all-reduce push epilogue on it (bitwise the gated low-M GEMM + push). +* ``features.ull2_mla_tail``: the senders stage the task plan before the grid wait, fold any + number of split tasks with one code path and pad the latent rows for the V up-projection + (``k3_helix_tail`` ``ull2``; bitwise the default). Read at every call; FP16 split partials + (``features.fp4_mla_f16_partials``) always take that path. For a batch of at most + :data:`_LATE_TAIL_MAX_TOKENS` tokens the decode before the tail then does not release it early + (:func:`decode_releases_tail_early`). + +Kimi K3 MLA context attention: one kernel chain per cache state. + +Every function takes the layer (:class:`~.kimi_k3_mla_attention.K3MLAAttention`) for its +weights and its context attention backend ``mha`` (``TrtllmAttention``, ``H`` context heads of +192 / 128 dims). :func:`forward_context` picks the chain: + +* **No cached prefix** (:func:`context_default`): kv_b over the batch's latent rows (a BF16 kv_b on + SM100f: two per-head ``trtllm::bmm_out`` writing K-nope into the strided K buffer and V into a + ``[K-nope | V]``-strided buffer; otherwise the kv_b GEMM and a copy of K-nope into K), then + ``TrtllmAttention.forward`` with ``latent_cache``, which applies the identity RoPE, appends the + cache and runs the context FMHA. Owner-only Helix prefill appends the KV it owns first + (``mla_rope_append_paged_kv_assign_q``) and passes K with its K-pe and no ``latent_cache``. +* **Chunked prefill over an FP8 / BF16 cache** (:func:`context_chunked_prefill`): + ``mla_rope_append_paged_kv_assign_q``; then per cached-prefix partition + ``load_chunked_kv_cache_for_mla``, the K / V expansion and ``TrtllmAttention.forward`` with a full + mask into a temporary output merged with ``merge_attention_for_mla``; the current chunk last, + causal. On an FP8 cache with the FP8 context FMHA (``fp8_cached_kv_context``) Q is quantized once + (``static_quantize_e4m3_per_tensor``) and each expansion is kv_b + ``mla_kv_to_fp8`` (FP8 K / V in + one pass), handed to the op as already-quantized inputs; otherwise kv_b and the BF16 K-nope / + K-pe concat. +* **Cached prefix without chunked prefill** (:func:`context_cached_kv`): the same expansion over the + whole paged cache (``load_paged_kv_cache_for_mla``) in one ``TrtllmAttention.forward``. +* **Chunked prefill over an FP4 cache**: the Rubin FP8 split-K context FMHA + (``forward_context_split_k_fp4``). + +Kimi K3 MLA cached-prefix context on Rubin: the FP8 split-K context FMHA. + +The chunked prefill of an MLA layer with an NVFP4 latent cache (:func:`forward_context_split_k_fp4`, +where :func:`uses_k3_split_k_context` holds: SM107, matching Q / KV heads, head dimensions +128 / 64 / 128, the CuTe DSL kernel imports, BF16 Q and output without an output scale): + +* ``TrtllmAttention.append_fp4_mla_context_cache`` rotates Q and the current K (identity RoPE) + and appends the chunk to the FP4 cache; +* each cached-prefix partition, and then the current chunk (causal), runs the Rubin FP8 split-K + context FMHA (``ops/attention/k3_context_fmha``, ``kernels/k3_context_fmha.py``), which reads + K-nope ``[T, H, 128]``, the shared K-pe ``[T, 1, 64]`` and V as they are and merges the + partition into the output in its epilogue (``merge_op_tensor`` selects copy / merge / skip per + request); +* its inputs are unit-scale E4M3: Q once per layer call, each partition's kv_b output and K-pe. + The Triton kernels of ``ops/attention/static_quantize_e4m3`` quantize Q and the K-pe rows (the + current chunk's read in place from ``latent_cache``; ``static_quantize_e4m3_per_tensor`` for K-pe + rows they do not support), with the bytes of ``static_quantize_e4m3_per_tensor``; +* where kv_b is the FP8 block-scale GEMM with its weight scale in the SM107 CuTe MXFP8 layout + (:func:`kv_b_takes_mxfp8`), each cached-prefix partition is gathered from the FP4 cache straight + into MXFP8 (``ops/attention/fp4_mla_gather_mxfp8``), the current chunk is quantized with + ``fp8_quantize_1x128_packed_ue8m0``, and kv_b runs as ``cute_dsl_mxfp8_gemm_rubin`` (BF16 + output); otherwise the BF16 gather and kv_b run. Either output is quantized with + ``static_quantize_e4m3_per_tensor``. + +Kimi K3 NVFP4 on Rubin (SM107): the MLA layer. + +``K3MLARuntime`` holds an MLA decoder layer's attention module (``mla.K3MLAAttention``) +and applies its TP output reduction. Also here: the quantization of a checkpoint attention +projection (``resolve_attention_quant_config``), the MLA layers whose KV cache is FP8 under +an NVFP4 KV cache, and the RoPE table size override. +""" + +from __future__ import annotations + +import contextlib +import functools +from typing import Callable, Dict, Optional, Tuple + +import torch +import triton + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.bmm_out import bmm_out +from tensorrt_llm._torch.attention.backends import ( + AttentionForwardArgs, + AttentionInputType, + AttentionMetadata, + TrtllmAttention, + TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.config import ( + FP4_MLA_K_RESIDUAL_DIM, + _fp4_mla_attention_backend, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.layout import ( + _get_fp4_mla_global_scale, + _get_fp4_mla_kv_cache_tensors, + _validate_fp4_mla_kv_storage_shape, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.metadata import ( + _materialize_fp4_mla_device_page_table_for_forward, +) +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.distributed import HelixAllToAllNative +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._torch.utils import maybe_compiled_cat, maybe_compiled_copy_ +from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm.logger import logger + +from .features import active as _features +from .kernels import fp8_static_quant +from .kernels.mla.mla_kv_fp8 import mla_kv_to_fp8 +from .layers import fork_join, rms_norm +from .ops.attention import _fp4_mla_decode as fp4_mla_decode +from .ops.attention import k3_fp4_mla_decode as _decode_op +from .ops.attention.k3_context_fmha import k3_context_fmha +from .ops.attention.static_quantize_e4m3 import static_quantize_e4m3 + +try: + from cuda.bindings import driver as cuda +except ImportError: + from cuda import cuda + +import copy +import os +import weakref + +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch import kimi_k3_cache_policy as _kimi_k3_cache_policy +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.custom_ops import BufferKind +from tensorrt_llm._torch.distributed.mnnvl_memory import HelixCpMnnvlMemory, MnnvlMemory +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph +from tensorrt_llm._torch.utils import AuxStreamType, get_model_extra_attrs +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + +from .layers import ( + Bf16Linear, + Fp8BlockScaleLinear, + LamportPending, + RMSNormWeight, + TpAllReduce, + finish_lamport_push, + lamport_push_target, + lamport_push_tensor, +) + +# Head dimensions (K-nope, K-pe, V) the split-K context FMHA is built for. +_SPLIT_K_HEAD_DIMS = (128, 64, 128) + + +def split_k_context_unavailable_reason(attn) -> Optional[str]: + """Why ``attn`` cannot run the Rubin FP8 split-K context FMHA, or None if it can.""" + sm_version = get_sm_version() + if sm_version != 107: + return f"requires SM107, got SM{sm_version}" + if attn.num_heads <= 0 or attn.num_heads != attn.num_kv_heads: + return ( + "requires matching positive Q/KV head counts per rank, " + f"got {attn.num_heads}/{attn.num_kv_heads}" + ) + head_dims = (attn.qk_nope_head_dim, attn.qk_rope_head_dim, attn.v_head_dim) + if head_dims != _SPLIT_K_HEAD_DIMS: + return f"requires head dimensions (K-nope, K-pe, V)={_SPLIT_K_HEAD_DIMS}, got {head_dims}" + from .kernels import k3_context_fmha_op + + if k3_context_fmha_op.IS_K3_CONTEXT_FMHA_AVAILABLE: + return None + if k3_context_fmha_op.K3_CONTEXT_FMHA_IMPORT_ERROR is not None: + return f"CuTe DSL kernel import failed: {k3_context_fmha_op.K3_CONTEXT_FMHA_IMPORT_ERROR!r}" + return "CuTe DSL kernel lacks split-K epilogue merge support" + + +def supports_k3_fp8_split_k_context(attn) -> bool: + """Whether the attention backend ``attn`` can run the Rubin FP8 split-K context FMHA.""" + reason = split_k_context_unavailable_reason(attn) + if reason is None: + return True + logger.info_once( + f"Kimi K3 MLA: Rubin FP8 split-K context FMHA disabled: {reason}; NVFP4 KV-cache layers " + "cannot attend a cached prefix (chunked prefill raises).", + key=("k3_mla_ctx_split_k_unavailable", reason), + ) + return False + + +def run_k3_split_k_context_partition( + attn, + q: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + softmax_stats: torch.Tensor, + merge_op: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + *, + max_q_len: int, + max_kv_len: int, + is_causal: bool, +) -> None: + """Attend FP8 Q to one FP8 KV partition and merge it into ``output`` / ``softmax_stats``. + + ``q`` ``[T_q, H * 192]``, ``k_nope`` / ``v`` ``[T_kv, H * 128]`` with unit inner stride, + ``k_pe`` ``[T_kv, 64]``, ``output`` ``[T_q, H * 128]`` BF16, ``softmax_stats`` + ``[T_q, H, 2]`` FP32; the softmax scale is the attention backend's. + """ + from .kernels import k3_context_fmha_op + + sm_scale = 1.0 / (attn.q_scaling * (attn.qk_nope_head_dim + attn.qk_rope_head_dim) ** 0.5) + num_heads = attn.num_heads + hints = "on" if k3_context_fmha_op.K3_CONTEXT_FMHA_SCHEDULING_HINTS else "off" + logger.info_once( + "Kimi K3 MLA context FMHA: k3_context_fmha (Rubin FP8 split-K, epilogue " + f"merge) serves {'causal current-chunk' if is_causal else 'cached-prefix'} partitions, " + f"{num_heads} heads per rank, CuTe DSL scheduling hints {hints}.", + key=("k3_mla_ctx_split_k_fmha", is_causal, num_heads), + ) + k3_context_fmha( + q.view(q.shape[0], num_heads, 192), + k_nope.view(k_nope.shape[0], num_heads, 128), + k_pe.view(k_pe.shape[0], 1, 64), + v.view(v.shape[0], num_heads, 128), + output.view(q.shape[0], num_heads, 128), + cu_q_seqlens, + cu_kv_seqlens, + softmax_stats, + merge_op, + max_q_len, + max_kv_len, + sm_scale, + is_causal, + ) + + +def load_fp4_mla_chunked_kv_cache_mxfp8( + metadata, + layer_idx: int, + *, + num_ctx_cached_tokens: int, + cu_chunked_seq_len: torch.Tensor, + chunked_global_offset: torch.Tensor, + chunked_max_seq_len: int, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]: + """Gather one cached-prefix partition of dense FP4 MLA V2 storage as MXFP8 kv_b_proj input. + + The partition and metadata are those of ``load_fp4_mla_chunked_kv_cache``. Returns + ``((E4M3 [M, 512], UE8M0 [pad_up(M, 128) * 16]), E4M3 [M, 64])``: the BF16-rounded latent + quantized 1x128 with scales in the K32 R128c4 layout ``cute_dsl_mxfp8_gemm_rubin`` reads, and + the BF16-rounded RoPE part plus its residual in unit-scale E4M3. + """ + if kv_lora_rank != 512 or qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM: + raise ValueError( + "MXFP8 gather requires kv_lora_rank=512 and " + f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got {kv_lora_rank} and {qk_rope_head_dim}." + ) + if num_ctx_cached_tokens < 0: + raise ValueError("FP4 MLA chunk gather token count must be non-negative.") + if chunked_max_seq_len < 0: + raise ValueError("FP4 MLA chunk gather max sequence length must be non-negative.") + if num_ctx_cached_tokens > 0 and chunked_max_seq_len == 0: + raise ValueError("A non-empty FP4 MLA chunk gather requires a positive max length.") + num_contexts = int(metadata.num_contexts) + if num_ctx_cached_tokens > 0 and num_contexts <= 0: + raise ValueError("A non-empty FP4 MLA chunk gather requires context requests.") + tensors = (cu_chunked_seq_len, chunked_global_offset) + if any(not isinstance(tensor, torch.Tensor) or not tensor.is_cuda for tensor in tensors): + raise ValueError("FP4 MLA chunk gather metadata must use CUDA tensors.") + if ( + cu_chunked_seq_len.dtype != torch.int64 + or cu_chunked_seq_len.ndim != 1 + or cu_chunked_seq_len.numel() < num_contexts + 1 + or not cu_chunked_seq_len.is_contiguous() + ): + raise ValueError( + "FP4 MLA chunk gather requires a contiguous int64 cumulative-length tensor." + ) + if ( + chunked_global_offset.dtype != torch.int64 + or chunked_global_offset.ndim != 1 + or chunked_global_offset.numel() < num_contexts + or not chunked_global_offset.is_contiguous() + ): + raise ValueError("FP4 MLA chunk gather requires contiguous int64 global offsets.") + if cu_chunked_seq_len.device != chunked_global_offset.device: + raise ValueError("FP4 MLA chunk gather metadata tensors must share one device.") + + device = cu_chunked_seq_len.device + compressed_kv = torch.empty( + (num_ctx_cached_tokens, kv_lora_rank), dtype=torch.float8_e4m3fn, device=device + ) + compressed_sf = torch.empty( + (triton.cdiv(num_ctx_cached_tokens, 128) * 128 * (kv_lora_rank // 32),), + dtype=torch.uint8, + device=device, + ) + k_pe = torch.empty( + (num_ctx_cached_tokens, qk_rope_head_dim), dtype=torch.float8_e4m3fn, device=device + ) + if num_ctx_cached_tokens == 0: + return (compressed_kv, compressed_sf), k_pe + + state = metadata.fp4_mla_state + if not bool(getattr(state, "device_page_table", False)): + raise RuntimeError("FP4 MLA chunk gather requires fixed-stride device page metadata.") + _materialize_fp4_mla_device_page_table_for_forward(metadata) + page_table_stride = int(state.page_table_stride) + page_ids = state._paged_kv_indices + if ( + page_table_stride <= 0 + or not isinstance(page_ids, torch.Tensor) + or page_ids.dtype != torch.int32 + or not page_ids.is_cuda + or page_ids.device != device + or page_ids.numel() < num_contexts * page_table_stride + ): + raise RuntimeError("FP4 MLA chunk gather received invalid device page metadata.") + + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + head_dim = kv_lora_rank + qk_rope_head_dim + storage_head_dim = _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=_fp4_mla_attention_backend(), + ) + if storage_head_dim != head_dim + FP4_MLA_K_RESIDUAL_DIM: + raise RuntimeError( + "FP4 MLA chunk gather requires K residual storage for BF16 reconstruction." + ) + from .ops.attention.fp4_mla_gather_mxfp8 import fp4_mla_gather_mxfp8 + + logger.info_once( + "Kimi K3 MLA context gather: cached-prefix chunks are read from the FP4 cache as MXFP8 " + "latent + FP8 K-pe (fp4_mla_gather_mxfp8).", + key="k3_mla_ctx_gather_mxfp8", + ) + fp4_mla_gather_mxfp8( + compressed_kv, + compressed_sf, + k_pe, + kv_cache, + sf_cache.view(torch.float8_e4m3fn), + page_ids, + cu_chunked_seq_len, + chunked_global_offset, + _get_fp4_mla_global_scale(metadata, kv_cache.device), + chunked_max_seq_len, + metadata.page_size, + page_table_stride=page_table_stride, + num_contexts=num_contexts, + ) + return (compressed_kv, compressed_sf), k_pe + + +def uses_k3_split_k_context(layer, q: torch.Tensor, output: torch.Tensor) -> bool: + """Whether the FP4-cache chunked prefill of ``layer`` runs on the split-K context FMHA: BF16 Q + and output, an NVFP4 latent cache on ``TrtllmAttention`` and + :func:`supports_k3_fp8_split_k_context`.""" + mha = layer.mha + return ( + q.dtype == torch.bfloat16 + and output.dtype == torch.bfloat16 + and layer.out_scale is None + and isinstance(mha, TrtllmAttention) + and mha.uses_fp4_mla_attention + and supports_k3_fp8_split_k_context(mha) + ) + + +def kv_b_takes_mxfp8(layer) -> bool: + """Whether kv_b runs as the SM107 MXFP8 GEMM on an MXFP8 cached-prefix gather. + + Needs SM107 with the CuTe DSL Rubin kernels, the 512 / 64 latent / RoPE dimensions of the + gather, and a bias-free FP8 block-scale kv_b that takes a prequantized CuTe-layout input, with + its UE8M0 weight scale in the CuTe MXFP8 layout. + """ + kv_b_proj = layer.kv_b_proj + weight = getattr(kv_b_proj, "weight", None) + weight_scale = getattr(kv_b_proj, "weight_scale", None) + return ( + get_sm_version() == 107 + and IS_CUTLASS_DSL_RUBIN_AVAILABLE + and layer.kv_lora_rank == 512 + and layer.qk_rope_head_dim == 64 + and getattr(kv_b_proj, "bias", None) is None + and getattr(kv_b_proj, "supports_prequantized_input", False) + and weight is not None + and weight.dtype == torch.float8_e4m3fn + and weight_scale is not None + and weight_scale.dtype == torch.uint8 + ) + + +def forward_context_split_k_fp4( + layer, + q: torch.Tensor, + compressed_kv: torch.Tensor, + latent_cache: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + output: torch.Tensor, +) -> torch.Tensor: + """FP4 MLA chunked prefill of ``layer`` with every partition on the split-K context FMHA.""" + mha = layer.mha + if latent_cache is None: + raise RuntimeError("FP4 MLA chunked prefill requires latent_cache.") + kv_b_mxfp8 = kv_b_takes_mxfp8(layer) + logger.info_once( + "Kimi K3 MLA cached-prefix context: NVFP4 KV-cache layers run the Rubin FP8 split-K " + "context FMHA (k3_context_fmha); cached prefix: " + + ( + "MXFP8 gather + FP8-output MXFP8 kv_b GEMM" + if kv_b_mxfp8 + else "BF16 gather + kv_b_proj + FP8 quantization" + ) + + "; Q FP8 quantization: static_quantize_e4m3 (Triton).", + key=("k3_mla_ctx_split_k", kv_b_mxfp8), + ) + lora, rope = layer.kv_lora_rank, layer.qk_rope_head_dim + heads = layer.num_heads_tp + + # Rotate Q and the current K in place and append the chunk to the FP4 cache; the current + # chunk's rotated K-pe is then read from latent_cache. + mha.append_fp4_mla_context_cache(q, latent_cache, attn_metadata) + current_k_pe = latent_cache.reshape(-1, lora + rope)[:, lora:] + fp8_unit_scale = mha.kv_cache_scaling_factor + q_fp8 = static_quantize_e4m3(q.contiguous(), fp8_unit_scale) + context_cu_q_seqlens = attn_metadata.mla_prepare_ctx_cu_seqlens() + if context_cu_q_seqlens is None: + raise RuntimeError("K3 split-K context FMHA requires context requests.") + + # The first partition of every request writes the output and its stats (merge op 2). + softmax_stats = torch.empty( + (attn_metadata.num_ctx_tokens, heads, 2), dtype=torch.float, device=q.device + ) + kv_split = [heads * layer.qk_nope_head_dim, heads * layer.v_head_dim] + + def expand_kv( + partition_compressed_kv: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + partition_k_pe: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if kv_b_mxfp8: + if isinstance(partition_compressed_kv, tuple): + activation, activation_scale = partition_compressed_kv + else: + activation, activation_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + partition_compressed_kv.contiguous() + ) + # The GEMM writes BF16 only; FP8 at unit scale from its BF16 rounding. + partition_kv, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + activation, + layer.kv_b_proj.weight, + activation_scale, + layer.kv_b_proj.weight_scale, + ), + fp8_unit_scale, + ) + else: + partition_kv, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + layer.kv_b_proj(partition_compressed_kv), fp8_unit_scale + ) + if partition_k_pe.dtype == torch.float8_e4m3fn: + k_pe_source = "FP8 from the gather" + elif fp8_static_quant.supports_rows(partition_k_pe): + k_pe_source = "static_quantize_e4m3_rows (Triton)" + partition_k_pe = static_quantize_e4m3(partition_k_pe, fp8_unit_scale, rows=True) + else: + k_pe_source = "static_quantize_e4m3_per_tensor" + partition_k_pe, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + partition_k_pe.contiguous(), fp8_unit_scale + ) + logger.info_once( + "Kimi K3 MLA context expand: per-chunk kv_b_proj runs as " + + ( + "cute_dsl_mxfp8_gemm_rubin (MXFP8 in, BF16 out) + static_quantize_e4m3_per_tensor" + if kv_b_mxfp8 + else "kv_b_proj (BF16) + static_quantize_e4m3_per_tensor" + ) + + f"; K-pe: {k_pe_source}; K-nope / K-pe / V go to the split-K FMHA unconcatenated.", + key=("k3_mla_ctx_split_k_expand", kv_b_mxfp8, k_pe_source), + ) + partition_k_nope, partition_v = partition_kv.split(kv_split, -1) + return partition_k_nope, partition_k_pe, partition_v + + def run_partition( + partition: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + cu_kv_seqlens: torch.Tensor, + max_kv_len: int, + merge_op: torch.Tensor, + is_causal: bool, + ) -> None: + partition_k_nope, partition_k_pe, partition_v = partition + run_k3_split_k_context_partition( + mha, + q_fp8, + partition_k_nope, + partition_k_pe, + partition_v, + output, + softmax_stats, + merge_op, + context_cu_q_seqlens, + cu_kv_seqlens, + max_q_len=attn_metadata.max_ctx_seq_len, + max_kv_len=max_kv_len, + is_causal=is_causal, + ) + + for loop_idx in range(attn_metadata.chunked_loop_num): + cu_chunked_seq_len = attn_metadata.cu_chunked_seq_len[loop_idx] + total_chunked_tokens = int( + attn_metadata.host_cu_chunked_seq_len[loop_idx, attn_metadata.num_contexts] + ) + chunked_max_seq_len = int(attn_metadata.max_chunk_len_per_loop[loop_idx]) + if kv_b_mxfp8: + chunked_compressed_kv, chunked_k_pe = load_fp4_mla_chunked_kv_cache_mxfp8( + attn_metadata, + mha.layer_idx, + num_ctx_cached_tokens=total_chunked_tokens, + cu_chunked_seq_len=cu_chunked_seq_len, + chunked_global_offset=attn_metadata.chunked_global_offset[loop_idx], + chunked_max_seq_len=chunked_max_seq_len, + kv_lora_rank=lora, + qk_rope_head_dim=rope, + ) + else: + chunked_compressed_kv, chunked_k_pe = mha.load_chunked_kv_cache_for_mla( + metadata=attn_metadata, + num_ctx_cached_tokens=total_chunked_tokens, + cu_chunked_seq_len=cu_chunked_seq_len, + chunked_global_offset=attn_metadata.chunked_global_offset[loop_idx], + chunked_max_seq_len=chunked_max_seq_len, + out_dtype=q.dtype, + ) + partition = expand_kv(chunked_compressed_kv, chunked_k_pe) + chunked_compressed_kv = None + chunked_k_pe = None + run_partition( + partition, + cu_chunked_seq_len, + chunked_max_seq_len, + attn_metadata.merge_op_tensor[loop_idx], + is_causal=False, + ) + + run_partition( + expand_kv(compressed_kv, current_k_pe), + context_cu_q_seqlens, + attn_metadata.max_ctx_seq_len, + attn_metadata.merge_op_tensor[attn_metadata.chunked_loop_num], + is_causal=True, + ) + return output + + +__all__ = [ + "forward_context_split_k_fp4", + "kv_b_takes_mxfp8", + "load_fp4_mla_chunked_kv_cache_mxfp8", + "run_k3_split_k_context_partition", + "split_k_context_unavailable_reason", + "supports_k3_fp8_split_k_context", + "uses_k3_split_k_context", + "context_cached_kv", + "context_chunked_prefill", + "context_default", + "forward_context", + "project_context_kv", + "HelixGenHooks", + "attn_forward_gen", + "forward_absorption_generation", + "forward_gen_q_fused", + "gen_attention", + "helix_exchange", + "helix_gen_zero_kv_mask", + "helix_tail_applies", + "latent_row", + "project_helix_v_partials", + "q_fused_applies", + "tail_output_buffer", + "K3MLAAttention", + "identity_rope_table", + "ull3_gate_place", +] + + +def project_context_kv(layer, compressed_kv: torch.Tensor) -> torch.Tensor: + """kv_b over the context heads: ``[T, H * qk_nope | H * v_head_dim]`` (one GEMM, or under Helix + this rank's head range of the K and V halves, concatenated).""" + if layer.num_context_heads == layer.num_heads_tp: + return layer.kv_b_proj(compressed_kv) + first_head = layer.mapping.cp_rank * layer.num_context_heads + last_head = first_head + layer.num_context_heads + key = layer.forward_slice( + layer.kv_b_proj, + compressed_kv, + first_head * layer.qk_nope_head_dim, + last_head * layer.qk_nope_head_dim, + ) + value_start = layer.num_heads_tp * layer.qk_nope_head_dim + value = layer.forward_slice( + layer.kv_b_proj, + compressed_kv, + value_start + first_head * layer.v_head_dim, + value_start + last_head * layer.v_head_dim, + ) + # The [all K | all V] row stride the context FP8 preprocessing reads. + return torch.cat((key, value), dim=-1) + + +def _expand_kv_bf16(layer, compressed_kv: torch.Tensor, k_pe: torch.Tensor): + """kv_b, then K ``[T, H * 192]`` (K-nope and the shared K-pe, concatenated) and V.""" + heads = layer.num_context_heads + kv = project_context_kv(layer, compressed_kv) + k_nope, v = kv.split([heads * layer.qk_nope_head_dim, heads * layer.v_head_dim], -1) + k_nope = k_nope.view(-1, heads, layer.qk_nope_head_dim) + k_pe = k_pe.view(-1, 1, layer.qk_rope_head_dim) + k = maybe_compiled_cat((k_nope, k_pe.expand(-1, heads, -1)), dim=-1) + return k.view(-1, heads * layer.qk_head_dim), v + + +def _expand_kv_fp8(layer, compressed_kv: torch.Tensor, k_pe: torch.Tensor): + """kv_b, then FP8 K ``[T, H * 192]`` and V ``[T, H * 128]`` in one pass: the bytes the + attention op's context copy writes for the BF16 K and V.""" + return mla_kv_to_fp8( + project_context_kv(layer, compressed_kv), + k_pe.view(-1, layer.qk_rope_head_dim), + layer.num_context_heads, + layer.qk_nope_head_dim, + layer.v_head_dim, + ) + + +def _fp8_q_args(layer, q_fp8): + """Attention-op arguments that hand it ``q_fp8`` as the already-quantized FP8 Q.""" + if q_fp8 is None: + return {} + return {"quant_q_buffer": q_fp8, "quant_scale_qkv": layer.mha.kv_scale_orig_quant} + + +def _quantize_q_fp8(layer, q: torch.Tensor) -> torch.Tensor: + """Q in FP8 as the attention op's context copy quantizes it (unit scale).""" + q_fp8, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + q, layer.mha.kv_scale_orig_quant + ) + return q_fp8 + + +@functools.cache +def _warmup_expand_kv_fp8( + num_heads, qk_nope_head_dim, qk_rope_head_dim, kv_lora_rank, v_head_dim, device +): + """Compile ``mla_kv_to_fp8`` before the first cached-KV context step (k_pe contiguous for + loaded chunks, a ``latent_cache`` view for the current tokens).""" + num_tokens = 2 + kv = torch.zeros( + num_tokens, num_heads * (qk_nope_head_dim + v_head_dim), dtype=torch.bfloat16, device=device + ) + latent = torch.zeros( + num_tokens, kv_lora_rank + qk_rope_head_dim, dtype=torch.bfloat16, device=device + ) + for k_pe in (latent[:, kv_lora_rank:].contiguous(), latent[:, kv_lora_rank:]): + mla_kv_to_fp8(kv, k_pe, num_heads, qk_nope_head_dim, v_head_dim) + + +@functools.cache +def _warmup_expand_kv_bf16( + num_heads, qk_nope_head_dim, qk_rope_head_dim, kv_lora_rank, v_head_dim, dtype, device +): + """Compile ``maybe_compiled_cat`` for both K layouts (a loaded chunk's contiguous K-pe and the + current tokens' ``latent_cache`` view) with the token dim dynamic; one token count still + specializes, so it is warmed separately.""" + + def warmup(num_tokens): + k_nope = torch.empty( + num_tokens, num_heads * (qk_nope_head_dim + v_head_dim), dtype=dtype, device=device + )[:, : num_heads * qk_nope_head_dim].view(num_tokens, num_heads, qk_nope_head_dim) + chunked_k_pe = torch.empty( + num_tokens, 1, qk_rope_head_dim, dtype=dtype, device=device + ).expand(-1, num_heads, -1) + k_pe = torch.empty( + num_tokens, 1, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + )[:, :, -qk_rope_head_dim:].expand(-1, num_heads, -1) + torch._dynamo.maybe_mark_dynamic(k_nope, 0) + torch._dynamo.maybe_mark_dynamic(chunked_k_pe, 0) + torch._dynamo.maybe_mark_dynamic(k_pe, 0) + maybe_compiled_cat((k_nope, chunked_k_pe), dim=-1) + maybe_compiled_cat((k_nope, k_pe), dim=-1) + + warmup(2) + warmup(1) + + +def _kv_b_strided_out(layer) -> bool: + """Whether the no-prefix expansion writes K-nope / V straight into the attention inputs with + per-head BMMs: a bias-free BF16 kv_b on SM100f.""" + w = layer.kv_b_proj.weight + return ( + is_sm_100f() + and w.dtype == torch.bfloat16 + and tuple(w.shape) + == (layer.num_heads_tp * (layer.qk_nope_head_dim + layer.v_head_dim), layer.kv_lora_rank) + and getattr(layer.kv_b_proj, "bias", None) is None + ) + + +def context_default( + layer, + q: torch.Tensor, + compressed_kv: torch.Tensor, + md: AttentionMetadata, + output: torch.Tensor, + latent_cache: torch.Tensor, +) -> torch.Tensor: + """Context attention without a cached prefix (see the module docstring).""" + heads = layer.num_context_heads + nope, rope, vdim = layer.qk_nope_head_dim, layer.qk_rope_head_dim, layer.v_head_dim + append_owned_kv = md.helix_context_local_slots is not None + k_pe = None + if append_owned_kv: + layer.mha.mla_rope_append_paged_kv_assign_q(q, latent_cache, md) + k_pe = latent_cache[:, layer.kv_lora_rank :] + k = torch.empty_like(q).view(-1, heads, layer.qk_head_dim) + if _kv_b_strided_out(layer): + num_tokens = compressed_kv.shape[0] + w = layer.kv_b_proj.weight + nope_rows = layer.num_heads_tp * nope + w_k = w[:nope_rows].view(layer.num_heads_tp, nope, layer.kv_lora_rank) + w_v = w[nope_rows:].view(layer.num_heads_tp, vdim, layer.kv_lora_rank) + if heads != layer.num_heads_tp: + first = layer.mapping.cp_rank * heads + w_k, w_v = w_k[first : first + heads], w_v[first : first + heads] + ckv = compressed_kv.unsqueeze(0).expand(heads, num_tokens, layer.kv_lora_rank) + # [H, T, nope] view of K with strides (192, H * 192, 1). + k_nope_out = k[..., :nope].transpose(0, 1) + # The context FP8 quantize kernel reads V with the [K-nope | V] row stride of the kv_b + # output: V is the second half of such a buffer (its first half is never written). + kv_buf = torch.empty(num_tokens, heads * (nope + vdim), dtype=q.dtype, device=q.device) + v = kv_buf[:, heads * nope :] + v_out = v.unflatten(-1, (heads, vdim)).transpose(0, 1) + bmm_out(ckv, w_k.transpose(1, 2), k_nope_out) + bmm_out(ckv, w_v.transpose(1, 2), v_out) + else: + kv = project_context_kv(layer, compressed_kv) + k_nope, v = kv.split([heads * nope, heads * vdim], -1) + maybe_compiled_copy_(k[..., :nope], k_nope.view(-1, heads, nope)) + # With the RoPE fused in the op the K-pe slot stays unwritten; the op reads latent_cache. + if append_owned_kv: + k[..., nope:] = k_pe.view(-1, 1, rope) + k = k.view(-1, heads * layer.qk_head_dim) + return layer.mha.forward( + q, + k, + v, + md, + forward_args=AttentionForwardArgs( + attention_input_type=AttentionInputType.context_only, + latent_cache=None if append_owned_kv else latent_cache, + out_scale=None, + output=output, + ), + ) + + +def context_cached_kv( + layer, + q: torch.Tensor, + latent_cache: torch.Tensor, + md: AttentionMetadata, + output: torch.Tensor, +) -> torch.Tensor: + """Context attention over the whole paged cache (cached prefix, no chunked prefill).""" + mha = layer.mha + mha.mla_rope_append_paged_kv_assign_q(q, latent_cache, md) + full_compressed_kv, full_k_pe = mha.load_paged_kv_cache_for_mla(md, q.dtype) + assert full_compressed_kv.shape == ( + md.num_ctx_cached_tokens + md.num_ctx_tokens, + layer.kv_lora_rank, + ) + assert full_k_pe.shape == (md.num_ctx_cached_tokens + md.num_ctx_tokens, layer.qk_rope_head_dim) + assert full_compressed_kv.is_contiguous() and full_k_pe.is_contiguous() + q_fp8 = None + if layer.fp8_cached_kv_context: + q_fp8 = _quantize_q_fp8(layer, q) + full_k, full_v = _expand_kv_fp8(layer, full_compressed_kv, full_k_pe) + else: + full_k, full_v = _expand_kv_bf16(layer, full_compressed_kv, full_k_pe) + full_compressed_kv = full_k_pe = None + return mha.forward( + q, + full_k, + full_v, + md, + forward_args=AttentionForwardArgs( + attention_input_type=AttentionInputType.context_only, + latent_cache=None, + out_scale=None, + output=output, + **_fp8_q_args(layer, q_fp8), + ), + ) + + +def context_chunked_prefill( + layer, + q: torch.Tensor, + compressed_kv: torch.Tensor, + latent_cache: torch.Tensor, + md: AttentionMetadata, + output: torch.Tensor, +) -> torch.Tensor: + """Chunked-prefill context attention over an FP8 / BF16 cache (see the module docstring).""" + mha = layer.mha + heads = layer.num_context_heads + mha.mla_rope_append_paged_kv_assign_q(q, latent_cache, md) + fp8 = layer.fp8_cached_kv_context + expand = _expand_kv_fp8 if fp8 else _expand_kv_bf16 + # With FP8 attention inputs, Q is quantized once for every partition. + q_fp8 = _quantize_q_fp8(layer, q) if fp8 else None + q_args = _fp8_q_args(layer, q_fp8) + softmax_stats = torch.empty((md.num_ctx_tokens, heads, 2), dtype=torch.float, device="cuda") + temp_softmax_stats = torch.empty( + (md.num_ctx_tokens, heads, 2), dtype=torch.float, device="cuda" + ) + temp_output = q.new_empty((q.size(0), heads * layer.v_head_dim), dtype=q.dtype) + buffer_batch_size = md.runtime_features.chunked_prefill_buffer_batch_size + + # Each partition attends with its own KV lengths; the originals are restored at the end. + origin_kv_lens_cuda_runtime = md.kv_lens_cuda_runtime + origin_kv_lens_runtime = md.kv_lens_runtime + origin_ctx_total_kv_len = md.host_total_kv_lens[0].item() + + for loop_idx in range(md.chunked_loop_num): + total_chunked_tokens = md.host_cu_chunked_seq_len[loop_idx, md.num_contexts] + chunked_compressed_kv, chunked_k_pe = mha.load_chunked_kv_cache_for_mla( + metadata=md, + num_ctx_cached_tokens=total_chunked_tokens, + cu_chunked_seq_len=md.cu_chunked_seq_len[loop_idx], + chunked_global_offset=md.chunked_global_offset[loop_idx], + chunked_max_seq_len=md.max_chunk_len_per_loop[loop_idx], + out_dtype=q.dtype, + helix_plan=( + md.helix_mla_chunk_plans[loop_idx] if md.helix_mla_chunk_plans is not None else None + ), + ) + chunked_k, chunked_v = expand(layer, chunked_compressed_kv, chunked_k_pe) + chunked_compressed_kv = chunked_k_pe = None + md.kv_lens_runtime = md.host_chunked_seq_len[loop_idx] + md.kv_lens_cuda_runtime = md.chunked_seq_len[loop_idx] + md.host_total_kv_lens[0] = total_chunked_tokens + temp_output = mha.forward( + q, + chunked_k, + chunked_v, + md, + forward_args=AttentionForwardArgs( + attention_input_type=AttentionInputType.context_only, + latent_cache=None, + out_scale=None, + attention_mask=PredefinedAttentionMask.FULL, + softmax_stats_tensor=temp_softmax_stats, + chunked_prefill_buffer_batch_size=buffer_batch_size, + output=temp_output, + **q_args, + ), + ) + mha.merge_attention_for_mla( + output, + temp_output, + softmax_stats, + temp_softmax_stats, + md.merge_op_tensor[loop_idx], + md, + ) + + # The current chunk, causal; its rotated K-pe is in latent_cache. + _, k_pe = latent_cache.view([-1, layer.kv_lora_rank + layer.qk_rope_head_dim]).split( + [layer.kv_lora_rank, layer.qk_rope_head_dim], -1 + ) + k, v = expand(layer, compressed_kv, k_pe) + md.kv_lens_runtime = md.prompt_lens_cpu_runtime + md.kv_lens_cuda_runtime = md.prompt_lens_cuda_runtime + md.host_total_kv_lens[0] = md.prompt_lens_cpu_runtime[: md.num_contexts].sum().item() + temp_output = mha.forward( + q, + k, + v, + md, + forward_args=AttentionForwardArgs( + attention_input_type=AttentionInputType.context_only, + latent_cache=None, + out_scale=None, + softmax_stats_tensor=temp_softmax_stats, + chunked_prefill_buffer_batch_size=buffer_batch_size, + output=temp_output, + **q_args, + ), + ) + mha.merge_attention_for_mla( + output, + temp_output, + softmax_stats, + temp_softmax_stats, + md.merge_op_tensor[md.chunked_loop_num], + md, + ) + md.kv_lens_runtime = origin_kv_lens_runtime + md.kv_lens_cuda_runtime = origin_kv_lens_cuda_runtime + md.host_total_kv_lens[0] = origin_ctx_total_kv_len + return output + + +def forward_context( + layer, + q: torch.Tensor, + compressed_kv: torch.Tensor, + md: AttentionMetadata, + output: torch.Tensor, + latent_cache: torch.Tensor, +) -> torch.Tensor: + """Context attention of the context rows into ``output``: picks the chain by cache state.""" + mha = layer.mha + if md.helix_context_local_slots is not None: + if mha.uses_fp4_mla_attention: + raise NotImplementedError( + "Owner-only Helix prefill does not support dense FP4 MLA KV cache" + ) + if md.num_ctx_cached_tokens > 0 and ( + not md.runtime_features.chunked_prefill or get_sm_version() < 100 + ): + raise NotImplementedError( + "Helix cached context requires chunked prefill on SM100 or newer" + ) + if mha.has_cached_kv_for_mla_context_warmup(md): + dims = ( + layer.num_context_heads, + layer.qk_nope_head_dim, + layer.qk_rope_head_dim, + layer.kv_lora_rank, + layer.v_head_dim, + ) + if layer.fp8_cached_kv_context: + _warmup_expand_kv_fp8(*dims, q.device) + else: + _warmup_expand_kv_bf16(*dims, q.dtype, q.device) + chunked = mha.is_chunked_prefill_for_mla_context(md) + if chunked and get_sm_version() >= 100: + if mha.uses_fp4_mla_attention: + if not uses_k3_split_k_context(layer, q, output): + raise NotImplementedError( + "Kimi K3 MLA: an FP4 MLA cache attends a cached prefix only through the Rubin " + "FP8 split-K context FMHA (SM107, BF16 Q and output, see ctx_split_k)" + ) + return forward_context_split_k_fp4(layer, q, compressed_kv, latent_cache, md, output) + return context_chunked_prefill(layer, q, compressed_kv, latent_cache, md, output) + if chunked or mha.has_cached_kv_for_mla_context(md): + if mha.uses_fp4_mla_attention: + raise RuntimeError( + "FP4 MLA attends a cached prefix only through chunked-prefill partitions on " + "SM100+; enable chunked prefill." + ) + return context_cached_kv(layer, q, latent_cache, md, output) + return context_default(layer, q, compressed_kv, md, output, latent_cache) + + +# Batches of at most this many tokens (helix n4 batches included) start the tail after the decode +# instead of at the decode's entry (the decode then releases it from its epilogue, once its last KV +# tile is done; ``ops/attention/_fp4_mla_decode``). +_LATE_TAIL_MAX_TOKENS = 32 + + +def ull2_tail_enabled() -> bool: + return _features().ull2_mla_tail + + +def decode_releases_tail_early(num_tokens: int) -> bool: + """Whether the FP4 decode before the fused tail releases it at the decode's entry + (``DeferredPackedReduce.release_early``, with ``features.fp4_mla_reduce_pdl``). + + With ``features.ull2_mla_tail`` and at most :data:`_LATE_TAIL_MAX_TOKENS` tokens it does not: + a tail grid that is resident during the whole decode folds and projects 2-4 us slower once + the decode completes than one launched at its end, more than the early launch hides. Larger + batches launch most of their tail CTAs at the decode's end anyway and gain from the others. + """ + return not (ull2_tail_enabled() and num_tokens <= _LATE_TAIL_MAX_TOKENS) + + +def _check_cu_result(cu_func_ret): + """The values a CUDA driver call returned after its status, raising if it failed.""" + if isinstance(cu_func_ret, tuple): + cu_result, *others = cu_func_ret + if cu_result != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(cu_result) + if len(others) == 1: + return others[0] + elif len(others) > 1: + return tuple(others) + else: + return None + else: + if cu_func_ret != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(cu_func_ret) + return None + + +# Receive buffer per rank in MiB (two halves). +_TAIL_MB = 16 + + +class K3HelixTail: + """The fused Helix tail over the CP group of ``mapping`` (one receive workspace per mapping).""" + + _cache: Dict[Mapping, "K3HelixTail"] = {} + + def __init__(self, mapping: Mapping): + from .kernels.k3_helix_tail import SENTINEL + + self.mapping = mapping + MnnvlMemory.initialize() + nbytes = _TAIL_MB << 20 + logger.info(f"Rank {mapping.cp_rank} allocating a {nbytes >> 20} MiB Helix tail buffer") + self.workspace = HelixCpMnnvlMemory(mapping, nbytes) + self.ws_addr = self.workspace.ptr + self.rank_stride = self.workspace.rank_stride + self.half_bytes = nbytes // 2 // 16 * 16 + # Armed through the driver: the kernel only uses raw addresses (see HelixAllToAllFused). + _check_cu_result( + cuda.cuMemsetD8Async( + cuda.CUdeviceptr(self.ws_addr + mapping.cp_rank * self.rank_stride), + SENTINEL & 0xFF, + nbytes, + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + ) + self.state = torch.zeros(8, dtype=torch.int64, device="cuda") + num_sms = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + self.max_ctas = num_sms + self.use_pdl = get_env_enable_pdl() + # Tests on one GPU: this rank's records go to its own buffer (k3_helix_tail loopback). + self.loopback = False + torch.cuda.synchronize() + HelixCpMnnvlMemory.get_comm(mapping).barrier() + + @staticmethod + def get(mapping: Mapping) -> "K3HelixTail": + if mapping not in K3HelixTail._cache: + K3HelixTail._cache[mapping] = K3HelixTail(mapping) + return K3HelixTail._cache[mapping] + + @staticmethod + def applies( + mapping: Mapping, num_tokens: int, num_heads: int, latent_dim: int, v_head_dim: int + ) -> bool: + """Whether the fused tail serves a generation batch of ``num_tokens`` tokens.""" + from .kernels import k3_helix_tail as kernel + + if get_sm_version() not in (100, 103, 107) or num_tokens <= 0: + return False + if not mapping.has_cp_helix() or mapping.cp_config.get("use_nccl_for_alltoall", True): + return False + if not kernel.can_run(mapping.cp_size, num_heads, latent_dim, v_head_dim): + return False + nbytes = _TAIL_MB << 20 + record = kernel.record_bytes(num_heads // mapping.cp_size, v_head_dim) + return num_tokens * mapping.cp_size * record <= nbytes // 2 // 16 * 16 + + def __call__( + self, + rows: torch.Tensor, + stats: torch.Tensor, + w_uv: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor], + fold=None, + out: Optional[torch.Tensor] = None, + mx_out=None, + timestamps: Optional[torch.Tensor] = None, + mx_atom: bool = False, + ) -> torch.Tensor: + """Fold, project, exchange and combine one layer's partials. + + Args: + rows: BF16 ``[T, H * 512]`` decode output (latent rows ``(token, head)``); with + ``fold`` only the rows the decode stored directly are read. + stats: FP32 ``[T, H, 2]`` ``(max, sum)`` of ``rows``; read only without ``fold``. + w_uv: BF16 ``[H, 128, 512]`` V up-projection of every head. + zero_kv_mask: Optional bool ``[T]``, True where this rank owns no KV of the token. + fold: ``k3_helix_tail.FoldInputs`` of the deferred split reduction, or None. + out: Optional BF16 ``[T, H / P * 128]`` output. + mx_out: Optional ``(gate, xq, xs)``: also store the MXFP8 form of + ``bf16(out * gate)``, o_proj's gated operand (see ``k3_helix_tail``). + timestamps: Optional int64 phase stamps (microbenchmarks, see ``k3_helix_tail``). + mx_atom: (c240) ``xs`` in the mid-M GEMM's scale layout (see ``k3_helix_tail``). + + Returns: + The combined BF16 ``[T, H / P * 128]`` output of this rank's heads. + """ + from .ops.attention.k3_helix_tail import k3_helix_tail + + num_tokens = rows.shape[0] + num_heads = w_uv.shape[0] + cp_size = self.mapping.cp_size + if out is None: + out = rows.new_empty((num_tokens, num_heads // cp_size * w_uv.shape[1])) + if zero_kv_mask is not None: + zero_kv_mask = zero_kv_mask[:num_tokens].contiguous() + k3_helix_tail( + rows.reshape(num_tokens * num_heads, -1), + stats.reshape(num_tokens * num_heads, 2), + w_uv, + out, + fold=fold, + zero_kv_mask=zero_kv_mask, + cp_rank=self.mapping.cp_rank, + cp_size=cp_size, + ws_addr=self.ws_addr, + rank_stride=self.rank_stride, + half_bytes=self.half_bytes, + state=self.state, + max_ctas=self.max_ctas, + use_pdl=self.use_pdl, + mx_out=mx_out, + timestamps=timestamps, + loopback=self.loopback, + mx_atom=mx_atom, + ull2=ull2_tail_enabled(), + ) + return out + + +def fold_inputs_of(deferred, rows: torch.Tensor, stats: torch.Tensor): + """``FoldInputs`` of the reduction a ``DeferredPackedReduce`` recorded for this decode (its + output must be ``rows`` and its statistics ``stats``), or None if it recorded none. A recorded + call on other buffers is launched instead (the caller then reads ``rows`` and ``stats``).""" + from .kernels.k3_helix_tail import FoldInputs + + args = deferred.args + if args is None: + return None + own_buffers = ( + args["output"].data_ptr() == rows.data_ptr() + and args["output"].numel() == rows.numel() + and args["stats"] is not None + and args["stats"].data_ptr() == stats.data_ptr() + ) + if not own_buffers: + deferred.run() + return None + return FoldInputs( + args["partial"], + args["maxima"], + args["sums"], + args["plan"], + args["q_scale"], + args["kv_scale"], + num_requests=args["num_requests"], + valid_m=args["valid_m"], + physical_m=args["physical_m"], + sm_scale=args["sm_scale"], + output_scale=args["output_scale"], + ) + + +_gen_F = _features() + +# Largest generation batch (tokens) the fused query serves (the kernel takes up to 16). +_Q_FUSED_MAX_TOKENS = 16 + + +class HelixGenHooks: + """What the side-gate forward (``K3MLAAttention._forward_side_gate``) hands the Helix + generation chain for one forward, and what the chain hands back. + + ``zero_kv_mask_fn`` returns the zero-local-KV mask computed ahead on the side stream (called + where the exchange needs it); ``pre_exchange_fn`` forks the output gate; ``pre_tail_fn`` joins + it before the fused tail and returns it (None if not forked yet); ``gate_place`` is the + ``features.ull3_mla_gate_at`` placement. Back: ``tail_mx``, the MXFP8 o_proj operand + ``(xq, xs)`` the fused tail stored, and ``gen_output``, the combined all-to-all output.""" + + __slots__ = ( + "zero_kv_mask_fn", + "pre_exchange_fn", + "pre_tail_fn", + "gate_place", + "tail_mx", + "gen_output", + ) + + def __init__( + self, + zero_kv_mask_fn: Callable[[], Optional[torch.Tensor]], + pre_exchange_fn: Callable[[], None], + pre_tail_fn: Callable[[], Optional[torch.Tensor]], + gate_place: str, + ) -> None: + self.zero_kv_mask_fn = zero_kv_mask_fn + self.pre_exchange_fn = pre_exchange_fn + self.pre_tail_fn = pre_tail_fn + self.gate_place = gate_place + self.tail_mx: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + self.gen_output: Optional[torch.Tensor] = None + + +def helix_gen_zero_kv_mask(md: AttentionMetadata, num_tokens: int) -> Optional[torch.Tensor]: + """Per-token mask of the ``num_tokens`` generation tokens this CP rank owns no KV for. + + With valid speculative verify groups KV ownership is per token (``helix_kv_bounds``); + otherwise the per-sequence ``kv_lens_cuda == 0`` of the generation requests, expanded over + their query lengths (``seq_lens_cuda``) with a static output size. None without those + buffers.""" + helix_kv_bounds = getattr(md, "helix_kv_bounds", None) + if helix_kv_bounds is not None and getattr(md, "_helix_spec_tokens_valid", False): + return helix_kv_bounds[:num_tokens] == 0 + kv_lens = getattr(md, "kv_lens_cuda", None) + seq_lens = getattr(md, "seq_lens_cuda", None) + if kv_lens is None or seq_lens is None: + return None + seqs = slice(md.num_contexts, md.num_contexts + md.num_generations) + return torch.repeat_interleave(kv_lens[seqs] == 0, seq_lens[seqs], output_size=num_tokens) + + +def helix_exchange( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + mapping, + num_heads_tp_cp: int, + value_dim: int, + zero_kv_mask: Optional[torch.Tensor], +) -> torch.Tensor: + """All-to-all of this rank's partials over the CP group (fifo v2 on the MNNVL workspace; the + sender zeroes the ``zero_kv_mask`` rows) and the combine: ``[T, num_heads_tp_cp * value_dim]`` + BF16.""" + num_tokens = partial_o.shape[0] + cp_size = mapping.cp_size + partial_o_out, stats_out = HelixAllToAllNative.get(mapping).alltoall_native( + partial_o.view(num_tokens, cp_size, num_heads_tp_cp * value_dim), + softmax_stats.view(num_tokens, cp_size, num_heads_tp_cp * 2), + None if zero_kv_mask is None else zero_kv_mask[:num_tokens], + ) + return torch.ops.trtllm.helix_post_process_native( + partial_o_out.view(num_tokens, cp_size, num_heads_tp_cp, value_dim), + stats_out.view(num_tokens, cp_size, num_heads_tp_cp, 2), + 1.0, + 1, + ) + + +def project_helix_v_partials(layer, partial_o: torch.Tensor) -> torch.Tensor: + """Every tp-local head's V up-projection of this CP rank's partials (it commutes with the + per-head weighted combine): ``[T, H * kv_lora_rank]`` -> ``[T, H * v_head_dim]``.""" + num_tokens = partial_o.shape[0] + heads = layer.num_heads_tp + partial_v = partial_o.new_empty((num_tokens, heads, layer.v_head_dim)) + bmm_out( + partial_o.view(num_tokens, heads, layer.kv_lora_rank).transpose(0, 1), + layer.helix_v_b_proj_full.transpose(1, 2), + partial_v.transpose(0, 1), + ) + return partial_v.view(num_tokens, heads * layer.v_head_dim) + + +def tail_output_buffer( + buf: Optional[torch.Tensor], num_tokens: int, cols: int +) -> Optional[torch.Tensor]: + """``buf`` as the fused tail's BF16 ``[T, cols]`` output, or None where the tail cannot write + it as is.""" + if ( + buf is None + or buf.dtype != torch.bfloat16 + or not buf.is_contiguous() + or buf.numel() != num_tokens * cols + ): + return None + return buf.view(num_tokens, cols) + + +def helix_tail_applies(layer, num_tokens: int) -> bool: + """Whether the Helix tail of this batch is one ``k3_helix_tail`` kernel (``helix_tail``).""" + w_uv = layer.helix_v_b_proj_full + return ( + isinstance(w_uv, torch.Tensor) + and w_uv.dtype == torch.bfloat16 + and w_uv.is_contiguous() + and K3HelixTail.applies( + layer.mapping, num_tokens, layer.num_heads_tp, layer.kv_lora_rank, layer.v_head_dim + ) + ) + + +def gen_attention( + layer, q: torch.Tensor, md: AttentionMetadata, q_rows_external: bool = False, **kwargs +) -> torch.Tensor: + """The absorbed decode: K3's packed FP4 decode on an FP4 cache with the CuTe DSL FP4 MLA + backend, ``TrtllmAttention.forward`` otherwise.""" + if fp4_mla_decode.applies(layer.mqa): + assert kwargs.get("out_scale") is None + return _decode_op.k3_fp4_mla_decode( + layer.mqa, + q, + md, + softmax_stats_tensor=kwargs.get("softmax_stats_tensor"), + q_rows_external=q_rows_external, + ) + return layer.mqa.forward(q, None, None, md, forward_args=AttentionForwardArgs(**kwargs)) + + +def _helix_gen( + layer, + q: torch.Tensor, + md: AttentionMetadata, + helix_v_up: bool, + q_rows_external: bool, + hooks: Optional[HelixGenHooks], + **kwargs, +) -> torch.Tensor: + """Decode with softmax statistics, (V up-projection,) exchange and combine.""" + softmax_stats = torch.empty( + (q.shape[0], layer.num_heads_tp, 2), device=q.device, dtype=torch.float32 + ) + kwargs["softmax_stats_tensor"] = softmax_stats + partial_o = gen_attention(layer, q, md, q_rows_external, **kwargs) + assert partial_o.shape[-1] == layer.num_heads_tp * layer.kv_lora_rank + if hooks is not None: + zero_kv_mask = hooks.zero_kv_mask_fn() + else: + zero_kv_mask = helix_gen_zero_kv_mask(md, partial_o.shape[0]) + value_dim = layer.kv_lora_rank + if helix_v_up: + partial_o = project_helix_v_partials(layer, partial_o) + value_dim = layer.v_head_dim + if hooks is not None: + hooks.pre_exchange_fn() + return helix_exchange( + partial_o, softmax_stats, layer.mapping, layer.num_heads_tp_cp, value_dim, zero_kv_mask + ) + + +def _helix_gen_fused_tail( + layer, + q: torch.Tensor, + md: AttentionMetadata, + q_rows_external: bool, + output: Optional[torch.Tensor], + hooks: Optional[HelixGenHooks], + **kwargs, +) -> torch.Tensor: + """The decode with its split reduction deferred, then one tail kernel (fold, V + up-projection, exchange, combine) writing ``output`` where it can. With the side gate forked + early (gate place other than "exchange") the tail also stores o_proj's MXFP8 operand + ``bf16(out * gate)`` where o_proj's gated MXFP8 GEMM takes it.""" + from .kernels.fp4_mla_packed_reduce import DeferredPackedReduce + + num_tokens = q.shape[0] + softmax_stats = torch.empty( + (num_tokens, layer.num_heads_tp, 2), device=q.device, dtype=torch.float32 + ) + kwargs["softmax_stats_tensor"] = softmax_stats + gate_early = hooks is not None and hooks.gate_place != "exchange" + if gate_early: + hooks.pre_exchange_fn() + with DeferredPackedReduce(release_early=decode_releases_tail_early(num_tokens)) as deferred: + partial_o = gen_attention(layer, q, md, q_rows_external, **kwargs) + assert partial_o.shape[-1] == layer.num_heads_tp * layer.kv_lora_rank + fold = fold_inputs_of(deferred, partial_o, softmax_stats) + if hooks is not None: + zero_kv_mask = hooks.zero_kv_mask_fn() + else: + zero_kv_mask = helix_gen_zero_kv_mask(md, num_tokens) + gate = None + if hooks is not None: + if gate_early: + gate = hooks.pre_tail_fn() + else: + hooks.pre_exchange_fn() + cols = layer.num_heads_tp_cp * layer.v_head_dim + out = tail_output_buffer(output, num_tokens, cols) + mx_out = None + mx_applies = getattr(layer.o_proj, "gated_mxfp8_applies", None) + if gate is not None and mx_applies is not None: + if out is None: + out = partial_o.new_empty((num_tokens, cols)) + if mx_applies(out, gate): + xq = torch.empty((num_tokens, cols), dtype=torch.float8_e4m3fn, device=q.device) + xs = torch.empty((num_tokens, cols // 32), dtype=torch.uint8, device=q.device) + mx_out = (gate, xq, xs) + hooks.tail_mx = (xq, xs) + return K3HelixTail.get(layer.mapping)( + partial_o, + softmax_stats, + layer.helix_v_b_proj_full, + zero_kv_mask, + fold=fold, + out=out, + mx_out=mx_out, + ) + + +def attn_forward_gen( + layer, + q: torch.Tensor, + md: AttentionMetadata, + *, + helix_v_up: bool, + q_rows_external: bool = False, + output: Optional[torch.Tensor] = None, + hooks: Optional[HelixGenHooks] = None, + **kwargs, +) -> torch.Tensor: + """The absorbed decode of ``q`` (``[T, H * 576]``), with the Helix exchange under Helix CP. + + Returns ``[T, H * kv_lora_rank]`` without Helix, ``[T, H / P * v_head_dim]`` with the V + up-projection before the exchange (``helix_v_up``), ``[T, H / P * kv_lora_rank]`` otherwise. + ``output``: the layer's buffer for the combined rows, which the fused tail writes.""" + num_tokens = q.shape[0] + if helix_v_up and helix_tail_applies(layer, num_tokens): + out = _helix_gen_fused_tail(layer, q, md, q_rows_external, output, hooks, **kwargs) + elif layer.mapping.has_cp_helix(): + out = _helix_gen(layer, q, md, helix_v_up, q_rows_external, hooks, **kwargs) + else: + return gen_attention(layer, q, md, q_rows_external, **kwargs) + if helix_v_up and hooks is not None: + hooks.gen_output = out + return out + + +def forward_absorption_generation( + layer, + q: torch.Tensor, + latent_cache: torch.Tensor, + md: AttentionMetadata, + output: torch.Tensor, + hooks: Optional[HelixGenHooks] = None, +) -> torch.Tensor: + """Absorbed generation of the ``q`` rows (``[T, H * 192]``, q_b's output) into ``output`` + (``[T, H / P * v_head_dim]``); ``latent_cache`` ``[T, 576]`` is the normalized latent and + K-pe the cache update appends.""" + num_tokens = q.shape[0] + heads, lora, rope = layer.num_heads_tp, layer.kv_lora_rank, layer.qk_rope_head_dim + device = q.device + q_nope, q_pe = q.view([-1, heads, layer.qk_head_dim]).split( + [layer.qk_nope_head_dim, rope], dim=-1 + ) + num_seqs = md.num_seqs + # The RoPE kernel writes seqQOffset[batch + 1] and scans total / per-sequence lengths over the + # presented rows: one per token under the ragged (token-major) verify presentation. + cu_rows = ( + max(num_seqs, num_tokens) + 1 if getattr(md, "is_ragged_verify", False) else num_seqs + 1 + ) + cu_q_seqlens = torch.empty(cu_rows, dtype=torch.int32, device=device) + cu_kv_seqlens = torch.empty(cu_rows, dtype=torch.int32, device=device) + fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=device) + mla_bmm1_scale = mla_bmm2_scale = quant_q_buffer = None + if layer.quantized_kv: + mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=device) + mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=device) + quant_q_buffer = torch.empty( + num_tokens, heads, lora + rope, dtype=torch.uint8, device=device + ) + # [T, H, kv_lora_rank | rope]: the W_UK output and the rotated q_pe the RoPE kernel writes. + fused_q = torch.empty([num_tokens, heads, lora + rope], dtype=q.dtype, device=device) + rope_kwargs = {"fuse_fp4_q_quant": True} if layer.fp4_kv else {} + + def rope_and_cache_update(): + layer.mqa.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + md, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + **rope_kwargs, + ) + + # [H, T, qk_nope] x [H, qk_nope, kv_lora_rank] -> [H, T, kv_lora_rank], written into fused_q. + q_nope_out = fused_q[..., :lora].transpose(0, 1) + if layer.k_b_proj_trans.dtype == torch.bfloat16: + + def w_uk(): + bmm_out(q_nope.transpose(0, 1), layer.k_b_proj_trans.transpose(1, 2), q_nope_out) + + else: + w_uk_bf16 = layer.k_b_proj_trans_dequant + if w_uk_bf16 is None: + raise NotImplementedError( + "Kimi K3 MLA: the FP8 absorption BMM runs on the dequantized weight (SM100f)" + ) + + def w_uk(): + torch.bmm(q_nope.transpose(0, 1), w_uk_bf16.transpose(1, 2), out=q_nope_out) + + # The quantized-cache RoPE kernel writes the quantized query too: same stream as the decode. + rope_stream = None if layer.quantized_kv else layer.aux_stream + fork_join(w_uk, rope_and_cache_update, layer.ln_events[0], layer.ln_events[1], rope_stream) + fused_q = fused_q.view([num_tokens, heads * (lora + rope)]) + + helix_v_up = layer.helix_v_up_before_exchange + attn_out_latent = attn_forward_gen( + layer, + fused_q, + md, + helix_v_up=helix_v_up, + output=output if helix_v_up else None, + hooks=hooks, + attention_input_type=AttentionInputType.generation_only, + out_scale=None, + latent_cache=latent_cache, + q_pe=q_pe, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + mla_bmm1_scale=mla_bmm1_scale, + mla_bmm2_scale=mla_bmm2_scale, + quant_q_buffer=quant_q_buffer, + ) + fused_q = None + if helix_v_up: + # The combine produced [T, H / P * v_head_dim]; the fused tail wrote it into output (the + # copy is then a no-op on the same storage). + output.view(attn_out_latent.shape).copy_(attn_out_latent) + return output + + heads_out = layer.num_heads_tp_cp + assert attn_out_latent.shape == (num_tokens, heads_out * lora) + latent = attn_out_latent.view([-1, heads_out, lora]) + attn_output = output.view([num_tokens, heads_out, layer.v_head_dim]) + # [H, T, kv_lora_rank] x [H, kv_lora_rank, v_head_dim] -> [H, T, v_head_dim]. + if layer.v_b_proj.dtype == torch.bfloat16: + bmm_out(latent.transpose(0, 1), layer.v_b_proj.transpose(1, 2), attn_output.transpose(0, 1)) + else: + torch.bmm( + latent.transpose(0, 1), + layer.v_b_proj_dequant.transpose(1, 2), + out=attn_output.transpose(0, 1), + ) + return output + + +# --------------------------------------------------------------------------------------------- +# Fused query: generation-only Helix batches of FP4-KV layers. +# --------------------------------------------------------------------------------------------- + + +def _flashinfer_cute_norm(device: torch.device) -> bool: + """Whether FlashInfer's ``rmsnorm`` runs its CuTe DSL ``RMSNormKernel`` on ``device`` (the + arithmetic ``k3_mla_latent_row`` and ``k3_mla_q_fused`` reproduce), not its CUDA JIT kernel.""" + try: + from flashinfer.norm import _use_cuda_norm + except ImportError: + return False + return not _use_cuda_norm(device) + + +def q_fused_w_uk(layer) -> Optional[torch.Tensor]: + """The BF16 absorption weight ``[H, kv_lora_rank, qk_nope_head_dim]`` the W_UK BMM reads (the + weight, or its dequantized copy).""" + w_kc = layer.k_b_proj_trans + if w_kc is not None and w_kc.dtype != torch.bfloat16: + w_kc = layer.k_b_proj_trans_dequant + return w_kc + + +def q_fused_qb_exponents(layer) -> Optional[torch.Tensor]: + """Block exponents of the q_b weight (``k3_mla_q_fused.block_exponents``), cached per weight + scale; None if they would have to be built during CUDA graph capture.""" + scale = getattr(layer.q_b_proj, "weight_scale", None) + if not isinstance(scale, torch.Tensor): + return None + cached = layer._q_fused_qb_exp + if cached is not None and cached[0] is scale: + return cached[1] + if torch.cuda.is_current_stream_capturing(): + return None + from .kernels import k3_mla_q_fused + + weight = layer.q_b_proj.weight + try: + exponents = k3_mla_q_fused.block_exponents(scale, *weight.shape) + except ValueError: + return None + layer._q_fused_qb_exp = (scale, exponents) + return exponents + + +def q_fused_applies(layer, hidden_states: torch.Tensor, md: AttentionMetadata) -> bool: + """Whether this forward takes :func:`forward_gen_q_fused`: a generation-only Helix batch of an + FP4-KV layer (V up-projection before the exchange) whose query fits the fused kernel.""" + num_tokens = md.num_tokens + if ( + hidden_states.dim() != 2 + or hidden_states.shape[0] != num_tokens + or md.num_contexts != 0 + or md.num_generations <= 0 + or num_tokens % md.num_generations != 0 + or not layer.helix_v_up_before_exchange + or not layer.mqa.uses_fp4_mla_attention + or getattr(md, "helix_position_offsets", None) is None + or getattr(md, "helix_is_inactive_rank", None) is None + or layer.mqa.rotary_cos_sin is None + or layer.mqa.rotary_cos_sin.dtype != torch.float32 + or not layer.mqa.rotary_cos_sin.is_contiguous() + ): + return False + from .kernels import k3_mla_q_fused + + if not ( + layer.num_heads_tp == k3_mla_q_fused.NUM_HEADS + and layer.q_lora_rank == k3_mla_q_fused.Q_LORA_RANK + and layer.kv_lora_rank == k3_mla_q_fused.KV_LORA_RANK + and layer.qk_nope_head_dim == k3_mla_q_fused.NOPE_DIM + and layer.qk_rope_head_dim == k3_mla_q_fused.ROPE_DIM + and num_tokens <= min(k3_mla_q_fused.MAX_TOKENS, _Q_FUSED_MAX_TOKENS) + ): + return False + weight = layer.q_b_proj.weight + w_uk = q_fused_w_uk(layer) + if w_uk is None or not k3_mla_q_fused.weights_supported(weight, w_uk): + return False + return q_fused_qb_exponents(layer) is not None + + +def latent_row(layer, compressed_kv: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: + """``concat([kv_a RMSNorm(compressed_kv), k_pe], -1)``: FlashInfer's norm written straight into + the row and ``k_pe`` copied beside it.""" + norm = layer.kv_a_layernorm + if not (IS_FLASHINFER_AVAILABLE and compressed_kv.dtype == torch.bfloat16): + return torch.concat( + [rms_norm(compressed_kv, norm.weight, norm.variance_epsilon), k_pe], dim=-1 + ) + from flashinfer.norm import rmsnorm + + lora = compressed_kv.shape[-1] + latent = torch.empty( + (compressed_kv.shape[0], lora + k_pe.shape[-1]), + dtype=compressed_kv.dtype, + device=compressed_kv.device, + ) + rmsnorm( + compressed_kv, + norm.weight, + norm.variance_epsilon, + out=latent[:, :lora], + enable_pdl=get_env_enable_pdl(), + ) + latent[:, lora:].copy_(k_pe) + return latent + + +def _q_fused_latent_row(layer, compressed_kv: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: + """:func:`latent_row` in one ``k3_mla_latent_row`` launch where FlashInfer's norm is its CuTe + DSL kernel (whose arithmetic it repeats, bit for bit).""" + from .kernels import k3_mla_q_fused + from .ops.attention.k3_mla_latent_row import k3_mla_latent_row + + norm = layer.kv_a_layernorm + if ( + IS_FLASHINFER_AVAILABLE + and _flashinfer_cute_norm(compressed_kv.device) + and k3_mla_q_fused.latent_row_supported(compressed_kv, k_pe, norm.weight) + ): + return k3_mla_latent_row( + compressed_kv, + k_pe, + norm.weight, + norm.variance_epsilon, + use_pdl=get_env_enable_pdl(), + ) + return latent_row(layer, compressed_kv, k_pe) + + +def forward_gen_q_fused( + layer, + hidden_states: torch.Tensor, + md: AttentionMetadata, + hooks: Optional[HelixGenHooks] = None, +) -> torch.Tensor: + """Absorbed Helix generation of an FP4-KV layer with the fused query; returns the combined + all-to-all output ``[T, H / P * v_head_dim]``. + + After the fused q_a / kv_a projection the KV half (``k3_mla_latent_row``, then the FP4 cache + update without its query rows) runs on the aux stream (with multi-stream on) and the main + stream runs ``k3_mla_q_fused`` (on the raw q_a rows with the RMSNorm inside, or after the + separate q_a RMSNorm), which writes the FP4 query rows into the Q buffers the cache update + recorded on the metadata. The decode joins both; it reads only those buffers (the BF16 + ``fused_q`` it is handed provides shapes).""" + from tensorrt_llm._torch.attention.backends.fp4_mla.cache_update import ( + _prepare_fp4_mla_q_buffers, + ) + from tensorrt_llm._torch.attention.backends.fp4_mla.config import _fp4_mla_packed_rows_per_seq + from tensorrt_llm._torch.attention.backends.fp4_mla.layout import _get_fp4_mla_q_global_scale + + from .kernels import k3_mla_q_fused as q_fused_kernel + from .ops.attention.fp4_mla_kv_append import fp4_mla_kv_append + from .ops.attention.k3_mla_q_fused import k3_mla_q_fused + + num_tokens = md.num_tokens + heads, lora, rope = layer.num_heads_tp, layer.kv_lora_rank, layer.qk_rope_head_dim + device = hidden_states.device + gen_len = num_tokens // md.num_generations + # The Q buffers the KV half records on the metadata, created (zero-filled, once) on this + # stream, which writes them. + _prepare_fp4_mla_q_buffers( + md, num_tokens, heads, device, rows_per_seq=_fp4_mla_packed_rows_per_seq(heads, gen_len) + ) + q, compressed_kv, k_pe = layer.project_qkv_a(hidden_states).split( + [layer.q_lora_rank, lora, rope], -1 + ) + kv_fork, kv_ready = layer._kv_events + kv_stream = layer.aux_stream if do_multi_stream() and layer.aux_stream is not None else None + if kv_stream is not None: + kv_fork.record() + with torch.cuda.stream(kv_stream) if kv_stream is not None else contextlib.nullcontext(): + if kv_stream is not None: + kv_fork.wait() + latent_cache = _q_fused_latent_row(layer, compressed_kv, k_pe) + fp4_mla_kv_append(layer.mqa, latent_cache, md, heads) + if kv_stream is not None: + kv_ready.record() + norm = layer.q_a_layernorm + fuse_norm = ( + q_fused_kernel.NORM_FUSION_AVAILABLE + and IS_FLASHINFER_AVAILABLE + and _flashinfer_cute_norm(device) + and norm.weight.dtype == torch.bfloat16 + and norm.weight.is_contiguous() + and q_fused_kernel.input_supported(q) + ) + if not fuse_norm: + q = layer.q_a_norm(q) + assert q_fused_kernel.input_supported(q), "the fused query reads BF16 q_a rows" + state = md.fp4_mla_state + rows_per_seq = int(state.prequantized_q_rows_per_seq or 0) + k3_mla_q_fused( + q, + layer.q_b_proj.weight, + q_fused_qb_exponents(layer), + q_fused_w_uk(layer), + _get_fp4_mla_q_global_scale(md, device), + layer.mqa.rotary_cos_sin, + md.helix_position_offsets, + state.prequantized_q, + state.prequantized_q_sf.view(torch.uint8), + gen_len=gen_len, + rows_per_seq=rows_per_seq or gen_len * heads, + use_pdl=get_env_enable_pdl(), + norm=(norm.weight, norm.variance_epsilon) if fuse_norm else None, + # features.ull2_mla_decode_pdl: the decode (a programmatic dependent reading only the + # query after its wait) starts its prologue as this kernel's CTAs finish. + early_trigger=_gen_F.ull2_mla_decode_pdl, + ) + fp4_mla_decode.mark_query_launch(md) + if kv_stream is not None: + kv_ready.wait() + fused_q = torch.empty([num_tokens, heads, lora + rope], dtype=q.dtype, device=device) + q_pe = fused_q[..., lora:] + num_seqs = md.num_seqs + attn_out = attn_forward_gen( + layer, + fused_q.view([num_tokens, heads * (lora + rope)]), + md, + helix_v_up=True, + q_rows_external=True, + hooks=hooks, + attention_input_type=AttentionInputType.generation_only, + out_scale=None, + latent_cache=latent_cache, + q_pe=q_pe, + cu_q_seqlens=torch.empty(num_seqs + 1, dtype=torch.int32, device=device), + cu_kv_seqlens=torch.empty(num_seqs + 1, dtype=torch.int32, device=device), + fmha_scheduler_counter=torch.empty(1, dtype=torch.uint32, device=device), + mla_bmm1_scale=torch.empty(2, dtype=torch.float32, device=device), + mla_bmm2_scale=torch.empty(1, dtype=torch.float32, device=device), + quant_q_buffer=torch.empty( + num_tokens, heads, lora + rope, dtype=torch.uint8, device=device + ), + ) + assert attn_out.shape == (num_tokens, layer.num_heads_tp_cp * layer.v_head_dim) + return attn_out + + +_F = _features() + +# Generation-only Helix batches of at most this many rows (V up-projection before the +# all-to-all) take the side-gate forward: the zero-KV mask and the output gate +# sigmoid(g_proj(x)) run on a side stream, and o_proj gates its input in its GEMM. +_K3_MLA_GATE_SIDE_MAX_M = 128 +# features.mla_q_norm_pdl: the q_a RMSNorm releases the q_b GEMM at launch instead of after its +# store (``q_norm_pdl``; same kernel code), so q_b streams its weights while q_a and the norm run. +_K3_MLA_Q_NORM_PDL = _F.mla_q_norm_pdl +# features.ull3_mla_gate_at (side-gate forward with the fused Helix tail): when the side-stream +# output gate GEMM is forked. "decode" (default): right before the decode, joined before the +# tail; "start": at the layer start with the zero-KV mask, beside q_a / kv_a; "exchange": right +# after the decode is launched and joined before o_proj, which then gates its input itself. A +# value ":," picks the first placement up to T tokens (default 16) and the +# second above, e.g. "exchange:16,start". Scheduling only. +_ULL3_GATE_PLACES = ("decode", "start", "exchange") + + +def _parse_ull3_gate_at(value: str) -> Tuple[str, int, str]: + small, _, large = value.partition(",") + small, _, limit = small.partition(":") + large = large or small + if small not in _ULL3_GATE_PLACES or large not in _ULL3_GATE_PLACES: + raise ValueError( + f"features.ull3_mla_gate_at={value!r}: expected [:][,] with " + f"places {_ULL3_GATE_PLACES}" + ) + return small, int(limit) if limit else 16, large + + +_K3_ULL3_MLA_GATE_AT = _parse_ull3_gate_at(_F.ull3_mla_gate_at) + + +def ull3_gate_place(num_tokens: int) -> str: + """The features.ull3_mla_gate_at placement of a batch of ``num_tokens`` tokens.""" + small, limit, large = _K3_ULL3_MLA_GATE_AT + return small if num_tokens <= limit else large + + +# One side stream per device, shared by all MLA layers (their side work never overlaps). +_GATE_SIDE_STREAMS: dict[int, torch.cuda.Stream] = {} + + +def _gate_side_stream(device: torch.device) -> torch.cuda.Stream: + index = device.index if device.index is not None else torch.cuda.current_device() + stream = _GATE_SIDE_STREAMS.get(index) + if stream is None: + stream = torch.cuda.Stream(device=index, priority=0) + _GATE_SIDE_STREAMS[index] = stream + return stream + + +# --------------------------------------------------------------------------------------------- +# Identity RoPE table. +# --------------------------------------------------------------------------------------------- + +# Identity (cos, sin) = (1, 0) tables by (device index, float count), shared by every layer. +_IDENTITY_ROPE_TABLES: dict[Tuple[int, int], torch.Tensor] = {} + + +def identity_rope_table(num_floats: int, device: torch.device) -> torch.Tensor: + """A ``[1, num_floats]`` FP32 table of interleaved ``(cos, sin) = (1, 0)`` pairs: the MLA RoPE + kernels' ``rotary_cos_sin`` layout, with every rotation the identity.""" + device = torch.device(device) + index = device.index if device.index is not None else torch.cuda.current_device() + key = (index, num_floats) + table = _IDENTITY_ROPE_TABLES.get(key) + if table is None: + table = torch.zeros(1, num_floats, dtype=torch.float32, device=device) + table.view(-1)[0::2] = 1.0 + if not torch.cuda.is_current_stream_capturing(): + # Written before any kernel on another stream reads it. + torch.cuda.synchronize(device) + _IDENTITY_ROPE_TABLES[key] = table + return table + + +def _install_identity_rope(backend: TrtllmAttention) -> None: + """Give ``backend`` the identity table in place of its sinusoid table (same shape), and grow + it as an identity table where the backend needs more positions.""" + cos_sin = backend.rotary_cos_sin + if cos_sin is None: + raise RuntimeError("the MLA attention backend has no RoPE table to replace") + backend.rotary_cos_sin = identity_rope_table(cos_sin.numel(), cos_sin.device) + rope = backend.rope_params + floats_per_position = rope.dim * (2 if rope.duplicate_data else 1) + + def ensure_table_size(required_max_positions: int) -> None: + table_positions = backend.rotary_cos_sin.numel() // floats_per_position + rope.max_positions = max(rope.max_positions, table_positions, required_max_positions) + if required_max_positions > table_positions: + backend.rotary_cos_sin = identity_rope_table( + rope.max_positions * floats_per_position, backend.rotary_cos_sin.device + ) + + backend._ensure_rope_table_size = ensure_table_size + + +def _pos_embd_params(qk_rope_head_dim: int, max_position_embeddings: int): + """The RoPE configuration the MLA kernels index the table with: GPT-NeoX type, interleaved + (cos, sin) pairs duplicated over the rope dim (the FP4 MLA cache update requires it).""" + return PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams( + dim=qk_rope_head_dim, + theta=10000.0, + max_positions=max_position_embeddings, + original_max_positions=max_position_embeddings, + duplicate_data=True, + ), + is_neox=False, + ) + + +# --------------------------------------------------------------------------------------------- +# The attention part under breakable CUDA graphs: the layer is found by its index string, so the +# break's frozen arguments are tensors and a string. +# --------------------------------------------------------------------------------------------- + + +def _registered_layer(layer_idx: str): + """The step's attention metadata and the MLA layer registered under ``layer_idx``.""" + extra_attrs = get_model_extra_attrs() + assert extra_attrs is not None, "Model extra attrs is not set" + metadata_ref = extra_attrs.get("attention_metadata", None) + assert metadata_ref is not None, "Attention metadata is not set" + metadata = metadata_ref() + assert isinstance(metadata, TrtllmAttentionMetadata), "the MLA metadata must be TRTLLM's" + mla_layers = extra_attrs.get("mla_layers", None) + assert mla_layers is not None, "Attention layer is not registered" + layer_ref = mla_layers.get(layer_idx, None) + assert layer_ref is not None, f"Cannot find attention layer for layer {layer_idx}" + layer = layer_ref() + assert isinstance(layer, K3MLAAttention), "the registered MLA layer must be K3's" + return metadata, layer + + +def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: + """The attention output buffer of the registered layer ``layer_idx`` for these rows.""" + _, layer = _registered_layer(layer_idx) + return layer.create_output(hidden_states) + + +def mla_custom_op_inplace( + hidden_states: torch.Tensor, layer_idx: str, output: torch.Tensor +) -> None: + """The attention part of the registered layer ``layer_idx`` into ``output``, on the step's + metadata registered with it.""" + metadata, layer = _registered_layer(layer_idx) + layer.run_attention(hidden_states, metadata, output) + + +_bcg_mla_custom_op_inplace = eager_on_graph(mla_custom_op_inplace) + + +# --------------------------------------------------------------------------------------------- +# K3MLAAttention. +# --------------------------------------------------------------------------------------------- + + +class K3MLAAttention(nn.Module): + """Kimi K3 MLA (see the module docstring).""" + + def __init__( + self, + *, + hidden_size: int, + num_heads: int, + q_lora_rank: int, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + rms_norm_eps: Optional[float] = None, + dtype: Optional[torch.dtype] = None, + layer_idx: int = 0, + use_output_gate: bool = True, + max_position_embeddings: int = 8192, + model_config: ModelConfig, + aux_stream_dict: dict[AuxStreamType, torch.cuda.Stream], + mapping_with_cp: Optional[Mapping] = None, + ) -> None: + super().__init__() + config = model_config + self.layer_idx = layer_idx + self.layer_idx_str = str(layer_idx) + self.dtype = dtype + self.hidden_size = hidden_size + self.num_heads = num_heads + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.max_position_embeddings = max_position_embeddings + self.use_output_gate = use_output_gate + # The checkpoint's separate q_a / kv_a rows load into one fused projection. + self.fuse_qkv_a_proj = True + self.out_scale = None + if rms_norm_eps is None: + rms_norm_eps = getattr(config.pretrained_config, "rms_norm_eps", None) + if rms_norm_eps is None: + rms_norm_eps = 1e-6 + self.kv_cache_dtype = config.extra_attrs.get("kv_cache_dtype", "auto") + + # Each projection's checkpoint recipe: FP8 128 x 128 block scales or BF16. + projection_configs = config.quant_config_dict or {} + + def is_fp8(name: str) -> bool: + algo = getattr(projection_configs.get(name, config.quant_config), "quant_algo", None) + if algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError(f"Kimi K3 MLA: {name} has unsupported quantization {algo}") + return algo == QuantAlgo.FP8_BLOCK_SCALES + + fp8 = { + name: is_fp8(name) + for name in ("q_a_proj", "kv_a_proj_with_mqa", "q_b_proj", "kv_b_proj", "o_proj") + + (("g_proj",) if use_output_gate else ()) + } + self._reject_unsupported(config, mapping_with_cp, fp8) + + # Registered for the custom ops; a later stack (one-model speculative decoding) with the + # same local layer index registers as "_". + self.register_to_config = True + registry = config.extra_attrs.setdefault("mla_layers", {}) + suffix = 0 + while self.layer_idx_str in registry: + self.layer_idx_str = f"{layer_idx}_{suffix}" + suffix += 1 + registry[self.layer_idx_str] = weakref.ref(self) + + # Head sharding: q_b / kv_b over TP (attention DP: none), the post-exchange heads of + # o_proj / g_proj over TP x CP. + self.mapping = mapping_with_cp if mapping_with_cp is not None else config.mapping + tp_size, pp_size, cp_size = self.mapping.tp_size, self.mapping.pp_size, self.mapping.cp_size + dp_size = 1 + if self.mapping.enable_attention_dp: + dp_size, tp_size = tp_size, 1 + world_size = pp_size * dp_size * tp_size * cp_size + mapping = Mapping( + world_size=world_size, + tp_size=tp_size, + pp_size=pp_size * dp_size, + cp_size=cp_size, + cp_config=self.mapping.cp_config, + rank=self.mapping.rank, + gpus_per_node=self.mapping.gpus_per_node, + enable_attention_dp=self.mapping.enable_attention_dp, + ) + self.mapping_o = Mapping( + world_size=world_size, + tp_size=tp_size * cp_size, + pp_size=pp_size * dp_size, + cp_size=1, + rank=self.mapping.rank, + gpus_per_node=self.mapping.gpus_per_node, + enable_attention_dp=self.mapping.enable_attention_dp, + ) + assert num_heads % (tp_size * cp_size) == 0 + self.num_heads_tp = num_heads // tp_size + self.num_heads_tp_cp = self.num_heads_tp // cp_size + # Helix context attends this CP rank's head range only. + self.num_context_heads = ( + self.num_heads_tp_cp if self.mapping.has_cp_helix() else self.num_heads_tp + ) + + self.aux_stream = (aux_stream_dict or {}).get(AuxStreamType.Attention) + self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + + def projection(fp8: bool, in_features: int, out_features: int, mode, mp: Optional[Mapping]): + shard = {} if mp is None else {"tp_size": mp.tp_size, "tp_rank": mp.tp_rank} + if fp8: + return Fp8BlockScaleLinear(in_features, out_features, tp_mode=mode, **shard) + return Bf16Linear(in_features, out_features, dtype=dtype, tp_mode=mode, **shard) + + column, row = TensorParallelMode.COLUMN, TensorParallelMode.ROW + self.kv_a_proj_with_mqa = projection( + fp8["kv_a_proj_with_mqa"], + hidden_size, + q_lora_rank + kv_lora_rank + qk_rope_head_dim, + None, + None, + ) + if fp8["kv_a_proj_with_mqa"]: + # q_lora + kv_lora + rope rows (2112 for K3) leave a partial last 128-row tile, + # which the FP8 low-M GEMM serves only on request. + self.kv_a_proj_with_mqa.fp8w_lowm_partial_n_tile = True + self.q_a_layernorm = RMSNormWeight(q_lora_rank, rms_norm_eps, dtype=dtype) + self.q_b_proj = projection( + fp8["q_b_proj"], q_lora_rank, num_heads * self.qk_head_dim, column, mapping + ) + self.kv_a_layernorm = RMSNormWeight(kv_lora_rank, rms_norm_eps, dtype=dtype) + self.kv_b_proj = projection( + fp8["kv_b_proj"], + kv_lora_rank, + num_heads * (qk_nope_head_dim + v_head_dim), + column, + mapping, + ) + self._kv_b_fp8 = isinstance(self.kv_b_proj, Fp8BlockScaleLinear) + # This CP rank's post-exchange heads of the V absorption (bound by the loader). + self.v_b_proj = nn.Parameter( + torch.empty((self.num_heads_tp_cp, v_head_dim, kv_lora_rank), dtype=dtype), + requires_grad=False, + ) + # Helix only: the BF16 V weight of every tp-local head, applied to this rank's partials + # before the all-to-all (``project_helix_v_partials``). + self.register_buffer("helix_v_b_proj_full", None, persistent=False) + self.helix_v_up_before_exchange = False + self.o_proj = projection( + fp8["o_proj"], num_heads * v_head_dim, hidden_size, row, self.mapping_o + ) + self.attention_output_hidden_size = self.o_proj.in_features + if use_output_gate: + # o_proj's input sharding (under Helix the post-exchange 1/cp head chunk). + self.g_proj = projection( + fp8["g_proj"], hidden_size, num_heads * v_head_dim, column, self.mapping_o + ) + + quant_config = config.get_quant_config() + self.quant_config = quant_config + pos_embd_params = _pos_embd_params(qk_rope_head_dim, max_position_embeddings) + mla_dims = dict( + pos_embd_params=pos_embd_params, + quant_config=quant_config, + q_scaling=1.0, + is_mla_enable=True, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + predicted_tokens_per_seq=1, + skip_create_weights_in_init=True, + sparse_params=None, + skip_correction_threshold=config.skip_correction_threshold, + ) + self.mqa = create_attention( + "TRTLLM", + layer_idx, + self.num_heads_tp, + head_dim=kv_lora_rank + qk_rope_head_dim, + num_kv_heads=1, + hidden_size=hidden_size, + v_head_dim=kv_lora_rank, + dtype=dtype, + aux_stream=self.aux_stream, + rope_append=True, + kv_cache_dtype=self.kv_cache_dtype, + **mla_dims, + ) + self.mha = create_attention( + "TRTLLM", + layer_idx, + self.num_context_heads, + head_dim=self.qk_head_dim, + num_kv_heads=self.num_context_heads, + v_head_dim=v_head_dim, + **mla_dims, + ) + assert isinstance(self.mha, TrtllmAttention) and isinstance(self.mqa, TrtllmAttention) + if not self.mqa.support_fused_rope(): + raise NotImplementedError("Kimi K3 MLA requires the RoPE fused in the MLA kernels") + _install_identity_rope(self.mha) + _install_identity_rope(self.mqa) + + # Set in create_weights() from the backends' KV cache quantization. + self.quantized_kv = False + self.fp4_kv = False + self.fp8_cached_kv_context = False + self._weights_transformed = False + # features.ull4_mla_plan_early: event of the step metadata the model pre-launched. + self.step_metadata_ready: Optional[torch.cuda.Event] = None + # Side-gate forward: fork / join of the zero-KV mask and the gate on the side stream. + self._gate_events = tuple(torch.cuda.Event() for _ in range(4)) + # Fused query: fork of the KV half after q_a, join before the decode. + self._kv_events = (torch.cuda.Event(), torch.cuda.Event()) + # Set by the decoder layer for the next forward only: the MXFP8 form ``(xq, xs)`` of the + # layer input, which the fused q_a / kv_a projection consumes (see q_a_takes_mxfp8). + self.next_input_mxfp8: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + # Fused query: (q_b weight scale, its block exponents), built on first use. + self._q_fused_qb_exp: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + + if not config.skip_create_weights_in_init: + self.create_weights() + + @staticmethod + def _reject_unsupported( + config: ModelConfig, mapping_with_cp: Optional[Mapping], fp8: dict[str, bool] + ) -> None: + """Configurations whose kernel chain K3 does not carry.""" + mapping = mapping_with_cp if mapping_with_cp is not None else config.mapping + unsupported = [] + if config.attn_backend.upper() != "TRTLLM": + unsupported.append(f"attention backend {config.attn_backend} (TRTLLM only)") + if config.sparse_attention_config is not None: + unsupported.append("sparse attention") + if fp8["q_a_proj"] != fp8["kv_a_proj_with_mqa"]: + unsupported.append("q_a / kv_a with different recipes (no fused q_a / kv_a GEMM)") + if config.extra_attrs.get("kv_cache_dtype", "auto") == "fp8_ds_mla": + unsupported.append("the fp8_ds_mla KV cache") + if hasattr(config.pretrained_config, "llama_4_scaling"): + unsupported.append("Llama-4 attention scaling") + # The CuTe DSL GEMM / BMM opt-ins of main's MLA, where they would select another kernel: + # BF16 projections, BF16 absorption BMMs (BF16 kv_b, or the Helix V up-projection), FP8 + # absorption BMMs. + if config.use_cute_dsl_bf16_gemm and not all(fp8.values()): + unsupported.append("use_cute_dsl_bf16_gemm with BF16 projections") + if config.use_cute_dsl_bf16_bmm and (mapping.has_cp_helix() or not fp8["kv_b_proj"]): + unsupported.append("use_cute_dsl_bf16_bmm") + if config.use_cute_dsl_blockscaling_bmm and fp8["kv_b_proj"]: + unsupported.append("use_cute_dsl_blockscaling_bmm") + if mapping.has_cp_ulysses(): + unsupported.append("Ulysses CP") + if mapping.cp_size > 1 and not mapping.has_cp_helix(): + unsupported.append(f"CP type {mapping.cp_config.get('cp_type')}") + if mapping.has_cp_helix(): + if mapping.enable_attention_dp: + unsupported.append("Helix CP with attention DP") + if mapping.cp_config.get("use_nccl_for_alltoall", True): + unsupported.append("the NCCL Helix all-to-all (use_nccl_for_alltoall: false)") + elif mapping.cp_config.get("fifo_version", 2) != 2: + unsupported.append("the fifo-v1 Helix all-to-all") + if unsupported: + raise NotImplementedError("Kimi K3 MLA does not support: " + "; ".join(unsupported)) + + def create_weights(self) -> None: + """Backend KV cache quantization and the absorption weights (their dtype follows kv_b).""" + self.mha.update_quant_config(self.quant_config) + self.mqa.update_quant_config(self.quant_config) + self.quantized_kv = bool(self.mqa.has_fp8_kv_cache or self.mqa.has_fp4_kv_cache) + self.fp4_kv = bool( + self.mqa.support_fp4_kv_cache() + and self.quant_config is not None + and self.quant_config.layer_quant_mode.has_fp4_kv_cache() + ) + if self.fp4_kv and not self.mqa.rope_params.duplicate_data: + raise ValueError("FP4 MLA requires RopeParams.duplicate_data=True.") + # Cached-prefix context on the FP8 context FMHA takes FP8 Q / K / V (plain conversions at + # the op's unit KV scale). + mha = self.mha + scale = mha.kv_scale_orig_quant + self.fp8_cached_kv_context = bool( + mha.has_fp8_context_mla + and not self.mapping.has_cp_helix() + and (self.qk_nope_head_dim, self.qk_rope_head_dim, self.v_head_dim) == (128, 64, 128) + and isinstance(scale, torch.Tensor) + and bool(torch.all(scale == 1.0).item()) + ) + heads, heads_cp = self.num_heads_tp, self.num_heads_tp_cp + lora, nope, vdim = self.kv_lora_rank, self.qk_nope_head_dim, self.v_head_dim + self.k_b_proj_trans = nn.Parameter( + torch.empty( + (heads, lora, nope), dtype=torch.float8_e4m3fn if self._kv_b_fp8 else self.dtype + ), + requires_grad=False, + ) + self.k_b_proj_trans_dequant = None + self.v_b_proj_dequant = None + if self._kv_b_fp8: + self.k_b_proj_trans_scale = nn.Parameter( + torch.empty((heads, lora // 128, nope // 128), dtype=torch.float32), + requires_grad=False, + ) + self.v_b_proj_scale = nn.Parameter( + torch.empty((heads_cp, vdim // 128, lora // 128), dtype=torch.float32), + requires_grad=False, + ) + if is_sm_100f(): + # The absorption BMMs run in BF16 on dequantized copies. + assert self.dtype == torch.bfloat16 + self.k_b_proj_trans_dequant = nn.Parameter( + torch.empty((heads, lora, nope), dtype=self.dtype), requires_grad=False + ) + self.v_b_proj_dequant = nn.Parameter( + torch.empty((heads_cp, vdim, lora), dtype=self.dtype), requires_grad=False + ) + else: + self.k_b_proj_trans_scale = None + self.v_b_proj_scale = None + self._weights_transformed = False + + # ----------------------------------------------------------------------------------------- + # Post-load. + # ----------------------------------------------------------------------------------------- + + def _bind_helix_v_up_weight(self, weight: torch.Tensor) -> None: + """Apply the V up-projection with ``weight`` (BF16 ``[num_heads_tp, v_head_dim, + kv_lora_rank]``) before the Helix all-to-all.""" + assert self.mapping.has_cp_helix() + assert weight.dtype == torch.bfloat16 + assert tuple(weight.shape) == (self.num_heads_tp, self.v_head_dim, self.kv_lora_rank) + self.helix_v_b_proj_full = weight + self.helix_v_up_before_exchange = True + + def _build_helix_v_up_weight(self) -> None: + """Under Helix, every tp-local head's BF16 V weight from kv_b's loaded rows (BF16, or the + raw FP8 codes with their FP32 128 x 128 block scales: this runs before kv_b's own + transform, which the post-load walk visits after this module).""" + if not self.mapping.has_cp_helix() or self.dtype != torch.bfloat16: + return + weight = self.kv_b_proj.weight + nope_rows = self.num_heads_tp * self.qk_nope_head_dim + num_rows = nope_rows + self.num_heads_tp * self.v_head_dim + v_shape = (self.num_heads_tp, self.v_head_dim, self.kv_lora_rank) + if weight.is_meta or tuple(weight.shape) != (num_rows, self.kv_lora_rank): + return + if weight.dtype == torch.bfloat16: + v_weight = weight[nope_rows:].view(v_shape) + elif weight.dtype == torch.float8_e4m3fn: + scale = self.kv_b_proj.weight_scale + if scale.dtype != torch.float32 or tuple(scale.shape) != ( + -(-num_rows // 128), + -(-self.kv_lora_rank // 128), + ): + return + row_blocks = torch.arange(nope_rows, num_rows, device=weight.device) // 128 + v_scale = scale[row_blocks].repeat_interleave(128, dim=1)[:, : self.kv_lora_rank] + v_weight = (weight[nope_rows:].float() * v_scale).to(self.dtype).view(v_shape) + if self.v_b_proj_dequant is not None: + # This rank's dequantized V chunk holds the same values. + start = self.mapping.cp_rank * self.num_heads_tp_cp + self.v_b_proj_dequant = nn.Parameter( + v_weight[start : start + self.num_heads_tp_cp], requires_grad=False + ) + else: + return + self._bind_helix_v_up_weight(v_weight) + + def transform_weights(self) -> None: + if self._weights_transformed: + return + self._build_helix_v_up_weight() + self._weights_transformed = True + + def cache_derived_state(self) -> None: + """The scales of the q_b / kv_b output slices a Helix rank's context heads read.""" + if self.num_context_heads != self.num_heads_tp: + first_head = self.mapping.cp_rank * self.num_context_heads + last_head = first_head + self.num_context_heads + value_start = self.num_heads_tp * self.qk_nope_head_dim + for proj, start, end in ( + (self.q_b_proj, first_head * self.qk_head_dim, last_head * self.qk_head_dim), + ( + self.kv_b_proj, + first_head * self.qk_nope_head_dim, + last_head * self.qk_nope_head_dim, + ), + ( + self.kv_b_proj, + value_start + first_head * self.v_head_dim, + value_start + last_head * self.v_head_dim, + ), + ): + if not isinstance(proj, Bf16Linear): + proj.prepare_output_slice(start, end) + self._weights_transformed = True + + def post_load_weights(self) -> None: + self.transform_weights() + self.cache_derived_state() + + # ----------------------------------------------------------------------------------------- + # Front. + # ----------------------------------------------------------------------------------------- + + def q_a_mx_atom(self, num_rows: int) -> bool: + """Whether the fused q_a / kv_a GEMM takes the MXFP8 rows with the scales in the mid-M + GEMM's layout (``takes_mid_mxfp8_input``).""" + takes_mid = getattr(self.kv_a_proj_with_mqa, "takes_mid_mxfp8_input", None) + return takes_mid is not None and bool(takes_mid(num_rows)) + + def q_a_takes_mxfp8(self, num_rows: int) -> bool: + """Whether the fused q_a / kv_a GEMM of ``num_rows`` rows can take their MXFP8 form + (``next_input_mxfp8``).""" + takes = getattr(self.kv_a_proj_with_mqa, "takes_mxfp8_input", None) + return takes is not None and (bool(takes(num_rows)) or self.q_a_mx_atom(num_rows)) + + def project_qkv_a(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The fused q_a / kv_a projection ``[T, q_lora | kv_lora | rope]``: on the MXFP8 rows the + decoder layer handed over for these rows (``forward_mxfp8``), else the GEMM on the BF16 + rows (a BF16 weight: ``cublas_mm``).""" + mxfp8_input, self.next_input_mxfp8 = self.next_input_mxfp8, None + proj = self.kv_a_proj_with_mqa + if isinstance(proj, Bf16Linear): + return cublas_mm( + hidden_states, + proj.weight.t(), + None, + out_dtype=None, + output_buffer_kind=int(BufferKind.DEFAULT), + group=None, + ) + if mxfp8_input is not None and mxfp8_input[0].shape[0] == hidden_states.shape[0]: + return proj.forward_mxfp8(hidden_states, *mxfp8_input) + return proj(hidden_states) + + def q_a_norm(self, q: torch.Tensor) -> torch.Tensor: + """The q_a RMSNorm (FlashInfer); with ``features.mla_q_norm_pdl`` and PDL the same kernel + releasing its dependent (q_b) at launch.""" + norm = self.q_a_layernorm + if _K3_MLA_Q_NORM_PDL and IS_FLASHINFER_AVAILABLE and get_env_enable_pdl(): + from .kernels.mla.q_norm_pdl import can_use + from .ops.attention.rmsnorm_release_at_launch import rmsnorm_release_at_launch + + if can_use(q, norm.weight): + return rmsnorm_release_at_launch(q, norm.weight, norm.variance_epsilon) + return rms_norm(q, norm.weight, norm.variance_epsilon) + + @staticmethod + def forward_slice(proj, x: torch.Tensor, start: int, end: int) -> torch.Tensor: + """Output features ``[start, end)`` of a column-parallel projection, contiguous.""" + if isinstance(proj, Bf16Linear): + return torch.nn.functional.linear(x, proj.weight[start:end]) + return proj.forward_slice(x, start, end) + + def _project_q_by_phase( + self, q: torch.Tensor, num_ctx_tokens: int + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """q_b of the context rows (a Helix rank's context heads only) and of the generation rows.""" + if self.num_context_heads == self.num_heads_tp or num_ctx_tokens == 0: + projected = self.q_b_proj(q) + return projected[:num_ctx_tokens], projected[num_ctx_tokens:] + first = self.mapping.cp_rank * self.num_context_heads * self.qk_head_dim + last = first + self.num_context_heads * self.qk_head_dim + ctx = self.forward_slice(self.q_b_proj, q[:num_ctx_tokens], first, last) + gen = self.q_b_proj(q[num_ctx_tokens:]) if q.shape[0] > num_ctx_tokens else None + return ctx, gen + + def _forward_attention( + self, + hidden_states: torch.Tensor, + md: AttentionMetadata, + output: torch.Tensor, + hooks: Optional[HelixGenHooks], + ) -> None: + """Front, context and generation attention of the batch's rows into ``output``.""" + num_ctx_tokens = md.num_ctx_tokens + num_tokens = md.num_tokens + hidden_states = hidden_states[:num_tokens, ...] + q, compressed_kv, k_pe = self.project_qkv_a(hidden_states).split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 + ) + kv_norm = self.kv_a_layernorm + q, compressed_kv = fork_join( + lambda: self.q_a_norm(q), + lambda: rms_norm(compressed_kv, kv_norm.weight, kv_norm.variance_epsilon), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + (q_ctx, q_gen), latent_cache = fork_join( + lambda: self._project_q_by_phase(q, num_ctx_tokens), + lambda: torch.concat([compressed_kv, k_pe], dim=-1), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + if md.num_contexts > 0: + forward_context( + self, + q_ctx, + compressed_kv[:num_ctx_tokens, ...], + md, + output[:num_ctx_tokens, :], + latent_cache[:num_ctx_tokens, ...], + ) + if md.num_generations > 0: + assert q_gen is not None + forward_absorption_generation( + self, + q_gen, + latent_cache[num_ctx_tokens:, ...], + md, + output[num_ctx_tokens:num_tokens, :], + hooks, + ) + + def run_attention( + self, + hidden_states: torch.Tensor, + md: AttentionMetadata, + output: torch.Tensor, + hooks: Optional[HelixGenHooks] = None, + ) -> None: + """The attention part of a forward (everything before the output gate) into ``output``, + on the mixed-precision KV staging copy where the metadata has one.""" + with md.mixed_precision_kv_staging(self.layer_idx): + self._forward_attention(hidden_states, md, output, hooks) + + def create_output(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The attention output ``[rows, H / P * v_head_dim]`` for the layer input's rows.""" + return hidden_states.new_empty( + [hidden_states.shape[0], self.attention_output_hidden_size], dtype=hidden_states.dtype + ) + + # ----------------------------------------------------------------------------------------- + # Forward. + # ----------------------------------------------------------------------------------------- + + def _gate_and_o_proj(self, hidden_states: torch.Tensor, attn_out: torch.Tensor) -> torch.Tensor: + if self.use_output_gate: + attn_out = attn_out * self.g_proj(hidden_states).sigmoid() + return self.o_proj(attn_out) + + def _side_gate_applies(self, hidden_states: torch.Tensor, md: AttentionMetadata) -> bool: + """Whether this forward takes :meth:`_forward_side_gate`.""" + return ( + self.use_output_gate + and hidden_states.dtype == torch.bfloat16 + and hidden_states.dim() == 2 + and hidden_states.shape[0] == md.num_tokens + and md.num_contexts == 0 + and self.helix_v_up_before_exchange + and hasattr(self.o_proj, "gated_applies") + and hidden_states.shape[0] <= _K3_MLA_GATE_SIDE_MAX_M + ) + + def _forward_side_gate( + self, hidden_states: torch.Tensor, md: AttentionMetadata + ) -> torch.Tensor: + """Generation forward with the zero-KV mask and the output gate on a side stream. + + The mask (one small kernel) is forked once the layer input is written; the gate (one GEMM + with a sigmoid epilogue) is forked when the all-to-all is enqueued (or per + ``features.ull3_mla_gate_at``), so it runs on the SMs the all-to-all leaves idle. The main + stream waits for the mask where the exchange reads it and for the gate before o_proj; with + the fused Helix tail and an early gate, before the tail kernel, so o_proj depends on the + tail alone and gates its input in its GEMM (``forward_gated(_mxfp8)``).""" + num_tokens = hidden_states.shape[0] + mask_fork, mask_ready, gate_fork, gate_ready = self._gate_events + multi_stream = do_multi_stream() + side = _gate_side_stream(hidden_states.device) if multi_stream else None + gates = [] + mask_joined = [not multi_stream] + gate_joined = [not multi_stream] + + def gate_of(x): + forward_sigmoid = getattr(self.g_proj, "forward_sigmoid", None) + if forward_sigmoid is not None: + return forward_sigmoid(x) + return self.g_proj(x).sigmoid() + + if multi_stream: + mask_fork.record() + with torch.cuda.stream(side): + mask_fork.wait() + zero_kv_mask = helix_gen_zero_kv_mask(md, num_tokens) + mask_ready.record() + else: + zero_kv_mask = helix_gen_zero_kv_mask(md, num_tokens) + + def zero_kv_mask_fn(): + if not mask_joined[0]: + mask_ready.wait() + mask_joined[0] = True + return zero_kv_mask + + def pre_exchange(): + if gates: + return # already forked (gate place "start") + if multi_stream: + gate_fork.record() + with torch.cuda.stream(side): + gate_fork.wait() + gates.append(gate_of(hidden_states)) + gate_ready.record() + else: + gates.append(gate_of(hidden_states)) + + def pre_tail(): + if gates and not gate_joined[0]: + gate_ready.wait() + gate_joined[0] = True + return gates[0] if gates else None + + gate_place = ull3_gate_place(num_tokens) if multi_stream else "decode" + if gate_place == "start": + pre_exchange() + hooks = HelixGenHooks(zero_kv_mask_fn, pre_exchange, pre_tail, gate_place) + if q_fused_applies(self, hidden_states, md): + attn_out = forward_gen_q_fused(self, hidden_states, md, hooks) + else: + output = self.create_output(hidden_states) + self.run_attention(hidden_states, md, output, hooks) + attn_out = hooks.gen_output if hooks.gen_output is not None else output + if not mask_joined[0]: + mask_ready.wait() + if gates: + gate = gates[0] + if not gate_joined[0]: + gate_ready.wait() + else: + gate = gate_of(hidden_states) + if hooks.tail_mx is not None: + return self.o_proj.forward_gated_mxfp8(attn_out, gate, *hooks.tail_mx) + if self.o_proj.gated_applies(attn_out, gate): + return self.o_proj.forward_gated(attn_out, gate) + return self.o_proj(attn_out * gate) + + def prelaunch_step_metadata(self, attn_metadata: AttentionMetadata) -> bool: + """``features.ull4_mla_plan_early``: launch, on the current stream, the step metadata the + first FP4 MLA generation kernel would write (generation lengths fused with the packed + split plan of this layer's query heads, then the device page table). Returns whether + anything was launched; the caller sets :attr:`step_metadata_ready` to an event recorded + after it.""" + if not fp4_mla_decode.applies(self.mqa) or attn_metadata.num_contexts != 0: + return False + use_helix = ( + getattr(attn_metadata, "helix_position_offsets", None) is not None + or getattr(attn_metadata, "helix_is_inactive_rank", None) is not None + ) + return fp4_mla_decode.prelaunch_generation_metadata( + attn_metadata, + num_heads=self.num_heads_tp, + use_helix=use_helix, + device=self.o_proj.weight.device, + ) + + def forward( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + """o_proj's partial output of the layer for ``hidden_states`` (K3 is NoPE: + ``position_ids`` is None).""" + assert position_ids is None, "Kimi K3 MLA is NoPE" + if not isinstance(hidden_states, torch.Tensor): + raise TypeError(f"Kimi K3 MLA takes BF16 rows, got {type(hidden_states).__name__}") + ready = self.step_metadata_ready + if ready is not None: + # features.ull4_mla_plan_early: the step metadata written on the model's side stream. + self.step_metadata_ready = None + torch.cuda.current_stream().wait_event(ready) + try: + if self._side_gate_applies(hidden_states, attn_metadata): + return self._forward_side_gate(hidden_states, attn_metadata) + if q_fused_applies(self, hidden_states, attn_metadata): + attn_out = forward_gen_q_fused(self, hidden_states, attn_metadata) + return self._gate_and_o_proj(hidden_states, attn_out) + if self.register_to_config: + output = create_mla_outputs(hidden_states, self.layer_idx_str) + _bcg_mla_custom_op_inplace(hidden_states, self.layer_idx_str, output) + else: + output = self.create_output(hidden_states) + self.run_attention(hidden_states, attn_metadata, output) + return self._gate_and_o_proj(hidden_states, output) + finally: + # Consumed by the projection; never carried into another forward. + self.next_input_mxfp8 = None + + +# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table +# holds cos=1/sin=0), but the chunked-context path indexes the table by +# absolute position, so it must cover max_position_embeddings (~512MB per +# backend for the 1M-position checkpoint); a smaller table is read out of +# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context +# deployments. +_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" + + +_KIMI_K3_FP8_KV_LAYERS_ENV = "TLLM_KIMI_K3_FP8_KV_LAYERS" + + +def _kimi_k3_fp8_kv_layer_ids(cfg, kv_quant_algo: Optional[QuantAlgo]) -> frozenset: + """MLA layers whose KV cache falls back to FP8 under NVFP4 KV (``TLLM_KIMI_K3_FP8_KV_LAYERS``). + + On the Rubin branch the NVFP4 default is the last two MLA layers and the FP4 MLA cache + manager allocates those layers as FP8 (``kimi_k3_cache_policy.get_kimi_k3_fp8_kv_layer_ids``, + read by the manager and here). Engine requirement: main's cache policy has no FP8 selection + (only ``TLLM_KIMI_K3_BF16_KV_LAYERS``) and its cache manager allocates every MLA layer in the + model-wide dtype, so switching a layer to FP8 here alone would mismatch the cache. The + override is therefore taken only when the engine's cache policy provides the selector; + until then every layer keeps the model-wide KV dtype and a warning names the gap. + """ + selector = getattr(_kimi_k3_cache_policy, "get_kimi_k3_fp8_kv_layer_ids", None) + if selector is not None: + return selector(cfg, nvfp4_enabled=kv_quant_algo == QuantAlgo.NVFP4) + requested = os.environ.get(_KIMI_K3_FP8_KV_LAYERS_ENV) + if (requested is None and kv_quant_algo == QuantAlgo.NVFP4) or (requested or "").strip(): + logger.warning_once( + f"Kimi K3: FP8 KV-cache fallback layers ({_KIMI_K3_FP8_KV_LAYERS_ENV}=" + f"{requested if requested is not None else ''}) are not " + "applied: the KV cache manager has no per-layer FP8 allocation on this tree. Every " + "MLA layer uses the model-wide KV cache dtype.", + key="kimi_k3_fp8_kv_layers_unsupported", + ) + return frozenset() + + +def resolve_attention_quant_config( + config: ModelConfig | None, layer_idx: int, projection: str +) -> QuantConfig: + """Resolve a checkpoint projection, including mixed-precision exclusions.""" + if config is None: + return QuantConfig() + global_config = config.quant_config or QuantConfig() + names = [ + f"{prefix}layers.{layer_idx}.self_attn.{projection}" + for prefix in ("language_model.model.", "model.", "") + ] + if any(global_config.is_module_excluded_from_quantization(name) for name in names): + return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) + declarations = config.quant_config_dict or {} + matches = [declarations[name] for name in names if name in declarations] + if matches: + selected = matches[0] + if any(match.quant_algo != selected.quant_algo for match in matches[1:]): + raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") + elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: + selected = QuantConfig() + else: + selected = global_config + if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError( + f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " + f"quantization {selected.quant_algo}" + ) + if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): + raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") + result = copy.copy(selected) + result.kv_cache_quant_algo = global_config.kv_cache_quant_algo + return result + + +# --------------------------------------------------------------------------- +# MLA runtime. +# --------------------------------------------------------------------------- + + +class K3MLARuntime(nn.Module): + """Wraps K3 MLA and applies its external TP output reduction.""" + + def __init__( + self, + cfg: "PretrainedConfig", + layer_idx: int, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + mapping_with_cp: Optional[Mapping] = None, + reduce_output: bool = True, + ) -> None: + super().__init__() + + from .mla import K3MLAAttention + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # K3MLAAttention owns MLA projection/head sharding. Keep only the + # final output reduction in this wrapper so the output gate remains + # between attention and the row-parallel o_proj. + # Helix: mapping_with_cp (the CP original) activates the base MLA's + # helix machinery; this wrapper's allreduce over the repurposed + # mapping sums the base o_proj's tp*cp partials. reduce_output=False + # leaves that reduction to the caller (helix_a2a reduce-scatter). + mapping = model_config.mapping.mapping_for_full_rows() + reduce_output = reduce_output and not mapping.enable_attention_dp and mapping.tp_size > 1 + self._o_allreduce = ( + TpAllReduce(mapping, torch.bfloat16, model_config.allreduce_strategy) + if reduce_output + else None + ) + # Set by the decoder layer for one forward: a pushed o_proj all-reduce may be returned + # unfinished, as a LamportPending that A_mid finishes. + self.ar_finish_deferred = False + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in ( + "q_a_proj", + "kv_a_proj_with_mqa", + "q_b_proj", + "kv_b_proj", + "g_proj", + "o_proj", + ) + } + kv_quant_algo = ( + model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + fp8_kv_layer_ids = _kimi_k3_fp8_kv_layer_ids(cfg, kv_quant_algo) + if layer_idx in fp8_kv_layer_ids: + if kv_quant_algo != QuantAlgo.NVFP4: + raise ValueError( + f"{_KIMI_K3_FP8_KV_LAYERS_ENV} requires the model-wide KV cache dtype to be NVFP4." + ) + kv_quant_algo = QuantAlgo.FP8 + logger.info_once( + f"Kimi K3 MLA layers using FP8 KV-cache fallback: {sorted(fp8_kv_layer_ids)}.", + key="kimi_k3_fp8_kv_fallback_layers", + ) + for projection_config in attention_config.quant_config_dict.values(): + projection_config.kv_cache_quant_algo = kv_quant_algo + attention_config.quant_config = QuantConfig(kv_cache_quant_algo=kv_quant_algo) + attention_config._frozen = model_config._frozen + self.mixer = K3MLAAttention( + hidden_size=cfg.hidden_size, + num_heads=cfg.num_attention_heads, + q_lora_rank=cfg.q_lora_rank, + kv_lora_rank=cfg.kv_lora_rank, + qk_nope_head_dim=cfg.qk_nope_head_dim, + qk_rope_head_dim=cfg.qk_rope_head_dim, + v_head_dim=cfg.v_head_dim, + rms_norm_eps=cfg.rms_norm_eps, + dtype=torch.bfloat16, + layer_idx=layer_idx, + use_output_gate=cfg.mla_use_output_gate, + max_position_embeddings=max_positions, + model_config=attention_config, + aux_stream_dict=aux_stream_dict, + mapping_with_cp=mapping_with_cp, + ) + + def forward( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + o_proj = self.mixer.o_proj + # o_proj stores its partial straight into the all-reduce's Lamport buffers where the + # low-M GEMM serves it. + ar_push = None + if self._o_allreduce is not None and hasattr(o_proj, "next_ar_push"): + ar_push = lamport_push_target( + self._o_allreduce, + hidden_states.shape[0], + o_proj.out_features, + deferred=self.ar_finish_deferred, + ) + o_proj.next_ar_push = ar_push + # MLA.forward takes position_ids first; K3 is NoPE, so pass None. + out = self.mixer(None, hidden_states, attn_metadata) + if ar_push is not None: + o_proj.next_ar_push = None + if ar_push is not None and ar_push.pushed: + # A_mid may finish the all-reduce (a LamportPending). + out_shape = out.shape + out = finish_lamport_push(self._o_allreduce, ar_push, deferred=self.ar_finish_deferred) + if not isinstance(out, LamportPending): + out = out.view(out_shape) + elif ar_push is not None and ar_push.deferred: + # o_proj ran without the push epilogue; A_mid finishes. + lamport_push_tensor(self._o_allreduce, ar_push, out.reshape(ar_push.output.shape)) + out = finish_lamport_push(self._o_allreduce, ar_push, deferred=True) + elif self._o_allreduce is not None: + # Head-sharded TP: sum the row-sharded o_proj partials across + # the head-shard group. + out = self._o_allreduce(out) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/modeling.py new file mode 100644 index 000000000000..3ce416d9b1a7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/modeling.py @@ -0,0 +1,1282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 on Rubin (SM107): the model. + +``ModelingV2KimiK3Nvfp4Sm107Default`` serves every deployment ``routing.py`` admits on SM107: +one-rank DWDP context servers, generation over Helix CP8, over Helix CP16 with the routed +experts split MoE TP4 x EP4, and attention DP over 16 ranks. The layout comes from the engine's +``Mapping``; the deployment's optimizations come from its ``features.py`` profile, selected +before this module is imported. + +It is the vision-language model: the vision tower (MoonViT3d) and its projector, and the text +model ``KimiK3DefaultForCausalLM``: + +* 93 decoder layers (``KimiK3DefaultDecoderLayer``): 69 KDA (Kimi Delta Attention, linear + attention) layers and 24 MLA (absorbed MQA, NoPE) layers, per the 1-indexed + ``linear_attn_config.kda_layers`` / ``full_attn_layers`` schedule. +* Layer 0 has a dense SiTU MLP; layers 1..92 the latent MoE of ``moe_runtime.py``: 896 routed + experts at top-16 (sigmoid + ``e_score_correction_bias`` routing), 2 shared experts, and the + latent down / norm / up projections around the routed experts. +* Attention residuals (``attn_res.py``): before ``input_layernorm``, before + ``post_attention_layernorm`` and at the model output each token mixes a bank of earlier + sublayer outputs; a snapshot is appended whenever ``layer_idx % attn_res_block_size == 0``. + +Caching: per request, a KDA layer keeps a short-convolution window +``[3 * num_heads * head_dim, W]`` (bf16, sections ``[q | k | v]``) and a delta-rule recurrent +state ``[num_heads, head_dim, head_dim]`` (fp32, ``[H, V, K]``) in ``KimiK3HybridCacheManager`` +(KV cache manager V2); an MLA layer uses the paged KV cache with one KV head of +``kv_lora_rank + qk_rope_head_dim`` (576). + +Parallelism: attention is head-split over the TP group, its KV sequence split over the CP +group under Helix, or data-parallel under attention DP. The routed experts are split MoE TP x +EP over the TP group (EP only unless ``moe_tensor_parallel_size`` / +``moe_expert_parallel_size`` are set). Routing is computed on every rank, and the routed +partial sums are all-reduced in the latent space, before ``routed_expert_norm`` / +``routed_expert_up_proj``. Without attention DP the shared experts are tensor-parallel over +the TP group and reduce on an auxiliary stream while the routed experts run on the main +stream; under attention DP they are replicated. ``lm_head`` is vocab-parallel with a gather, +so every rank has the full logits. + +Speculative decoding: the one-engine DSpark / DFlash drafters of ``drafter.py`` / ``dflash.py``, and the suffix +automaton; KDA and MLA verify several tokens per step and keep the accepted state. Chunked +prefill continues a request's KDA state and attends over its cached MLA prefix. KV cache +block reuse is opt-in (``kv_cache_config.enable_block_reuse``); periodic KDA state snapshots +come from ``mamba_state_config.periodic_snapshot_interval``. Pipeline parallelism is not +supported. + +Checkpoint loading is in ``weights.py``. +""" + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import replace +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig, PreTrainedModel + +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_kimi_k3_vl import KimiK3VisionModel +from tensorrt_llm._torch.models.modeling_kimi_k25 import ( + _MEDIA_PLACEHOLDER_TOKEN_ID, + DISAGG, + _has_meta_tensors, +) +from tensorrt_llm._torch.models.modeling_multimodal_utils import ( + find_input_mm_embeds, + fuse_input_embeds, + get_multimodal_embeddings, +) +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, MetaInitException +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, RuntimeHooks +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.logger import logger + +from .attn_res import ( + _K3_ULL_TAIL_SANDWICH, + _apply_attn_res, + _apply_attn_res_add_and_rmsnorm, + _apply_attn_res_and_rmsnorm, + _defer_moe_residual, + _DeferredMoEResidual, + _DeferredRpTail, + _k3_attn_res_lamport_applies, + _k3_attn_res_pdl_moe_add, + _k3_attn_res_pdl_rp_tail, + _materialize_prefix_sum, +) +from .drafter import install_k3_dspark +from .features import active as _features +from .kda import K3KDALinearAttention, K3MambaMetadata +from .kernels.spec_decode_head import spec_decode_head +from .layers import ( + EmbeddingTable, + GatedMlp, + RMSNormWeight, + gather_embedding, + rms_norm, + subgroup_mapping, +) +from .mla import K3MLARuntime, resolve_attention_quant_config +from .moe_runtime import _K3_ULL2_COLL_CP_TAIL, K3MoERuntime +from .ops.attention import _fp4_mla_decode as fp4_mla_decode +from .text_config import _get_text_config, _is_kda_layer, _is_mla_layer +from .weights import KimiK3DefaultTextWeights, KimiK3DefaultWeights + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + +_F = _features() + +_K3_FUSE_MOE_RESIDUAL = _F.fuse_moe_residual + +# CUDA stream priority of the MoE input stream (latent_down chain, routed-output memset). +# Negative values rank above the main and shared-expert streams, so the latent chain that +# gates the routed GEMM1 gets free SMs first. +_K3_LATENT_STREAM_PRIORITY = 0 + +# CUDA stream priority of the stream shared by the attention-side aux work and the MoE shared +# expert. Negative values let the shared-expert chain take SMs ahead of the persistent routed +# GEMMs on the main stream, so its all-reduce does not land after them. +_K3_SHARED_STREAM_PRIORITY = 0 + +# BF16 shared/latent MLP weights are optionally quantized once after loading. +# Checkpoint-native attention FP8 is independent of this lossy conversion. +_KIMI_K3_FP8_WEIGHT_READ_MOE_MLP_ENV = "KIMI_K3_FP8_WEIGHT_READ_MOE_MLP" + + +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" +"""Which residual-stream value the DFlash/DSpark hidden-state tap captures. + +``1`` (default) captures the pre-norm attn_res mixture -- the value the next +consumer actually reads. ``0`` captures the raw running prefix sum instead. + +Both conventions exist in the wild and a drafter distilled against one scores +lower on the other with nothing raised, so this is a property of the DRAFTER +checkpoint, not a performance knob. SGLang (and therefore RadixArk/Kimi-K3-DSpark) +uses the mixture: ``kimi_k3.py _dspark_capture_stream`` -> ``attn_residual.py +aggregate_stream``. vLLM implements both and defaults to the prefix +(``VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=0``, ``models/kimi_k3/nvidia/model.py +_capture_aux_hidden_stream``), which is what a TorchSpec-distilled drafter may +have been trained against. Measured cost of getting it wrong on K3 + RadixArk: +AR 71.4% -> 66.9%. + +Per-checkpoint measurements on K3, GSM8K AL, n=200, TEP8: + +=================== ================== ============ ====== +drafter stream (default 1) prefix (0) delta +=================== ================== ============ ====== +RadixArk (GQA) 71.8% -- -- +Inferact (MLA) 65.7% 66.6% +0.9pt +=================== ================== ============ ====== + +So the default is right for RadixArk. For Inferact the prefix convention +matches its vLLM/TorchSpec lineage and measures better, but +0.9pt at n=200 is +inside this harness's noise band (it treats RadixArk's own 71-73% spread as +noise), so this is a direction, not a settled requirement -- unlike the 4.5pt +RadixArk case above, which was unambiguous. Both Inferact acc_len values (5.60 +and 5.66) sit on its model card's 5.64, so neither convention is grossly wrong +for it. + +Deriving this from checkpoint metadata is not possible today: neither published +drafter's config records which capture convention it was distilled against.""" + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + +# TRTLLM_K3_ATTN_RES_MXFP8=1: the pre-attention mix that folds in the previous MoE residual +# (k3_attn_res_moe_add_rmsnorm_mx) also stores the MXFP8 form of its output, and the +# attention's first projection (KDA fused QKVG, MLA fused q_a / kv_a) runs the low-M GEMM on it +# (k3_fp8w_lowm_mxfp8) instead of quantizing the BF16 rows itself: bitwise the same projection, +# generation-only batches where the projection runs the low-M kernel. +_K3_ATTN_RES_MXFP8 = _F.attn_res_mxfp8 + +# TRTLLM_K3_ULL_HYGIENE=1 (decode launch hygiene; the same bits): in CUDA-graph decode steps +# (multi-stream on) +# - the DSpark hidden-state taps (``maybe_capture_hidden_states``, a strided copy into the capture +# buffer) run on a side stream forked after their source and joined at the end of the model +# forward, instead of between the mix and the input norm of the tapped layer; +# - at an attention-residual block boundary served by the fused column-parallel / row-parallel tail +# mix (``_k3_attn_res_pdl_rp_tail``), the mix writes the updated prefix sum straight into the +# snapshot slot of the bank, so the snapshot copy (a memcpy graph node, which also cuts the PDL +# edge to the next layer's first GEMM) is gone. +_K3_ULL_HYGIENE = _F.ull_hygiene + + +class _SideCopies: + """``TRTLLM_K3_ULL_HYGIENE``: copies that nothing reads before the end of the forward run on a + side stream. Each is forked after its source; the sources stay referenced until + :meth:`join` makes the current stream wait for every pending copy, so the caching allocator + cannot hand their memory to a later kernel while a copy may still read it.""" + + def __init__(self, stream: torch.cuda.Stream): + self.stream = stream + self._pending: List[Tuple[torch.Tensor, torch.cuda.Event]] = [] + + @staticmethod + def active() -> bool: + """Side copies only where the model forks side streams (CUDA graph capture).""" + return do_multi_stream() + + def run(self, fn: Callable[[], None], source: torch.Tensor) -> None: + """Run ``fn`` (a copy reading ``source``) on the side stream.""" + fork = torch.cuda.Event() + fork.record() + done = torch.cuda.Event() + with torch.cuda.stream(self.stream): + self.stream.wait_event(fork) + fn() + done.record() + self._pending.append((source, done)) + + def join(self) -> None: + current = torch.cuda.current_stream() + for _, done in self._pending: + current.wait_event(done) + self._pending.clear() + + +# --------------------------------------------------------------------------- +# Decoder layer. +# --------------------------------------------------------------------------- + + +class KimiK3DefaultDecoderLayer(nn.Module): + # TRTLLM_K3_ULL_HYGIENE: the model's side-stream copies (``None``: off). + _side_copies: Optional[_SideCopies] = None + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + dtype = torch.bfloat16 + self.is_moe = ( + cfg.num_experts is not None + and layer_idx >= cfg.first_k_dense_replace + and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 + ) + + self.is_kda = _is_kda_layer(cfg, layer_idx) + is_mla = _is_mla_layer(cfg, layer_idx) + if self.is_kda == is_mla: + raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") + + if self.is_kda: + projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in projection_names + } + attention_config._frozen = model_config._frozen + self.linear_attn = K3KDALinearAttention( + cfg, + layer_idx, + mapping=model_config.mapping.mapping_for_full_rows(), + allreduce_strategy=model_config.allreduce_strategy, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + model_config=attention_config, + ) + else: + self.self_attn = K3MLARuntime( + cfg, + layer_idx, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + # CP original stashed by _setup_helix_mappings; None outside helix. + mapping_with_cp=getattr(model_config, "_helix_mapping_with_cp", None), + ) + + if self.is_moe: + self.block_sparse_moe = K3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) + else: + # The dense MLP (layer 0): BF16 gate_up GEMM, SiTU-and-mul, down GEMM and, on a TP + # group, the down projection's all-reduce. Attention DP replicates it (ranks own + # different tokens); otherwise it is TP over the largest group of consecutive ranks + # that divides the intermediate size and stays within one node. + mlp_mapping = model_config.mapping.mapping_for_full_rows() + if mlp_mapping.enable_attention_dp: + self.mlp_tp_size = 1 + else: + self.mlp_tp_size = math.gcd(cfg.intermediate_size, mlp_mapping.tp_size) + if self.mlp_tp_size > mlp_mapping.gpus_per_node: + self.mlp_tp_size = math.gcd(self.mlp_tp_size, mlp_mapping.gpus_per_node) + self.mlp = GatedMlp( + cfg.hidden_size, + cfg.intermediate_size, + activation="situ", + beta=getattr(cfg, "activation_situ_beta", None) or 1.0, + linear_beta=getattr(cfg, "activation_situ_linear_beta", None), + dtype=dtype, + mapping=subgroup_mapping(mlp_mapping, self.mlp_tp_size), + reduce_output=self.mlp_tp_size > 1, + allreduce_strategy=model_config.allreduce_strategy, + ) + + # FlashInfer RMSNorm on the K3 weights: the input norm after the pre-attention mix and the + # post-attention norm after A_mid (fused into the attention-residual kernels where they + # serve, ``trtllm::flashinfer_rmsnorm`` otherwise). + self.input_layernorm = RMSNormWeight(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype) + self.post_attention_layernorm = RMSNormWeight( + cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype + ) + + # Attention residual scheme (always on for K3): the score norms and the [1, H] score + # projections the attention-residual kernels read (``.weight`` / ``.eps``). + self.attn_res_block_size = cfg.attn_res_block_size + assert self.attn_res_block_size is not None, ( + "Kimi K3 runtime expects attn_res_block_size to be set" + ) + self.self_attention_res_norm = RMSNormWeight(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype) + self.mlp_res_norm = RMSNormWeight(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype) + self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + + def _attn_takes_mxfp8(self, num_rows: int, attn_metadata: AttentionMetadata) -> bool: + """Whether the pre-attention mix should also emit its MXFP8 form for the attention mixer + (``TRTLLM_K3_ATTN_RES_MXFP8``; generation-only batches): the KDA fused QKVG projection + (``qkvg_takes_mxfp8``) or the MLA fused q_a / kv_a projection (``q_a_takes_mxfp8``).""" + if not _K3_ATTN_RES_MXFP8 or attn_metadata.num_contexts != 0: + return False + if self.is_kda: + return self.linear_attn.qkvg_takes_mxfp8(num_rows) + q_a_takes_mxfp8 = getattr(self.self_attn.mixer, "q_a_takes_mxfp8", None) + return q_a_takes_mxfp8 is not None and q_a_takes_mxfp8(num_rows) + + def _attn_mx_atom(self, num_rows: int) -> bool: + """Whether that MXFP8 form goes to the mid-M GEMM, i.e. the mix stores its scales in that + GEMM's layout (k3_attn_res ``mx_atom``).""" + if self.is_kda: + mx_atom = getattr(self.linear_attn, "qkvg_mx_atom", None) + else: + mx_atom = getattr(self.self_attn.mixer, "q_a_mx_atom", None) + return mx_atom is not None and bool(mx_atom(num_rows)) + + def _snapshot_slot( + self, block_residual: torch.Tensor, num_snapshots: int + ) -> Optional[torch.Tensor]: + """``TRTLLM_K3_ULL_HYGIENE``: on a block-boundary layer, the bank slot this layer's + snapshot goes to, for the fused pre-attention mix to store the prefix sum into (no copy).""" + if ( + not _K3_ULL_HYGIENE + or self.layer_idx % self.attn_res_block_size != 0 + or num_snapshots >= block_residual.shape[0] + ): + return None + return block_residual[num_snapshots] + + def _capture_hidden_states(self, capture: Tuple[Any, int], value: torch.Tensor) -> None: + """The DSpark tap of ``capture[1]``; on the model's side stream with + ``TRTLLM_K3_ULL_HYGIENE`` (joined at the end of the model forward, before the drafter).""" + side = self._side_copies + if side is None or not side.active(): + capture[0].maybe_capture_hidden_states(capture[1], value, None) + return + side.run(lambda: capture[0].maybe_capture_hidden_states(capture[1], value, None), value) + + def forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, + ) -> Tuple[torch.Tensor, int]: + """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). + + ``block_residual`` is a preallocated snapshot bank in kernel-native + ``[K_max, M, H]`` layout. Returns the running prefix sum and the + number of valid bank rows. The prefix sum (input and output) can be a + ``_DeferredMoEResidual`` (``TRTLLM_K3_ATTN_RES_PDL``). + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. + """ + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + valid_block_residual = block_residual[:num_snapshots] + normed = None + attn_input_mxfp8 = None + if isinstance(hidden_states, _DeferredRpTail): + # The previous layer's row-parallel MoE tail: this mix finishes its all-reduce where + # it can; the DSpark tap path materializes it first. + if capture is None and num_snapshots > 0: + rows = hidden_states.residual.shape[0] + mxfp8 = self._attn_takes_mxfp8(rows, attn_metadata) + fused = _k3_attn_res_pdl_rp_tail( + hidden_states, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + mxfp8=mxfp8, + # TRTLLM_K3_ULL_TAIL_SANDWICH: release the staging KDA QKVG GEMM at kernel + # entry, as the MoE-residual mix does. + entry_trigger=(_K3_ULL_TAIL_SANDWICH or _K3_ULL2_COLL_CP_TAIL) + and mxfp8 + and self.is_kda + and self.linear_attn.qkvg_stages_weights(rows), + prefix_out=self._snapshot_slot(block_residual, num_snapshots), + ) + if fused is not None: + hidden_states, normed, *attn_input_mxfp8 = fused + if normed is None: + hidden_states = hidden_states.materialize() + elif isinstance(hidden_states, _DeferredMoEResidual): + # The previous layer's MoE residual update, folded into this mix when possible. + if capture is None and num_snapshots > 0: + rows = hidden_states.residual.shape[0] + mxfp8 = self._attn_takes_mxfp8(rows, attn_metadata) + fused = _k3_attn_res_pdl_moe_add( + hidden_states, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + mxfp8=mxfp8, + # Release the staging KDA QKVG GEMM at kernel entry. + entry_trigger=mxfp8 + and self.is_kda + and self.linear_attn.qkvg_stages_weights(rows), + mx_atom=mxfp8 and self._attn_mx_atom(rows), + ) + if fused is not None: + hidden_states, normed, *attn_input_mxfp8 = fused + if normed is None: + hidden_states = hidden_states.materialize() + prefix_sum = hidden_states + + if capture is not None and not _AUX_ATTN_RES_STREAM_ENABLED: + # prefix_only wants the incoming running prefix, which is already in + # hand -- no pre-norm intermediate is needed, so the tap is a plain + # read and this layer keeps the fused attn_res+RMSNorm path like the + # other 88. ``prefix_sum`` already has the previous layer's MLP + # output folded in, i.e. the quantity vLLM builds as + # ``prefix_sum + pending_mlp_out`` on its deferred-add path. + self._capture_hidden_states(capture, prefix_sum) + capture = None + + if capture is not None: + # The mixture tap needs the PRE-norm value (SGLang's + # aggregate_stream is a "softmax mixture, no output norm"), which + # the fused attn-res + RMSNorm kernel does not expose. Keep the two + # steps split on captured layers only and fuse everywhere else. + mixed = ( + _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + if num_snapshots > 0 + else hidden_states + ) + self._capture_hidden_states(capture, mixed) + hidden_states = rms_norm(mixed, self.input_layernorm.weight, self.input_layernorm.eps) + elif normed is not None: + hidden_states = normed + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + ) + else: + hidden_states = rms_norm( + hidden_states, self.input_layernorm.weight, self.input_layernorm.eps + ) + + if self.layer_idx % self.attn_res_block_size == 0: + snapshot = block_residual[num_snapshots] + if prefix_sum.data_ptr() != snapshot.data_ptr(): + # (TRTLLM_K3_ULL_HYGIENE: the fused mix above may have stored it there already.) + snapshot.copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + + # When A_mid below adds the attention output to the prefix sum, the attention block may + # hand back its pushed all-reduce unfinished (a LamportPending) for A_mid to finish and + # read from the workspace. A_mid must do so before the next all-reduce on the workspace + # re-arms that buffer: the shared expert, whose down-projection all-reduce is the next + # one, forks after A_mid, and its pushed finish kernel can only become resident once its + # producer has passed its own griddepcontrol.wait (every grid before, A_mid included, has + # then completed). + attn = self.linear_attn if self.is_kda else self.self_attn + attn.ar_finish_deferred = prefix_sum is not None and _k3_attn_res_lamport_applies( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + if self.is_kda: + if attn_input_mxfp8: + self.linear_attn.next_input_mxfp8 = tuple(attn_input_mxfp8) + hidden_states = self.linear_attn(hidden_states, attn_metadata) + else: + if attn_input_mxfp8: + self.self_attn.mixer.next_input_mxfp8 = tuple(attn_input_mxfp8) + hidden_states = self.self_attn(hidden_states, attn_metadata) + attn.ar_finish_deferred = False + + if prefix_sum is None: + prefix_sum = hidden_states + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + # Where the MoE front's head is the next launch after A_mid, A_mid releases it at entry. + releases_early = ( + getattr(self.block_sparse_moe, "ull_front_releases_early", None) + if self.is_moe + else None + ) + front_early = releases_early is not None and releases_early(prefix_sum) + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + hidden_states, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + entry_trigger=front_early, + ) + if self.is_moe and _K3_FUSE_MOE_RESIDUAL: + prefix_sum = self.block_sparse_moe( + hidden_states, + all_rank_num_tokens, + residual=prefix_sum, + defer_residual=_defer_moe_residual(prefix_sum.shape[0]), + # The row-parallel tail serves generation-only batches. + rp_tail_ok=getattr(attn_metadata, "num_contexts", 1) == 0, + ) + return prefix_sum, num_snapshots + if self.is_moe: + hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens) + else: + hidden_states = self.mlp(hidden_states) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, num_snapshots + + def skip_forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. + + ``modeling_utils.skip_forward()`` only drops a module's weights when it + finds this attribute, so without it the layer-wise benchmarks would + allocate all 93 layers instead of the profiled slice. + """ + return hidden_states, block_residual + + +# --------------------------------------------------------------------------- +# Model. +# --------------------------------------------------------------------------- + + +# TRTLLM_K3_ULL4_MLA_PLAN_EARLY=1: in a generation-only step (multi-stream) the model writes the +# FP4 MLA step metadata (generation lengths + packed split plan, device page table) on a side +# stream at the step start; the first FP4 MLA layer joins it instead of launching those kernels on +# its critical path. The same kernels with the same inputs: bitwise. +_K3_ULL4_MLA_PLAN_EARLY = _F.ull4_mla_plan_early + + +class KimiK3DefaultModelingV2(DecoderModel): + # TRTLLM_K3_ULL_HYGIENE: side-stream copies of the decode step (``None``: off). + _side_copies: Optional[_SideCopies] = None + + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + cfg = _get_text_config(model_config.pretrained_config) + self._text_cfg = cfg + dtype = torch.bfloat16 + + # Attention and MoE phases are sequential, so their branch-overlap + # roles share one stream; MoE-internal overlap roles remain separate. + aux_stream_list = [torch.cuda.Stream(priority=_K3_SHARED_STREAM_PRIORITY)] + aux_stream_list += [torch.cuda.Stream() for _ in range(2)] + aux_stream_list.append(torch.cuda.Stream(priority=_K3_LATENT_STREAM_PRIORITY)) + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + } + + # The whole vocabulary on every rank: a row gather (``embedding_gather`` with the profile's + # ``embedding_gather``, ``F.embedding`` otherwise). + self.embed_tokens = EmbeddingTable(cfg.vocab_size, cfg.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + [ + KimiK3DefaultDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) + for layer_idx in range(cfg.num_hidden_layers) + ] + ) + # TRTLLM_K3_ULL_HYGIENE: one side stream for the DSpark taps of every layer. + self._side_copies = _SideCopies(torch.cuda.Stream()) if _K3_ULL_HYGIENE else None + for layer in self.layers: + layer._side_copies = self._side_copies + # TRTLLM_K3_ULL4_MLA_PLAN_EARLY: the side stream and events of the step-metadata prelaunch. + self._mla_plan_stream = torch.cuda.Stream() if _K3_ULL4_MLA_PLAN_EARLY else None + self._mla_plan_events = ( + (torch.cuda.Event(), torch.cuda.Event()) if _K3_ULL4_MLA_PLAN_EARLY else None + ) + # The final norm, fused into the output attention-residual mix where that kernel serves. + self.norm = RMSNormWeight(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype) + self.output_attn_res_norm = RMSNormWeight(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype) + self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.num_attn_res_snapshots = ( + cfg.num_hidden_layers + cfg.attn_res_block_size - 1 + ) // cfg.attn_res_block_size + + # Which convention the drafter tap is on is not recoverable from the + # served output -- a mismatch only lowers acceptance -- so state it once + # at construction rather than leaving it to be inferred from an AL. + logger.info_once( + "Kimi K3 aux hidden capture: mode=" + f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " + f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", + key="kimi_k3_aux_capture_mode", + ) + + def _prelaunch_mla_step_metadata(self, attn_metadata: AttentionMetadata) -> None: + """TRTLLM_K3_ULL4_MLA_PLAN_EARLY: fork the step's FP4 MLA metadata kernels to a side stream + here; the first FP4 MLA layer joins it (``step_metadata_ready``).""" + if not do_multi_stream(): + return + mixer = None + for layer in self.layers: + if layer.is_kda: + continue + candidate = getattr(layer.self_attn, "mixer", None) + mqa = getattr(candidate, "mqa", None) + if getattr(mqa, "uses_fp4_mla_attention", False) and hasattr( + candidate, "prelaunch_step_metadata" + ): + mixer = candidate + break + if mixer is None: + return + fork, done = self._mla_plan_events + fork.record() + with torch.cuda.stream(self._mla_plan_stream): + fork.wait() + mixer.prelaunch_step_metadata(attn_metadata) + done.record() + mixer.step_metadata_ready = done + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata=None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + weight = self.embed_tokens.weight + if _F.embedding_gather and input_ids.dim() == 1 and weight.is_contiguous(): + # The F.embedding rows, copied with 16-byte accesses. + inputs_embeds = gather_embedding(input_ids, weight) + else: + inputs_embeds = F.embedding(input_ids, weight) + hidden_states = inputs_embeds + + # The FP4 MLA decode plan of the previous step is stale. + fp4_mla_decode.begin_step(attn_metadata) + if self._mla_plan_stream is not None and attn_metadata.num_contexts == 0: + self._prelaunch_mla_step_metadata(attn_metadata) + + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = ( + getattr(spec_metadata, "_capture_layer_set", None) + if spec_metadata is not None + else None + ) + for i, layer in enumerate(self.layers): + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + capture = None + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + capture = (spec_metadata, self.layers[i - 1].layer_idx) + hidden_states, num_snapshots = layer( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + capture=capture, + ) + hidden_states = _materialize_prefix_sum(hidden_states) + + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) + + output = _apply_attn_res_and_rmsnorm( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + if self._side_copies is not None: + # TRTLLM_K3_ULL_HYGIENE: the DSpark taps land before anything after the model (the + # drafter) reads the capture buffer. + self._side_copies.join() + return output + + +class KimiK3DefaultForCausalLM( + KimiK3DefaultTextWeights, SpecDecOneEngineForCausalLM[KimiK3DefaultModelingV2, Any] +): + """The Kimi K3 text model (KDA + MLA + MoE decoder layers) of ``ModelingV2KimiK3Nvfp4Sm107Default``. + + Registered under no architecture: the target class builds it as its ``llm``. Weight loading + is the ``KimiK3DefaultTextWeights`` mixin's (``weights.py``). + + Speculative decoding: ``SpecDecOneEngineForCausalLM`` builds the drafter and the worker; + :meth:`_install_private_spec` replaces them with the target's DSpark drafter and worker + (``drafter.py``, ``dflash.py``).""" + + mamba_metadata_cls = K3MambaMetadata + + @property + def runtime_hooks(self) -> Optional[RuntimeHooks]: + """The base hooks plus the ones the target's profile turns on: the speculative decode + step's head in one kernel (``embedding_gather``, ``decode_prologue``) and the decode + step's small input uploads, the one-model sampler's Philox seeds and offsets included, + as one upload and one scatter at the graph replay (``pack_step_inputs``, + ``pack_spec_rng_uploads``).""" + features = _features() + changes = {} + if features.embedding_gather: + changes["decode_prologue"] = spec_decode_head + if features.pack_step_inputs: + changes["pack_step_inputs"] = True + changes["pack_spec_rng_uploads"] = True + hooks = super().runtime_hooks + if not changes: + return hooks + return replace(hooks or DEFAULT_RUNTIME_HOOKS, **changes) + + def __init__(self, model_config: ModelConfig): + mlp_fp8 = os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_MOE_MLP_ENV, "0") + if mlp_fp8 not in ("0", "1"): + raise ValueError( + f"{_KIMI_K3_FP8_WEIGHT_READ_MOE_MLP_ENV} must be 0 or 1; got {mlp_fp8!r}" + ) + self._fp8_weight_read_moe_mlp = mlp_fp8 == "1" + cfg = _get_text_config(model_config.pretrained_config) + assert model_config.mapping.pp_size == 1, "Kimi K3 does not support pipeline parallelism" + spec_config = getattr(model_config, "spec_config", None) + + # Helix: swap in the repurposed mapping; restored after super().__init__. + self._setup_helix_mappings(model_config, cfg, spec_config) + # Supported spec-dec modes: + # - SA (suffix automaton): one-engine in-forward drafting, no draft + # weights; the KDA/MLA verify paths below implement multi-token + # verification for it. + # - DFlash: external-drafter parallel drafting; the drafter is a + # separate dense checkpoint (K2.7-Code-DFlash schema) consumed by + # the generic DFlashForCausalLM wrapper, and the target only has + # to expose per-layer hidden states via maybe_capture_hidden_states + # (see KimiK3DefaultModelingV2.forward). + # - DSpark: the same external-drafter flow with the Markov and + # confidence heads enabled (RadixArk/Kimi-K3-DSpark and friends). + # The target side is identical -- the capture in + # KimiK3DefaultModelingV2.forward is unconditional -- so this gate is the + # only place the mode has to be admitted. + # Modes needing draft heads (MTP/Eagle) are blocked until a + # draft-head checkpoint exists. + assert ( + spec_config is None + or spec_config.spec_dec_mode.is_sa() + or spec_config.spec_dec_mode.is_dflash() + or spec_config.spec_dec_mode.is_dspark() + ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" + super().__init__( + KimiK3DefaultModelingV2(model_config), + model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + # The drafter and worker of main's shell are built (under the repurposed helix mapping, + # which is still active here); the package's own replace them at this one point. + self._install_private_spec(model_config) + + # Restore the CP original: executor-side helix bookkeeping keys off + # has_cp_helix() at runtime. + if self.mapping_with_cp is not None: + model_config._frozen = False + model_config.mapping = self.mapping_with_cp + model_config._frozen = True + + def _install_private_spec(self, model_config: ModelConfig) -> None: + """Swap in the target's DSpark drafter and worker (``drafter.py``, ``dflash.py``). + + Called once from ``__init__``, right after ``SpecDecOneEngineForCausalLM.__init__`` and + before the helix CP mapping is restored, i.e. with the same mapping main's drafter was + built under. At that point main's shell has set ``self.spec_config``, + ``self.draft_config``, ``self.draft_model`` (from the draft-builder registry, or None for a + separate draft engine), ``self.spec_worker`` (``mapping_for_full_rows()`` mapping, + ``set_target_model(self)`` done), ``self.use_separate_draft_kv_cache`` and ``self.epilogue``. + + A private implementation must replace ``draft_model`` / ``spec_worker`` consistently: + swap them in ``self.epilogue`` (main appends both there), call ``set_target_model(self)`` on the + new worker, keep ``self.lm_head`` a stock ``LMHead`` and ``self.logits_processor`` exactly a + ``LogitsProcessor`` (main's vocab-parallel selection checks the type), and keep the + contract main's ``DFlashSpecMetadata.prepare()`` has with the worker (the attributes it + reads and the slot methods it calls, listed in ``drafter.py``). + + One-engine DSpark with a standalone MLA drafter gets the package's drafter + (``drafter.py``) and worker (``drafter.py`` on ``dflash.py``) + through :func:`~.drafter.install_k3_dspark`, which keeps that contract; every other + configuration keeps main's. + """ + install_k3_dspark(self, model_config) + + def _setup_helix_mappings( + self, + model_config: ModelConfig, + cfg: "PretrainedConfig", + spec_config: Optional["DecodingBaseConfig"], + ) -> None: + """Validate helix preconditions and stage the dual-mapping swap. + + DeepseekV3 pattern: the MLA layers keep the CP original; everything + else is built against the repurposed mapping (CP ranks become TP + ranks). Sets ``mapping_with_cp`` (restored after construction) and + ``_repurposed_tp_mapping`` (load_weights shard selection); both stay + None outside helix. + """ + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + if not model_config.mapping.has_cp_helix(): + return + # The K3 helix targets run head-TP attention, the residual stream and the MoE on every + # row of the CP group; there is no token-sliced (attention-DP) variant. + if model_config.mapping.enable_attention_dp: + raise NotImplementedError( + "Kimi K3 helix runs without attention DP: KDA and MLA are head-TP over the CP " + "group, so there is no request-level DP to add." + ) + if spec_config is not None: + # Helix supports only the standalone DSpark drafter (verified on + # the V2 superblock ledger); reject everything else loudly rather + # than let an unsupported spec mode run silently wrong. + if not spec_config.spec_dec_mode.is_dspark(): + raise ValueError( + "Kimi K3 helix supports speculative decoding only with " + f"DSpark (standalone drafter); got " + f"{spec_config.decoding_type!r}." + ) + # The SpeculationGate acceptance-rate trip permanently disables + # speculation mid-flight while enable_spec_decode stays True; + # in-flight helix requests then fall into the plain generation + # loop whose position math (total_input_len_cp + + # py_decoding_iter - 1) is stale once any draft token was + # accepted -> silently wrong RoPE positions and KV slots. Reject + # the trip wires until that loop is helix-group aware. + if ( + spec_config.acceptance_rate_window_size is not None + or spec_config.acceptance_rate_threshold is not None + ): + raise ValueError( + "Kimi K3 helix does not support the speculation " + "acceptance-rate gate (acceptance_rate_window_size / " + "acceptance_rate_threshold): dynamically disabling " + "speculation mid-flight leaves helix requests on a " + "single-token position formula." + ) + # max_concurrency is the same trip wire by another name: the + # drafter re-evaluates should_use_spec_decode on every scheduling + # iteration and flips enable_spec_decode off as soon as the active + # batch exceeds the cap. In-flight helix requests then take the + # plain generation loop, whose position formula counts ITERATIONS + # (total_input_len_cp + py_decoding_iter - 1) rather than + # committed tokens, so it is stale by however many draft tokens + # were accepted -- a wrong RoPE position, and across a ledger page + # boundary a KV write to the wrong CP rank. Mirror the drafter's + # own "unset" test (Drafter.should_use_spec_decode returns True + # when max_concurrency is None) so an unset value is not rejected. + if spec_config.max_concurrency is not None: + raise ValueError( + "Kimi K3 helix does not support the speculation " + "concurrency cutoff (max_concurrency): disabling " + "speculation above the cap leaves in-flight helix " + "requests on a position formula that assumes one " + "committed token per iteration, which accepted draft " + "tokens break." + ) + # draft_len_schedule is the user-facing alternative to + # max_concurrency (llm_args rejects setting both) and reaches the + # same end by a route that does not go through + # should_use_spec_decode at all: py_executor turns speculation off + # directly once the schedule yields draft_len 0 for the active + # batch size. Guarding only max_concurrency would leave this door + # open. Skip the schedule that llm_args synthesized from + # max_concurrency, so a config that set only that field raises the + # message above naming the field the user actually wrote. + if ( + spec_config.draft_len_schedule is not None + and not spec_config._translated_from_max_concurrency + ): + raise ValueError( + "Kimi K3 helix does not support the dynamic draft-length " + "schedule (draft_len_schedule): a batch size past the " + "last entry drops the draft length to 0 and turns " + "speculation off mid-run, leaving in-flight helix " + "requests on a position formula that assumes one " + "committed token per iteration." + ) + cp = model_config.mapping.cp_size + repurposed_tp = model_config.mapping.tp_size * cp + if cfg.num_attention_heads % repurposed_tp != 0: + raise ValueError( + f"Kimi K3 helix requires tp_size*cp_size ({repurposed_tp}) " + f"to divide the MLA head count ({cfg.num_attention_heads})." + ) + kda_heads = cfg.linear_attn_config["num_heads"] + if kda_heads % repurposed_tp != 0: + raise ValueError( + f"Kimi K3 helix requires tp_size*cp_size ({repurposed_tp}) to " + f"divide the KDA head count ({kda_heads})." + ) + # MoE splits apply to the repurposed tp*cp group (helix + # moe_world_size = tp*cp); default EP-only. The Mapping constructor + # skips its product check when both sizes are 1, so validate here. + moe_ep = repurposed_tp + if model_config.mapping.moe_tp_ep_user_specified: + moe_tp = model_config.mapping.moe_tp_size + moe_ep = model_config.mapping.moe_ep_size + if moe_tp * moe_ep != repurposed_tp: + raise ValueError( + f"Kimi K3 helix: moe_tensor_parallel_size ({moe_tp}) x " + f"moe_expert_parallel_size ({moe_ep}) must equal " + f"tp_size*cp_size ({repurposed_tp}): MoE runs on the " + "repurposed tp*cp group." + ) + if cfg.num_experts and cfg.num_experts % moe_ep != 0: + raise ValueError( + f"Kimi K3 helix requires the MoE EP size ({moe_ep}) to " + f"divide the routed expert count ({cfg.num_experts}): each " + "EP rank of the repurposed tp*cp group holds whole experts." + ) + self.mapping_with_cp = copy.deepcopy(model_config.mapping) + repurposed = model_config.mapping.repurpose_helix_cp_to_tp() + # repurpose passes resolved moe sizes, which the Mapping constructor + # mistakes for user-specified values; restore the flag. + repurposed.moe_tp_ep_user_specified = self.mapping_with_cp.moe_tp_ep_user_specified + # load_weights shard selection must use this tp_rank; the restored + # CP original's tp_rank is 0 on every rank. + self._repurposed_tp_mapping = repurposed + model_config._frozen = False + model_config.mapping = repurposed + # Side-channel for the MLA layers (avoids threading a kwarg through + # every intermediate signature). + model_config._helix_mapping_with_cp = self.mapping_with_cp + model_config._frozen = True + + @classmethod + def get_model_defaults(cls, llm_args) -> dict: + # - Keep block reuse opt-in because its combination with SA + # speculative decoding is not validated. Both settings use V2 + # with the default use_kv_cache_manager_v2=True. The "auto" + # setting follows the model preference and can resolve to V1. + # - tokens_per_block=64: keep the validated K3 page size. Earlier + # releases needed it to avoid a legacy MLA fallback that could not + # serve attention-DP's 96 query heads. The current FlashInfer + # TRTLLM-Gen backend declines H=96 at either page size, allowing + # the updated legacy TRTLLM-Gen autotuner to select a Q16 kernel. + return { + "kv_cache_config": { + "enable_block_reuse": False, + "tokens_per_block": 64, + } + } + + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config: Any = None, + ) -> Literal["PYTHON"]: + """Kimi K3 disaggregated serving requires the Python transceiver. + + Only the Python NIXL transceiver (KvCacheTransceiverV2) can move + the KDA recurrent state; the C++ transceiver has no KDA support. + Adopted when the user leaves + ``cache_transceiver_config.transceiver_runtime`` at 'auto' and the + effective backend is NIXL. An explicit non-Python runtime is + rejected by ``get_kv_cache_manager_cls`` rather than silently + routed to a path that cannot transfer the recurrent state. + """ + return "PYTHON" + + +class ModelingV2KimiK3Nvfp4Sm107Default(KimiK3DefaultWeights, PreTrainedModel): + """Kimi K3 NVFP4 on SM107, sharded as the engine's mapping says. + + The vision-language model: the vision tower (``_VISION_MODEL_CLS``, built in + :meth:`load_weights` when meta init deferred it) and the text model + ``KimiK3DefaultForCausalLM`` as ``llm``. The speculative-decoding and weight-loading + attributes the engine reads off the outer model (``draft_model``, ``lm_head``, + ``runtime_hooks``, ...) are the text model's. Weight loading is the + ``KimiK3DefaultWeights`` mixin's (``weights.py``). + """ + + _LANG_PREFIX = "language_model." + _VISION_MODEL_CLS = KimiK3VisionModel + mamba_metadata_cls = KimiK3DefaultForCausalLM.mamba_metadata_cls + + @classmethod + def get_preferred_kv_cache_manager_version( + cls, + pretrained_config: Any = None, + ) -> Literal["V2"]: + """Prefer KV cache manager V2 (as Kimi K2.5).""" + return "V2" + + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config: Any = None, + ) -> Literal["PYTHON"]: + """Default to the Python (v2) KV-cache transceiver (as Kimi K2.5). + + Only adopted when the user leaves + ``cache_transceiver_config.transceiver_runtime`` at 'auto' and the + effective backend is NIXL; otherwise the C++ transceiver is used. + """ + return "PYTHON" + + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + *args, + **kwargs, + ) -> None: + config = model_config.pretrained_config + # PreTrainedModel.__init__ resolves config._attn_implementation and + # rejects the default "sdpa" unless the class declares support. This + # wrapper never runs HF attention itself (TRT-LLM attention is wired + # inside the sub-modules), so declare support to keep HF init happy. + self._supports_sdpa = True + PreTrainedModel.__init__(self, config) + + # Re-entry guard: the tail of this __init__ repoints + # model_config.pretrained_config at the text_config and remaps the quant + # exclude list, so running the body twice on the same instance would + # fail the text_config assert below. If the text backbone already + # exists this is a re-init on a built instance: no-op. + if hasattr(self, "llm"): + return + + self.model_config = model_config + self._vlm_pretrained_config = config + + # --- Vision encoder (deferred under MetaInitMode, recreated in + # load_weights) --- + self.mm_encoder = None + if not DISAGG: + try: + mm_encoder = self._VISION_MODEL_CLS(model_config) + if _has_meta_tensors(mm_encoder): + logger.info("Vision encoder deferred to load_weights() (MetaInitMode active)") + else: + self.mm_encoder = mm_encoder + except MetaInitException: + logger.info("Vision encoder deferred to load_weights() (MetaInitMode active)") + + text_model_config = copy.copy(model_config) + assert hasattr(config, "text_config"), "Kimi K3 config must have text_config" + text_model_config._frozen = False + text_model_config.pretrained_config = config.text_config + + # Remap quant exclude_modules: language_model.X -> model.X + if text_model_config.quant_config.exclude_modules: + text_model_config.quant_config = copy.copy(text_model_config.quant_config) + p = self._LANG_PREFIX + mapped = [] + for m in text_model_config.quant_config.exclude_modules: + if m.startswith(p): + rest = m[len(p) :] + if rest.startswith("layers."): + rest = "model." + rest + mapped.append(rest) + else: + mapped.append(m) + text_model_config.quant_config.exclude_modules = mapped + + if not text_model_config.skip_create_weights_in_init: + text_model_config.skip_create_weights_in_init = True + text_model_config._frozen = True + + self.llm = KimiK3DefaultForCausalLM(text_model_config) + + self._media_placeholder_token_id = getattr( + config, "media_placeholder_token_id", _MEDIA_PLACEHOLDER_TOKEN_ID + ) + # Backing storage for the ``mm_token_ids`` property below — built once + # here so the executor doesn't re-allocate per step. + self._mm_token_ids = torch.tensor([self._media_placeholder_token_id], dtype=torch.int32) + + # Point model_config at the text_config so the executor reads generation + # params (eos_token_id, ...) from the text backbone, matching K2.5. + self.config = self.llm.config + model_config._frozen = False + model_config.pretrained_config = self.llm.config + model_config._frozen = True + + # Forward spec-dec / weight-loading attrs to self.llm: this wrapper is not + # itself a spec-dec model but is the outer model that those paths see. + @property + def draft_config(self): + return self.llm.draft_config + + @property + def draft_model(self): + return self.llm.draft_model + + @property + def model(self): + return self.llm.model + + @property + def lm_head(self): + return self.llm.lm_head + + @property + def vocab_size_padded(self) -> int: + return self.llm.vocab_size_padded + + @property + def runtime_hooks(self) -> Optional[RuntimeHooks]: + """The text model's step-loop hooks: the engine reads them off the outer model.""" + return self.llm.runtime_hooks + + def set_guided_decoder(self, *args, **kwargs): + return self.llm.set_guided_decoder(*args, **kwargs) + + @property + def multimodal_data_device_paths(self) -> List[str]: + return [ + "image.pixel_values", + "image.image_grid_thw", + "video.pixel_values_videos", + "video.video_grid_thw", + "multimodal_embedding", + ] + + @property + def mm_token_ids(self) -> torch.Tensor: + """Surface the in-vocab media placeholder to the model engine so + ``runners.prepare_multimodal_indices`` selects the ``torch.isin`` predicate + instead of the OOV (``>= vocab_size``) fallback (which would miss + Kimi's placeholder and force ``fuse_input_embeds`` through the + ``torch.where`` host-sync path on GPU input_ids). + """ + return self._mm_token_ids + + def infer_max_seq_len(self) -> int: + return self.llm.infer_max_seq_len() + + @torch.inference_mode() + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: Optional[bool] = False, + multimodal_params: Optional[List[MultimodalParams]] = None, + **kwargs, + ) -> torch.Tensor: + multimodal_params = multimodal_params or [] + mm_embeds: List[torch.Tensor] = [] + mm_token_ids = None + + if len(multimodal_params) > 0: + if DISAGG: + raise NotImplementedError("Disaggregated inference not yet supported for K2.5.") + mm_ctx_params = multimodal_params[: attn_metadata.num_contexts] + mm_embeds = get_multimodal_embeddings( + encoder_forward_fn=self.mm_encoder.forward, + multimodal_params=mm_ctx_params, + ) + mm_embeds = find_input_mm_embeds(mm_embeds, mm_ctx_params) + + # ``runners.prepare_multimodal_indices`` sees Kimi's in-vocab + # placeholder via ``self.mm_token_ids`` and emits indices that match + # ``find_input_mm_embeds``'s active-chunk slice; a chunked-prefill / + # KV-reuse mismatch surfaces as a row-count error inside + # ``fuse_input_embeds`` itself. + if len(mm_embeds) > 0: + mm_token_ids = self.mm_token_ids.to(input_ids.device, dtype=input_ids.dtype) + input_ids, inputs_embeds = fuse_input_embeds( + self.llm.model.embed_tokens, + input_ids, + mm_embeds, + mm_token_ids=mm_token_ids, + mm_token_indices=kwargs.get("mm_token_indices"), + text_token_indices=kwargs.get("text_token_indices"), + ) + + return self.llm.forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + **kwargs, + ) + + +__all__ = ["ModelingV2KimiK3Nvfp4Sm107Default"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/moe.py new file mode 100644 index 000000000000..d2bc2090ab8c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/moe.py @@ -0,0 +1,979 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi K3's routed experts of one MoE layer on this rank, as a chain of kernel launches. + +:class:`K3RoutedExperts` runs one forward on its NVFP4 weights (:class:`K3Nvfp4Experts`, the +``backend`` attribute DWDP binds) with one of the CuTe DSL kernel sets: + +``CUTEDSL`` (context):: + + routing (route tail | noaux_tc_route_sort | noaux_tc_op) + fp4_quantize [on the aux stream up to 1024 rows] + moe_sort [unless the routing wrote the tables] + moe_output_memset_inplace [aux stream] + cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_{rubin,blackwell} FC1 + SiTU + FP4 requant + cute_dsl_nvfp4_grouped_gemm_finalize_inplace_{rubin,blackwell} FC2 + weighted scatter-add + +``CUTEDSL_FC12`` (generation, SM107):: + + routing, fp4_quantize, moe_sort (as above) + cute_dsl_nvfp4_fc12_fused_rubin memset + FC1 + SiTU + FC2 + scatter-add + k3r_nvfp4_small_batch_moe instead, up to 32 rows (K3SmallBatchExperts) + +The routing is the caller's (:class:`SortedRouting`, or a ``route`` callable run per chunk): the +experts only run the kernels. The output is this rank's partial sum (its experts, its slice of the +intermediate dim); the caller reduces it. + +Kimi K3's NVFP4 routed-expert weights on one rank, in the layout the CuTe DSL MoE kernels read. + +:class:`K3Nvfp4Experts` holds one MoE layer's experts ``[slot_start, slot_end)`` (all of them under +DWDP once the setup widened the range) as the kernels take them: + +* ``w3_w1_weight`` ``[E, N1, H / 16]`` int64 (16 E2M1 values each) and ``w3_w1_weight_scale`` + ``[E, N1, H / 64]`` int32 (4 E4M3 block scales each, 128x4-swizzled), ``N1 = align128(2 I)``; + FC1 rows in the gate/up order of the kernels of this GPU (:func:`fc1_rows_for_kernels`). +* ``w2_weight`` ``[E, H, N1 / 32]`` int64 and ``w2_weight_scale`` ``[E, H, N1 / 128]`` int32 + (swizzled). +* ``fc31_input_scale`` / ``fc2_input_scale``: the reciprocal of the largest checkpoint + ``input_scale`` of the rank's experts (FC1 input, FC2 input); ``fc31_alpha`` / ``fc2_alpha`` + ``[E]``: the per-expert dequantization scales of the two GEMMs. + +The checkpoint streams in expert by expert (:meth:`load_expert`); :meth:`finalize` runs once the +last of the rank's experts is in. The parameter names are the ones DWDP rebinds and exchanges. + +Kimi K3's routed experts at small token counts: the small-batch NVFP4 kernel +(``ops.moe.k3_small_batch_moe``) on the expert weights of a ``CUTEDSL_FC12`` layer. + +:class:`K3SmallBatchExperts` serves the forwards of 1 to ``SMALL_BATCH_MOE_MAX_TOKENS`` tokens of a +:class:`K3RoutedExperts` that runs the FC12 kernels; the experts route and quantize the input, this +runs the kernel on their weights. Besides this rank's ``moe_sort`` tables the kernel takes a top-k +a routing kernel published token by token (:class:`PublishedSortTables`) and pushes its output into a +one-shot MNNVL all-reduce (:class:`LamportWorkspaceView`). + +What the Kimi K3 MoE routing hands to the routed experts of the same forward. + +``SortTables``: the six ``moe_sort`` outputs of this rank's experts at one tile size. +``SortedRouting``: the routing a kernel computed (ids, weights) together with this rank's +``SortTables``, or ``PublishedSortTables`` where the routing kernel published the routed ids token by +token and the experts build the tables themselves. +""" + +import os +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.autotuner import AutoTuner +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._torch.utils import ActivationType, Fp4QuantizedTensor +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +from .kernels.k3_small_batch_moe_op import ( + SMALL_BATCH_MOE_MAX_TOKENS, + new_small_batch_moe_counters, + small_batch_moe_route_word, + small_batch_moe_supports, + small_batch_moe_takes_published, +) +from .layers import load_weight_shard +from .ops.moe.k3_small_batch_moe import k3_small_batch_moe + + +@dataclass(frozen=True) +class SortTables: + """The outputs of ``trtllm::moe_sort`` (or of a routing kernel that writes the same tables) for + this rank's experts at ``tile_size`` rows per tile: ``(tile_idx_to_group_idx, + tile_idx_to_mn_limit, expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, + total_num_padded_tokens, num_non_exiting_tiles)``.""" + + tile_size: int + outputs: Tuple[torch.Tensor, ...] + + +@dataclass(frozen=True) +class PublishedSortTables: + """In place of :class:`SortTables`: the routing kernel published the routed ids token by + token into the experts' published-routing word, and the experts build this rank's tables + themselves.""" + + +@dataclass(frozen=True) +class SortedRouting: + """The routing of one forward: the top-k ids and weights, with this rank's ``moe_sort`` tables + where the routing kernel wrote them (``None``: the experts run ``moe_sort``).""" + + token_selected_experts: torch.Tensor # int32 [num_tokens, top_k] + token_final_scales: torch.Tensor # fp32 [num_tokens, top_k] + sort_tables: Optional[Union[SortTables, PublishedSortTables]] + + +#: Values per NVFP4 block scale. +SF_VEC_SIZE = 16 +#: FC1 rows (gate + up) are padded to a multiple of this. +_ROW_ALIGNMENT = 128 + + +def ep_partition(num_experts: int, ep_size: int, ep_rank: int) -> Tuple[int, int]: + """``[start, end)`` of ``ep_rank``'s experts: ranks below ``num_experts % ep_size`` hold one + expert more.""" + if num_experts < ep_size: + raise ValueError(f"{num_experts} experts cannot be split over {ep_size} ranks") + base, remainder = divmod(num_experts, ep_size) + start = ep_rank * base + min(ep_rank, remainder) + return start, start + base + (1 if ep_rank < remainder else 0) + + +def fc1_rows_for_kernels(x: torch.Tensor) -> torch.Tensor: + """FC1 rows (dim 0: the up half, then the gate half) in the order the CuTe DSL FC1 kernels of + this GPU read: Rubin (SM107) groups of 16 gate rows then 16 up rows, Blackwell groups of 64 + up rows then 64 gate rows. Weights and their block scales take the same order.""" + rows = x.shape[0] + group, gate_first = (16, True) if get_sm_version() == 107 else (64, False) + assert rows % (2 * group) == 0, (rows, group) + halves = x.view(2, rows // (2 * group), group, *x.shape[1:]) + if gate_first: + halves = halves.flip(0) + return halves.transpose(0, 1).contiguous().view(x.shape) + + +def _pad_to(src: torch.Tensor, dst: torch.Tensor) -> torch.Tensor: + """``src`` zero-padded at the end of both dims to ``dst``'s 2D shape.""" + rows, cols = src.shape + if (rows, cols) != tuple(dst.shape): + src = torch.nn.functional.pad( + src, (0, dst.shape[1] - cols, 0, dst.shape[0] - rows), "constant", 0 + ).contiguous() + return src + + +def _swizzle_block_scales(dst: torch.Tensor) -> None: + """Swizzle one expert's linear E4M3 block scales (int32-packed, in place) to the 128x4 layout.""" + swizzled = ( + torch.ops.trtllm.block_scale_interleave(dst.view(torch.uint8)) + .view(dst.dtype) + .reshape(dst.shape) + ) + torch.cuda.synchronize() + dst.copy_(swizzled) + + +class K3Nvfp4Experts(nn.Module): + """One MoE layer's NVFP4 routed experts on this rank (see the module docstring). + + ``tp_size`` / ``tp_rank``: MoE tensor parallelism, slicing every expert's intermediate dim (FC1 + rows, FC2 columns); ``ep_size`` / ``ep_rank``: which experts the rank holds. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + tp_size: int, + tp_rank: int, + ep_size: int, + ep_rank: int, + ) -> None: + super().__init__() + if hidden_size % 4 != 0: + raise ValueError(f"hidden_size {hidden_size} must be divisible by 4") + self.num_experts = num_experts + self.hidden_size = hidden_size + self.tp_size, self.tp_rank = tp_size, tp_rank + self.ep_size, self.ep_rank = ep_size, ep_rank + self.intermediate_size_per_partition = intermediate_size // tp_size + self.scaling_vector_size = SF_VEC_SIZE + # DWDP's fixup re-materializes activation constants only on modules that declare one. + self.activation = None + self.set_expert_range(*ep_partition(num_experts, ep_size, ep_rank)) + self._lock = threading.Lock() + self._staged: Optional[Dict] = None + self._streamed_expert_slots: set = set() + + def set_expert_range(self, start: int, end: int) -> None: + """Hold experts ``[start, end)`` (before the weights are created).""" + self.num_slots = self.num_experts + self.expert_size_per_partition = end - start + self.slot_start, self.slot_end = start, end + self.initial_local_expert_ids = list(range(start, end)) + self.initial_global_assignments = list(range(self.num_experts)) + + def create_weights(self) -> None: + num_local = self.expert_size_per_partition + n1 = -(-2 * self.intermediate_size_per_partition // _ROW_ALIGNMENT) * _ROW_ALIGNMENT + h = self.hidden_size + packed = 16 # E2M1 values per int64 + scales = 4 # E4M3 block scales per int32 + + def param(shape, dtype, fill=None): + data = ( + torch.empty(shape, dtype=dtype) + if fill is None + else torch.full(shape, fill, dtype=dtype) + ) + return nn.Parameter(data, requires_grad=False) + + self.w3_w1_weight = param((num_local, n1, h // packed), torch.int64) + self.w2_weight = param((num_local, h, n1 // 2 // packed), torch.int64) + self.w3_w1_weight_scale = param( + (num_local, n1, h // SF_VEC_SIZE // scales), torch.int32, fill=1 + ) + self.w2_weight_scale = param( + (num_local, h, n1 // 2 // SF_VEC_SIZE // scales), torch.int32, fill=1 + ) + self.fc31_input_scale = param((), torch.float32, fill=1.0) + self.fc2_input_scale = param((), torch.float32, fill=1.0) + self.fc31_alpha = param((num_local,), torch.float32, fill=1.0) + self.fc2_alpha = param((num_local,), torch.float32, fill=1.0) + + # ------------------------------------------------------------------ + # Loading: expert by expert, then one finalize per layer. + # ------------------------------------------------------------------ + def prepare_streaming_load(self) -> None: + """Create the per-layer staging the loader threads fill (before any of them runs).""" + self._staged = {"input_scales": {}, "weight_scale_2": {}} + self._streamed_expert_slots = set() + + def load_expert( + self, + *, + global_expert_id: int, + local_slot_id: int, + w1_weight: torch.Tensor, + w1_weight_scale: torch.Tensor, + w1_weight_scale_2: torch.Tensor, + w1_input_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + w2_weight_scale_2: torch.Tensor, + w2_input_scale: torch.Tensor, + w3_weight: torch.Tensor, + w3_weight_scale: torch.Tensor, + w3_weight_scale_2: torch.Tensor, + w3_input_scale: torch.Tensor, + ) -> None: + """Load one checkpoint expert (packed E2M1 ``uint8`` weights, E4M3 block scales, FP32 + global scales) into ``local_slot_id``: this rank's slice of the intermediate dim, FC1 as + ``[w3 | w1]`` padded to ``N1`` rows, block scales swizzled.""" + if not 0 <= local_slot_id < self.expert_size_per_partition: + raise IndexError( + f"local_slot_id={local_slot_id} is outside [0, {self.expert_size_per_partition})" + ) + if global_expert_id != self.initial_local_expert_ids[local_slot_id]: + raise ValueError( + f"local slot {local_slot_id} holds expert " + f"{self.initial_local_expert_ids[local_slot_id]}, got {global_expert_id}" + ) + if self._staged is None: + raise RuntimeError("prepare_streaming_load() must run before load_expert()") + if local_slot_id in self._streamed_expert_slots: + raise ValueError(f"NVFP4 local slot {local_slot_id} was loaded twice") + for name, value in (("w1", w1_weight), ("w2", w2_weight), ("w3", w3_weight)): + if value.dtype != torch.uint8: + raise TypeError(f"{name}_weight must hold packed NVFP4 uint8, got {value.dtype}") + + tp, rank = self.tp_size, self.tp_rank + column, row = TensorParallelMode.COLUMN, TensorParallelMode.ROW + + dst = self.w3_w1_weight.data[local_slot_id] + w1 = load_weight_shard(w1_weight, tp, rank, column, device=dst.device) + w3 = load_weight_shard(w3_weight, tp, rank, column, device=dst.device) + fc1 = torch.cat([w3.contiguous().view(dst.dtype), w1.contiguous().view(dst.dtype)], dim=0) + dst.copy_(_pad_to(fc1, dst), non_blocking=True) + + dst = self.w2_weight.data[local_slot_id] + w2 = load_weight_shard(w2_weight, tp, rank, row, device=dst.device) + dst.copy_(_pad_to(w2.contiguous().view(dst.dtype), dst), non_blocking=True) + + dst = self.w3_w1_weight_scale.data[local_slot_id] + w1_s = load_weight_shard(w1_weight_scale, tp, rank, column, device=dst.device) + w3_s = load_weight_shard(w3_weight_scale, tp, rank, column, device=dst.device) + fc1_s = torch.cat( + [w3_s.contiguous().view(dst.dtype), w1_s.contiguous().view(dst.dtype)], dim=0 + ) + dst.copy_(_pad_to(fc1_s, dst)) + _swizzle_block_scales(dst) + + dst = self.w2_weight_scale.data[local_slot_id] + w2_s = load_weight_shard(w2_weight_scale, tp, rank, row, device=dst.device) + cols = dst.shape[1] * 4 + if cols > w2_s.shape[1]: + w2_s = torch.nn.functional.pad( + w2_s, (0, cols - w2_s.shape[1]), "constant", 0 + ).contiguous() + dst.copy_(_pad_to(w2_s.view(dst.dtype), dst)) + + self._staged["weight_scale_2"][local_slot_id] = { + "w1": w1_weight_scale_2, + "w3": w3_weight_scale_2, + "w2": w2_weight_scale_2, + } + self._staged["input_scales"][global_expert_id] = { + "w1": w1_input_scale[...].reshape([]), + "w3": w3_input_scale[...].reshape([]), + "w2": w2_input_scale[...].reshape([]), + } + self._streamed_expert_slots.add(local_slot_id) + + def loaded_slots(self) -> set: + return self._streamed_expert_slots + + def finalize(self) -> None: + """After the last expert of the layer: swizzle the FC2 block scales, compute the global + input scales and the per-expert alphas, and put FC1 rows and block scales in the kernels' + gate/up order.""" + staged, self._staged = self._staged, None + if staged is None: + return + for slot in range(self.w2_weight_scale.shape[0]): + _swizzle_block_scales(self.w2_weight_scale.data[slot]) + + fc31_values, fc2_values = [], [] + for expert_id, entry in staged["input_scales"].items(): + assert torch.allclose(entry["w1"], entry["w3"]), ( + f"expert {expert_id}: w1 input_scale != w3 input_scale" + ) + fc31_values.append(entry["w1"]) + fc2_values.append(entry["w2"]) + if fc31_values: + self.fc31_input_scale.data.copy_(torch.stack(fc31_values).max().reciprocal()) + if fc2_values: + self.fc2_input_scale.data.copy_(torch.stack(fc2_values).max().reciprocal()) + + fc31_input_scale = self.fc31_input_scale.data + fc2_input_scale = self.fc2_input_scale.data + for slot, scales in staged["weight_scale_2"].items(): + w1_ws2, w3_ws2 = scales["w1"], scales["w3"] + if not torch.allclose(w1_ws2, w3_ws2): + logger.warning( + f"w1_weight_scale_2 != w3_weight_scale_2 ({w1_ws2} != {w3_ws2}), " + "selecting the larger value. Accuracy may be affected." + ) + w1_ws2 = torch.max(w1_ws2, w3_ws2) + gate_ws2 = 1.0 / w1_ws2[...].reshape([]) + self.fc31_alpha.data[slot].copy_(1.0 / (fc31_input_scale * gate_ws2)) + w2_ws2 = 1.0 / scales["w2"][...].reshape([]) + self.fc2_alpha.data[slot].copy_(1.0 / (fc2_input_scale * w2_ws2)) + + n = self.intermediate_size_per_partition * 2 + k = self.hidden_size + sf_cols = k // SF_VEC_SIZE + for slot in range(self.w3_w1_weight.shape[0]): + weight = self.w3_w1_weight.data[slot] + weight.copy_(fc1_rows_for_kernels(weight.view(torch.uint8)).view(weight.dtype)) + scale = self.w3_w1_weight_scale.data[slot] + linear = torch.ops.trtllm.block_scale_interleave_reverse( + scale.view(torch.uint8).view(-1, n, sf_cols) + ).view(n, sf_cols) + swizzled = torch.ops.trtllm.block_scale_interleave( + fc1_rows_for_kernels(linear).view(-1, n, sf_cols) + ).view(n, sf_cols) + scale.copy_(swizzled.view(scale.dtype)) + + +if TYPE_CHECKING: + from .layers import LamportWorkspaceView + + +class K3SmallBatchExperts: + """The small-batch path of one MoE layer's routed experts (``experts``).""" + + def __init__(self, experts: "K3RoutedExperts") -> None: + self.experts = experts + weights = experts.backend + self.max_tokens = 0 + if experts.kernels == "CUTEDSL_FC12" and small_batch_moe_supports( + weights.intermediate_size_per_partition, weights.hidden_size + ): + self.max_tokens = SMALL_BATCH_MOE_MAX_TOKENS + self._takes_published = bool(self.max_tokens) and small_batch_moe_takes_published( + weights.expert_size_per_partition + ) + # Created after loading (or by the first launch, before any graph capture); every launch + # leaves it zero. + self.counters: Optional[torch.Tensor] = None + + def post_load_weights(self) -> None: + if self.max_tokens and self.counters is None: + self.counters = new_small_batch_moe_counters(self.experts.backend.w3_w1_weight.device) + + def applies(self, num_tokens: int) -> bool: + """Whether a forward of ``num_tokens`` tokens runs here.""" + return 1 <= num_tokens <= self.max_tokens + + def published_routing_word(self, num_tokens: int) -> Optional[torch.Tensor]: + """The kernel's routing word where it serves ``num_tokens`` tokens in the token-indexed + layout (MoE-TP slices): a routing kernel publishes the top-k there + (:class:`PublishedSortTables`) and the kernel builds its tables from it.""" + if ( + not self.applies(num_tokens) + or not self._takes_published + or self.experts.backend.tp_size <= 1 + or self.counters is None + ): + return None + return small_batch_moe_route_word(self.counters) + + def run( + self, + x: torch.Tensor, + x_sf: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + sort_tables: Optional[Union[SortTables, PublishedSortTables]], + push: Optional["LamportWorkspaceView"], + ) -> torch.Tensor: + """The kernel on the experts' weights; writes every element of the output (and pushes it + into ``push``). ``sort_tables``: this rank's tables at tile 128 (else ``moe_sort`` runs + here), or the published top-k.""" + experts = self.experts + weights = experts.backend + assert token_final_scales.dtype == torch.float32, ( + f"the small-batch kernel requires float32 routing scales, got {token_final_scales.dtype}" + ) + published = isinstance(sort_tables, PublishedSortTables) + if published and weights.tp_size <= 1: + raise RuntimeError("published routing needs the token-indexed layout (MoE-TP)") + if published: + tile_idx_to_expert_idx = tile_idx_to_mn_limit = permuted_idx_to_expanded_idx = None + num_non_exiting_tiles = None + expanded_idx_to_permuted_idx = token_selected_experts.reshape(-1) + else: + if sort_tables is not None and sort_tables.tile_size == 128: + sort_outputs = sort_tables.outputs + else: + sort_outputs = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=weights.num_slots, + top_k=token_selected_experts.size(1), + local_expert_offset=weights.slot_start, + local_num_experts=weights.expert_size_per_partition, + tile_tokens_dim=128, + ) + ( + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + _total_num_padded_tokens, + num_non_exiting_tiles, + ) = sort_outputs + if self.counters is None: + self.counters = new_small_batch_moe_counters(x.device) + out = torch.empty( + (token_final_scales.size(0), weights.hidden_size), dtype=torch.bfloat16, device=x.device + ) + k3_small_batch_moe( + input=x.view(torch.float4_e2m1fn_x2), + input_scale=x_sf.view(torch.uint8), + fc1_weight=weights.w3_w1_weight.view(torch.float4_e2m1fn_x2), + fc1_weight_scale=weights.w3_w1_weight_scale.view(torch.uint8), + fc1_alpha=weights.fc31_alpha, + fc2_weight=weights.w2_weight.view(torch.float4_e2m1fn_x2), + fc2_weight_scale=weights.w2_weight_scale.view(torch.uint8), + fc2_alpha=weights.fc2_alpha, + output=out, + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + global_sf=weights.fc2_input_scale, + token_final_scales=token_final_scales, + counters=self.counters, + top_k=token_selected_experts.size(1), + situ_beta=experts.situ_beta, + situ_linear_beta=experts.situ_linear_beta, + token_indexed=weights.tp_size > 1, + die_major=False, + published_routing=published, + expert_offset=weights.slot_start, + push_mc=push.mc_base if push is not None else 0, + push_flags=push.flags if push is not None else None, + push_rank=push.rank if push is not None else 0, + push_world=push.world if push is not None else 0, + sm_reserve=experts.sm_reserve, + ) + return out + + +#: The kernel sets :class:`K3RoutedExperts` runs, by ``moe_config.backend``. +KERNEL_SETS = ("CUTEDSL", "CUTEDSL_FC12") + +# ``TLLM_MOE_ASYNC_INPUT_QUANT=0`` keeps the CUTEDSL input quantization on the main stream. +_ASYNC_INPUT_QUANT = os.environ.get("TLLM_MOE_ASYNC_INPUT_QUANT", "1") == "1" +_ASYNC_INPUT_QUANT_MAX_TOKENS = 1024 + +#: A route: ``route(router_logits, local_expert_offset, local_num_experts)``. +Route = Callable[[torch.Tensor, int, int], SortedRouting] + + +def _wait(event: Optional[torch.cuda.Event]) -> None: + if event is not None: + event.wait() + + +def _tile_runners(): + """The AutoTuner adapters of the two kernel sets: the routing tile is the tuned tactic.""" + from tensorrt_llm._torch.moe.fused_moe.cutedsl.fused_fc12 import CuteDslFc12FusedMoENvfp4Runner + from tensorrt_llm._torch.moe.fused_moe.cutedsl.runners import CuteDslFusedMoENvfp4Runner + + return { + "CUTEDSL": (CuteDslFusedMoENvfp4Runner, "CuteDslFusedMoE::run_moe_nvfp4"), + "CUTEDSL_FC12": (CuteDslFc12FusedMoENvfp4Runner, "CuteDslFc12FusedMoE::run_moe_nvfp4"), + } + + +class K3RoutedExperts(nn.Module): + """One MoE layer's routed experts on this rank (see the module docstring). + + ``mapping``: the routed experts' mapping (``moe_tp_size`` x ``moe_ep_size`` over the TP + group). ``aux_stream``: the stream the CUTEDSL output memset and input quantization fork to. + ``max_chunk_tokens``: rows per launch chain (``moe_max_num_tokens``). + """ + + def __init__( + self, + *, + kernels: str, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + mapping: Mapping, + layer_idx: int, + situ_beta: float, + situ_linear_beta: float, + aux_stream: Optional[torch.cuda.Stream], + max_chunk_tokens: int, + ) -> None: + super().__init__() + if kernels not in KERNEL_SETS: + raise ValueError(f"Kimi K3 routed experts run {KERNEL_SETS}, not {kernels!r}") + if kernels == "CUTEDSL_FC12" and ( + get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE + ): + raise ValueError("CUTEDSL_FC12 runs on SM107 with the Rubin CuTe DSL only") + if mapping.enable_attention_dp and mapping.tp_size > 1: + raise NotImplementedError( + "Kimi K3 CuTe DSL routed experts take the tokens of their own rank only " + "(no attention-DP dispatch)" + ) + self.kernels = kernels + self.layer_idx = layer_idx + self.top_k = top_k + self.hidden_size = hidden_size + self.situ_beta, self.situ_linear_beta = float(situ_beta), float(situ_linear_beta) + self.moe_ep_size = mapping.moe_ep_size + self.max_chunk_tokens = max(1, int(max_chunk_tokens)) + self.backend = K3Nvfp4Experts( + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + tp_size=mapping.moe_tp_size, + tp_rank=mapping.moe_tp_rank, + ep_size=mapping.moe_ep_size, + ep_rank=mapping.moe_ep_rank, + ) + self.dwdp_manager = None + self.enable_dwdp = False + self._decide_dwdp() + self.backend.create_weights() + # SMs the FC12 grids leave to kernels running beside them (``reserved_sms``). + self.sm_reserve = 0 + self._aux_stream = aux_stream + self._main_event = torch.cuda.Event() + self._memset_event = torch.cuda.Event() + self._quant_event = torch.cuda.Event() + self.small_batch = K3SmallBatchExperts(self) + + def _decide_dwdp(self) -> None: + """Under a DWDP manager, store only this rank's range of its expert partition and register + the layer; DWDP widens the range to all experts once it is set up.""" + from tensorrt_llm._torch.pyexecutor.dwdp import get_global_dwdp_manager + + manager = get_global_dwdp_manager() + if manager is None: + return + start, end = manager.expert_partition(self.backend.num_experts).range_of(manager.dwdp_rank) + self.backend.set_expert_range(start, end) + self.dwdp_manager = manager + self.enable_dwdp = True + manager.add_layer(layer_idx=self.layer_idx) + + @property + def consumes_sort_tables(self) -> bool: + """Whether the kernels can run on tables a routing kernel wrote (no slot remap).""" + return self.backend.num_slots == self.backend.num_experts + + def post_load_weights(self) -> None: + self.small_batch.post_load_weights() + + # ------------------------------------------------------------------ + # Input quantization. + # ------------------------------------------------------------------ + def quantize_input(self, x: Union[torch.Tensor, Fp4QuantizedTensor]): + """``(x, x_sf)``: packed E2M1 ``[M, H / 2]`` and linear E4M3 block scales ``[M, H / 16]`` + of ``x`` (BF16, quantized against ``fc31_input_scale``; or an unswizzled + ``Fp4QuantizedTensor`` produced upstream with the same scale).""" + rows = x.shape[0] + if isinstance(x, Fp4QuantizedTensor): + assert not x.is_sf_swizzled, "the routed experts take linear (unswizzled) block scales" + x, x_sf = x.fp4_tensor, x.scaling_factor + else: + x, x_sf = fp4_quantize(x, self.backend.fc31_input_scale, SF_VEC_SIZE, False, False) + return x, x_sf.view(rows, -(-self.hidden_size // SF_VEC_SIZE)) + + def _aux(self) -> torch.cuda.Stream: + if self._aux_stream is None: + self._aux_stream = torch.cuda.Stream() + return self._aux_stream + + def _quantize_on_aux(self, x: torch.Tensor, input_ready_event: Optional[torch.cuda.Event]): + """Quantize ``x`` on the aux stream (it overlaps the routing); the returned event marks + the result ready.""" + main_stream = torch.cuda.current_stream() + aux = self._aux() + self._main_event.record() + x.record_stream(aux) + with torch.cuda.stream(aux): + self._main_event.wait() + _wait(input_ready_event) + x_q, x_sf = self.quantize_input(x) + self._quant_event.record() + x_q.record_stream(main_stream) + x_sf.record_stream(main_stream) + return x_q, x_sf, self._quant_event + + # ------------------------------------------------------------------ + # Forward. + # ------------------------------------------------------------------ + def forward( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + *, + route: Route, + routing: Optional[SortedRouting] = None, + input_ready_event: Optional[torch.cuda.Event] = None, + output_push=None, + all_rank_num_tokens: Optional[List[int]] = None, + ) -> torch.Tensor: + """This rank's routed output ``[M, hidden]`` (BF16) of ``x`` (``[M, hidden]`` BF16 or + NVFP4), routed by ``routing`` or else by ``route`` on ``router_logits``. + + ``input_ready_event``: ``x`` is written on another stream, joined before its first read. + ``output_push``: also push the rows into that one-shot all-reduce's slots (small-batch + kernel only). + """ + num_tokens = router_logits.shape[0] + if ( + routing is not None + and routing.sort_tables is not None + and not self.consumes_sort_tables + ): + raise ValueError("sort tables of a routing need experts without a slot remap") + if self.enable_dwdp: + self.dwdp_manager.wait_and_bind(self.backend, self.layer_idx) + if self.small_batch.applies(num_tokens): + x_q, x_sf, routing = self._prologue(x, router_logits, route, routing, input_ready_event) + out = self.small_batch.run( + x_q, + x_sf, + routing.token_selected_experts, + routing.token_final_scales, + routing.sort_tables, + output_push, + ) + else: + if output_push is not None: + raise RuntimeError( + f"an output push needs the small-batch kernel, not {num_tokens} rows" + ) + out = self._forward_chunked( + x, router_logits, route, routing, input_ready_event, all_rank_num_tokens + ) + if self.enable_dwdp: + self.dwdp_manager.finish_layer(self.layer_idx) + return out + + def _prologue(self, x, router_logits, route, routing, input_ready_event): + """Route and quantize: ``(x, x_sf, routing)``; a BF16 ``x`` is joined before the routing, + an NVFP4 one after it.""" + if isinstance(x, torch.Tensor): + _wait(input_ready_event) + input_ready_event = None + if routing is None: + routing = route( + router_logits, self.backend.slot_start, self.backend.expert_size_per_partition + ) + _wait(input_ready_event) + x_q, x_sf = self.quantize_input(x) + return x_q, x_sf, routing + + def _forward_chunked( + self, x, router_logits, route, routing, input_ready_event, all_rank_num_tokens + ): + num_tokens = router_logits.shape[0] + rows = max(all_rank_num_tokens) if self.enable_dwdp and all_rank_num_tokens else num_tokens + num_chunks = max(1, -(-rows // self.max_chunk_tokens)) + if num_chunks == 1: + return self._forward_chunk(x, router_logits, route, routing, input_ready_event) + if routing is not None: + raise RuntimeError("a caller routing covers one launch chain, not several") + _wait(input_ready_event) + base, extra = divmod(num_tokens, num_chunks) + sizes = [base + 1] * extra + [base] * (num_chunks - extra) + if isinstance(x, Fp4QuantizedTensor): + x_chunks = [ + Fp4QuantizedTensor(codes, scales, is_sf_swizzled=False) + for codes, scales in zip( + x.fp4_tensor.split(sizes), x.scaling_factor.view(num_tokens, -1).split(sizes) + ) + ] + else: + x_chunks = x.split(sizes) + outs = [ + self._forward_chunk(xc, lc, route, None, None) + for xc, lc in zip(x_chunks, router_logits.split(sizes)) + ] + return torch.cat(outs) + + def _forward_chunk(self, x, router_logits, route, routing, input_ready_event): + x_ready = None + if ( + self.kernels == "CUTEDSL" + and _ASYNC_INPUT_QUANT + and isinstance(x, torch.Tensor) + and x.shape[0] <= _ASYNC_INPUT_QUANT_MAX_TOKENS + and (do_multi_stream() or torch.cuda.is_current_stream_capturing()) + ): + x_q, x_sf, x_ready = self._quantize_on_aux(x, input_ready_event) + if routing is None: + routing = route( + router_logits, self.backend.slot_start, self.backend.expert_size_per_partition + ) + else: + x_q, x_sf, routing = self._prologue(x, router_logits, route, routing, input_ready_event) + tables = routing.sort_tables + if isinstance(tables, PublishedSortTables): + raise RuntimeError("published routing needs the small-batch kernel") + return self._run_tuned( + x_q, x_sf, routing.token_selected_experts, routing.token_final_scales, tables, x_ready + ) + + def _run_tuned(self, x, x_sf, ids, scales, tables: Optional[SortTables], x_ready): + """The AutoTuner's routing tile for this shape, then the kernel chain at it (on + ``tables`` where they are at that tile).""" + b = self.backend + moe_output = torch.empty( + (scales.size(0), self.hidden_size), dtype=torch.bfloat16, device=x.device + ) + if ids.size(0) == 0: + _wait(x_ready) + return moe_output + runner_cls, key = _tile_runners()[self.kernels] + extra = {} + if self.kernels == "CUTEDSL_FC12": + extra["workload_identity"] = ( + int(ActivationType.SiTu), + self.situ_beta, + self.situ_linear_beta, + ) + runner = runner_cls( + forward_impl=self._run, + num_experts=b.num_slots, + top_k=ids.size(-1), + num_local_experts=b.expert_size_per_partition, + local_expert_offset=b.slot_start, + enable_finalize_fusion=True, + enable_alltoall=False, + **extra, + ) + self._pending_x_ready = x_ready + inputs = [x, ids, scales, x_sf, moe_output, b] + _, tactic = AutoTuner.get().choose_one(key, [runner], runner.get_tuning_config(), inputs) + if tables is not None and runner.tile_size_of(tactic) == tables.tile_size: + return self._run(*inputs, tile_size=tables.tile_size, sort_tables=tables) + return runner(inputs, tactic=tactic) + + def _join_x_ready(self) -> None: + event, self._pending_x_ready = getattr(self, "_pending_x_ready", None), None + _wait(event) + + def _run( + self, + x: torch.Tensor, + ids: torch.Tensor, + scales: torch.Tensor, + x_sf: torch.Tensor, + moe_output: torch.Tensor, + weights: K3Nvfp4Experts, + enable_alltoall: bool = False, + tile_size: int = 128, + sort_tables: Optional[SortTables] = None, + **unused, + ) -> torch.Tensor: + """The kernel chain at ``tile_size`` (on ``sort_tables`` if given, else ``moe_sort``).""" + del enable_alltoall, unused + top_k = ids.size(1) + offset, num_local = weights.slot_start, weights.expert_size_per_partition + if sort_tables is not None: + assert sort_tables.tile_size == tile_size + outs = sort_tables.outputs + else: + outs = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=weights.num_slots, + top_k=top_k, + local_expert_offset=offset, + local_num_experts=num_local, + tile_tokens_dim=tile_size, + ) + tile_to_expert, tile_mn_limit, expanded_to_permuted, permuted_to_expanded, _, num_tiles = ( + outs + ) + fc1_weight = weights.w3_w1_weight.view(torch.float4_e2m1fn_x2) + fc1_scale = weights.w3_w1_weight_scale.view(torch.uint8) + fc2_weight = weights.w2_weight.view(torch.float4_e2m1fn_x2) + fc2_scale = weights.w2_weight_scale.view(torch.uint8) + if self.kernels == "CUTEDSL_FC12": + self._join_x_ready() + torch.ops.trtllm.cute_dsl_nvfp4_fc12_fused_rubin( + input=x.view(torch.float4_e2m1fn_x2), + fc1_weight=fc1_weight, + input_scale=x_sf.view(torch.uint8), + fc1_weight_scale=fc1_scale, + fc1_alpha=weights.fc31_alpha, + tile_idx_to_group_idx=tile_to_expert, + tile_idx_to_mn_limit=tile_mn_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=num_tiles, + global_sf=weights.fc2_input_scale, + fc2_weight=fc2_weight, + fc2_weight_scale=fc2_scale, + fc2_alpha=weights.fc2_alpha, + output=moe_output, + token_final_scales=scales, + expanded_idx_to_permuted_idx=expanded_to_permuted, + num_experts=weights.num_slots, + top_k=top_k, + num_local_experts=num_local, + local_expert_offset=offset, + tile_size=tile_size, + swiglu_limit=float("inf"), + ep_size=self.moe_ep_size, + enable_alltoall=False, + scaling_vector_size=SF_VEC_SIZE, + activation_type=int(ActivationType.SiTu), + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, + reserved_sms=self.sm_reserve, + ) + return moe_output + + rubin = get_sm_version() == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE + memset_stream = self._aux() + self._main_event.record() + moe_output.record_stream(memset_stream) + with torch.cuda.stream(memset_stream): + self._main_event.wait() + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_mn_limit, + expanded_idx_to_permuted_idx=expanded_to_permuted, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=num_tiles, + tile_tokens_dim=tile_size, + top_k=top_k, + ep_size=self.moe_ep_size, + enable_alltoall=False, + ) + self._memset_event.record() + fc1 = ( + torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin + if rubin + else torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell + ) + fc1_kwargs = {"output_tensor": None, "output_sf_tensor": None} if rubin else {} + self._join_x_ready() + h, h_sf = fc1( + input=x.view(torch.float4_e2m1fn_x2), + weight=fc1_weight, + input_scale=x_sf.view(torch.uint8), + weight_scale=fc1_scale, + alpha=weights.fc31_alpha, + tile_idx_to_group_idx=tile_to_expert, + tile_idx_to_mn_limit=tile_mn_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=num_tiles, + global_sf=weights.fc2_input_scale, + num_experts=weights.num_slots, + top_k=top_k, + num_local_experts=num_local, + local_expert_offset=offset, + tile_size=tile_size, + **fc1_kwargs, + activation_type=ActivationType.SiTu, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, + ) + self._memset_event.wait() + fc2 = ( + torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin + if rubin + else torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell + ) + fc2_kwargs = {} if rubin else {"expert_counts": None, "expert_capacity": 0} + fc2( + input=h.view(torch.float4_e2m1fn_x2), + weight=fc2_weight, + input_scale=h_sf.view(torch.uint8), + weight_scale=fc2_scale, + alpha=weights.fc2_alpha, + output=moe_output, + tile_idx_to_group_idx=tile_to_expert, + tile_idx_to_mn_limit=tile_mn_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=num_tiles, + token_final_scales=scales, + num_experts=weights.num_slots, + top_k=top_k, + num_local_experts=num_local, + local_expert_offset=offset, + tile_size=tile_size, + output_dtype=torch.bfloat16, + **fc2_kwargs, + ) + return moe_output + + +__all__ = [ + "K3RoutedExperts", + "KERNEL_SETS", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/moe_runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/moe_runtime.py new file mode 100644 index 000000000000..d6a9a03a8bf7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/moe_runtime.py @@ -0,0 +1,2525 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 on Rubin (SM107): the MoE block. + +``K3MoERuntime`` runs one layer's latent MoE: the router (``K3MoEGate``), the latent +down projection, the routed experts (the CuTe DSL kernel chain of ``moe.py``, +the MegaMoE module, or the small-batch kernels), the latent all-reduce, +``routed_expert_norm`` / ``routed_expert_up_proj``, and the shared experts +(``_K3SharedExpert``) on an auxiliary stream. Its decode tails leave the residual update to +the next layer's attention-residual mix (``attn_res.py``). + +Also here: the routed-expert checkpoint specs (key families and the loader of each +quantization), the workspaces the latent all-reduce and the column-parallel tail share +across layers, and the MoE feature constants, read from ``features.active()`` at import. +""" + +from __future__ import annotations + +import copy +import functools +import os +from contextlib import contextmanager +from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.noaux_tc_op import noaux_tc_op +from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._torch.utils import ActivationType, AuxStreamType, Fp4QuantizedTensor +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from .attn_res import _K3_ULL_TAIL_SANDWICH, _DeferredMoEResidual, _DeferredRpTail +from .features import active as _features +from .fp8_weight_read import ( + _K3_FP8W_LOWM_GEMM, + _K3_LD_FUSED, + _Fp8BlockScaleWeightReadLinear, + _k3_lowm_gemm_can_run, +) +from .layers import ( + Bf16Linear, + GatedMlp, + LamportPending, + LamportWorkspaceView, + RMSNormWeight, + TpAllReduce, + finish_lamport_push, + fork_join, + gated_activation, + lamport_push_target, + subgroup_mapping, +) +from .layers import mnnvl_workspace_mpi_comm as _k3_mnnvl_mpi_comm +from .moe import KERNEL_SETS as _K3_ROUTED_KERNEL_SETS +from .moe import K3Nvfp4Experts, K3RoutedExperts, PublishedSortTables, SortedRouting, SortTables +from .ops.comm.k3_lamport_push import k3_lamport_push +from .ops.comm.k3_latent_ar_rmsnorm import k3_latent_ar_rmsnorm +from .ops.gemm.k3_fp8w_lowm_gemm_push import k3_fp8w_lowm_gemm_push as k3_fp8w_lowm_gemm_push_op +from .ops.gemm.k3_fp8w_lowm_mxfp8_push import k3_fp8w_lowm_mxfp8_push as k3_fp8w_lowm_mxfp8_push_op +from .ops.gemm.k3_fp8w_mid_gemm import k3_fp8w_mid_gemm as k3_fp8w_mid_gemm_op +from .ops.gemm.k3_latent_down_fp4 import k3_latent_down_fp4 as k3_latent_down_fp4_op +from .ops.moe.k3_moe_residual_add import k3_moe_residual_add as k3_moe_residual_add_op +from .ops.moe.k3_moe_tail1 import k3_moe_tail1 as k3_moe_tail1_op +from .ops.moe.k3_route_tail import k3_route_tail as k3_route_tail_op +from .ops.moe.k3_router_logits import k3_router_logits as k3_router_logits_op +from .ops.moe.k3_shared_persist import k3_shared_persist as k3_shared_persist_op +from .ops.moe.k3_ull_front import k3_ull_front as k3_ull_front_op +from .text_config import _is_kda_layer, _resolve_kimi_situ_betas + +_F = _features() + +# Decode-batch routing tail of the CuteDSL MoE (see k3_route_tail): replaces noaux_tc_op -> +# moe_sort with one launch on a cluster of at most eight CTAs (same logits, same routing). +# Largest token count the routing tail serves: up to 256 tokens one cluster, above that (at most +# 2048 tokens, see k3_route_tail.wide_sort_fits) a grid of CTAs whose last CTA does the sort. +_K3_ROUTE_TAIL_MAX_M = 2048 + +_K3_LD_FUSED_MIN_M = 1 +_K3_LD_FUSED_MAX_M = 4096 + +# For up to _K3_ROUTER_FIRST_MAX_M tokens, the router GEMM is issued (captured) on the main stream +# before latent_down is issued on the latent stream. Both depend only on the MoE input (the latent +# fork event is recorded before the router), so the graph has the same edges as the reverse +# order; only the node creation order differs, which decides which of the two the +# graph launches first once the MoE input's producer releases its dependents (PDL). At decode +# batch sizes latent_down's low-M GEMM is 112 CTAs x ~225 KB of shared memory, and with up to 16 +# tokens the router GEMM (dsv3 router MMA kernel) is 112 CTAs whose registers fill most of an SM, +# so no SM holds one of each: launched second, the router gets only the SMs latent_down leaves +# and runs in several waves, and the routing (which gates the routed FC1 together with +# latent_down) ends late. Launched first, the router takes its SMs and latent_down starts on the +# rest. nsys node tracing turns latent_down's cross-stream programmatic edge into a full edge (it +# then waits for the MoE input's producer to complete), so such traces show the router first in +# either order. +_K3_ROUTER_FIRST_MAX_M = 64 + +# c240: SM-count target of the shared expert's cuBLAS GEMMs (0: off). FC12's persistent grid needs +# whole SMs, so shared-expert GEMMs still resident when routing finishes delay its launch; with +# TRTLLM_K3_SHARED_SM=N they run on N SMs (torch SM carveout, captured into CUDA graphs) and the +# routed experts' backend keeps N SMs free of its persistent grids +# (``reserve_sms_for_concurrent_kernels``, K3MoERuntime). +_K3_SHARED_SM = _F.shared_sm +# Two-domain devices: die-major expert weights. Not available: the routed experts keep the loaded +# layout (``_k3_routed_backend_opt_ins`` logs it once). +_K3_FC12_DIE_MAJOR = _F.fc12_die_major + + +@contextmanager +def _k3_shared_sm_scope(): + """Run the enclosed cuBLAS GEMMs on ``TRTLLM_K3_SHARED_SM`` SMs (no-op when 0).""" + if _K3_SHARED_SM <= 0: + yield + return + num_sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + prev = torch._C._get_sm_carveout_experimental() + torch._C._set_sm_carveout_experimental(max(0, num_sms - _K3_SHARED_SM)) + try: + yield + finally: + torch._C._set_sm_carveout_experimental(prev) + + +# Small-batch MoE front end (b1-b4 decode), bit-identical to the path it replaces: +# TRTLLM_K3_ROUTER_EXACT=kda: in the MoE blocks of KDA layers, the router logits of M <= 16 +# tokens come from k3_router_exact (bit-identical to dsv3_router_gemm_op; its weight is copied +# into shared memory before griddepcontrol.wait). The faster router also starts the route tail +# earlier, inside the burst of latent_down and the shared gate_up; after an MLA attention block +# the route tail then stretches by ~3 us (b2 gen-only, GB300) and the routing chain gets longer +# than with dsv3_router_gemm_op, while after a KDA attention block the front gets shorter. +_K3_ROUTER_EXACT_KDA = _F.router_exact_kda +# TRTLLM_K3_LATENT_AR_OWN_WS=1: for up to _K3_LATENT_AR_OWN_WS_MAX_M tokens, the routed latent +# all-reduce + RMSNorm (the same MNNVL Lamport kernels) runs on a workspace of its own instead of the +# one it shares with the shared expert's all-reduce. It then no longer has to follow the shared-expert +# chain: the main stream issues it and latent_up right after the routed FC2 and joins the shared +# stream after latent_up. Applies where latent_up does not add the shared output itself (no fused +# latent_up + residual variant); the output is bit-identical. +# Up to _K3_LATENT_AR_ONESHOT_MAX_BYTES (rows x latent x ranks x 2 B; the workspace is sized for it, +# the other MNNVL all-reduces keep the global one-shot limit: 6291456 runs the 96-row latent +# all-reduce of 8 ranks, 5.5 MB, one-shot) that all-reduce runs k3_latent_ar_norm, bit-identical to +# the C++ one-shot kernel but triggering its programmatic dependents (latent_up) right after its +# griddepcontrol.wait instead of after the cross-rank reduction: its row kernel where that serves +# the size (one 512-thread CTA per row instead of a cluster of four; the rank's own slot is summed +# from registers instead of being polled back from the multicast loopback; slots that arrived are +# not reloaded), else the cluster kernel; elsewhere the C++ kernel (trtllm::mnnvl_fusion_allreduce) +# runs on the same workspace. +# Where latent_up runs the mid-M path of the k3_fp8w_lowm_gemm entry (above the low-M cut-over, e.g. +# the 96-row helix8 b16 mtp5 verify step), the row kernel also stores the MXFP8 form of the +# normalized latent in that GEMM's scale layout and latent_up runs the mid-M GEMM on it +# (k3_fp8w_lowm_gemm.run_mid_prequantized) instead of quantizing the BF16 rows in a kernel of its +# own: bitwise the same projection, one kernel less after the all-reduce. +_K3_LATENT_AR_OWN_WS = _F.latent_ar_own_ws +_K3_LATENT_AR_OWN_WS_MAX_M = 128 +_K3_LATENT_AR_ONESHOT_MAX_BYTES = 6291456 +# The routed latent all-reduce's own MNNVL workspaces, by all-reduce mapping. +_K3_LATENT_AR_WORKSPACES: Dict[Mapping, Dict[str, Any]] = {} + + +def _k3_latent_ar_workspace( + mapping: Mapping, dtype: torch.dtype, hidden_size: int, max_tokens: int +) -> Dict[str, Any]: + """MNNVL Lamport workspace of the routed latent all-reduce alone (three buffers and the state + flags, laid out as ``get_or_scale_allreduce_mnnvl_workspace`` lays out the shared one), sized + for up to ``max_tokens`` rows of ``hidden_size``. Collective over the all-reduce group: every + rank builds it at the same point of model construction.""" + workspace = _K3_LATENT_AR_WORKSPACES.get(mapping) + if workspace is not None: + if workspace["max_tokens"] < max_tokens: + raise ValueError( + f"Kimi K3 latent all-reduce workspace serves {workspace['max_tokens']} tokens, " + f"{max_tokens} requested" + ) + return workspace + from tensorrt_llm._torch.distributed.ops import ( + MNNVLAllReduce, + get_or_scale_allreduce_mnnvl_workspace, + ) + from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer + + # The shared workspace's communicator (split by the all-reduce group) is the same group. + comm = _k3_mnnvl_mpi_comm(get_or_scale_allreduce_mnnvl_workspace(mapping, dtype)) + need = max( + MNNVLAllReduce.get_required_workspace_size(m, hidden_size, mapping.tp_size, dtype) + for m in range(1, max_tokens + 1) + ) + # One-shot rows under the latent all-reduce's own limit: every rank's rows in one buffer. + row_bytes = hidden_size * mapping.tp_size * torch.tensor([], dtype=dtype).element_size() + one_shot_rows = min(max_tokens, _K3_LATENT_AR_ONESHOT_MAX_BYTES // row_bytes) + need = max(need, one_shot_rows * row_bytes) + buffer_bytes = -(-need // (1 << 20)) << 20 + use_fabric_handle = ( + os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node() + ) + handle = McastGPUBuffer( + 3 * buffer_bytes, + mapping.tp_size, + mapping.tp_rank, + mapping.local_rank, + use_fabric_handle, + comm.py2f(), + ) + # Lamport-clean (-0.0) before any rank writes into it. + buffer = handle.get_uc_buffer(mapping.tp_rank, (3 * buffer_bytes // 4,), torch.float32, 0) + buffer.fill_(-0.0) + torch.cuda.synchronize() + comm.Barrier() + # [current index, dirty index, bytes per buffer, dirty stages, bytes to clear x4, arrivals] + flags = torch.tensor( + [0, 2, buffer_bytes, 0, 0, 0, 0, 0, 0], dtype=torch.uint32, device=buffer.device + ) + workspace = { + "handle": handle, + "buffer": buffer.view(dtype).view(3, -1), + # The multicast view of the same buffers (stores reach every rank's copy). + "mc": handle.get_mc_buffer((3 * buffer_bytes // 4,), torch.float32, 0), + "flags": flags, + "max_tokens": max_tokens, + # Bytes of one of the three Lamport buffers. + "buffer_bytes": buffer_bytes, + "rank": mapping.tp_rank, + "world_size": mapping.tp_size, + } + _K3_LATENT_AR_WORKSPACES[mapping] = workspace + return workspace + + +# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections +# instead of the min-latency fused GEMM op (read once at import). +_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( + os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" +) + + +class K3MoEGate(nn.Module): + """Kimi K3 gate: the router weight, the score-correction bias and the routing kernels. + + Routing is ungrouped sigmoid top-k with the bias added for the selection and the selected + scores renormalized and scaled (``noaux_tc``); :meth:`route`, :meth:`route_with_sort` and + :meth:`route_tail` are three kernels of the same routing. + """ + + def __init__( + self, + config: Any, + *, + logits_gemm_dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + + if self.moe_router_activation_func != "sigmoid" or not self.moe_renormalize: + raise ValueError( + "Kimi K3 routing takes sigmoid scores with top-k weight renormalization" + ) + if self.num_expert_group != 1 or self.topk_group != 1: + raise ValueError("Kimi K3 routing is ungrouped (num_expert_group = topk_group = 1)") + if self.num_experts > 1024 or self.top_k > 32: + raise ValueError( + f"the fused routing kernel serves up to 1024 experts and top-32, got " + f"{self.num_experts} and top-{self.top_k}" + ) + + # The checkpoint stores the gate weight in bf16. Storing it in bf16 + # permits the single bf16xbf16 router GEMM while retaining fp32 output. + weight_dtype = logits_gemm_dtype or torch.float32 + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32, device=device) + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_2d.contiguous(), + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return torch.nn.functional.linear( + hidden_2d.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + + def route(self, logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """``(ids int32, weights fp32)``, each ``[num_tokens, top_k]``: ``noaux_tc_op``.""" + values, indices = noaux_tc_op( + logits, + self.e_score_correction_bias, + self.num_expert_group, + self.topk_group, + self.top_k, + self.routed_scaling_factor, + ) + return indices.to(torch.int32), values.to(torch.float32) + + def route_with_sort( + self, logits: torch.Tensor, local_expert_offset: int, local_num_experts: int + ) -> Optional[SortedRouting]: + """The routing with the ``moe_sort`` tables of experts ``[local_expert_offset, + local_expert_offset + local_num_experts)`` in one launch (``noaux_tc_route_sort``, bitwise + ``noaux_tc_op`` + ``moe_sort``), or ``None`` where that kernel does not serve.""" + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + if not IS_CUTLASS_DSL_AVAILABLE: + return None + from tensorrt_llm._torch.moe.kernels.blackwell import kimi_k3_routing + + if self.top_k != kimi_k3_routing.TOP_K or not kimi_k3_routing.serves( + logits.shape[0], logits.shape[1], local_expert_offset, local_num_experts, logits.device + ): + return None + ids, weights, *tables = torch.ops.trtllm.noaux_tc_route_sort( + logits, + self.e_score_correction_bias, + float(self.routed_scaling_factor), + local_expert_offset, + local_num_experts, + ) + return SortedRouting(ids, weights, SortTables(kimi_k3_routing.TILE_SIZE, tuple(tables))) + + def route_tail( + self, logits: torch.Tensor, local_expert_offset: int, local_num_experts: int + ) -> SortedRouting: + """The CuTe routing tail (``k3_route_tail``, top-k bitwise ``noaux_tc_op``) with the + ``moe_sort`` tables of experts ``[local_expert_offset, local_expert_offset + + local_num_experts)``.""" + from .kernels import k3_route_tail + + # Above the cluster kernel's tokens, the multi-CTA sort's persistent workspace. + workspace = ( + k3_route_tail.wide_workspace(logits.device) + if logits.shape[0] > k3_route_tail.MAX_M + else None + ) + outs = k3_route_tail_op( + logits, + self.e_score_correction_bias, + float(self.routed_scaling_factor), + local_expert_offset, + local_num_experts, + k3_route_tail.TILE, + workspace=workspace, + ) + return SortedRouting(outs[0], outs[1], SortTables(k3_route_tail.TILE, tuple(outs[2:8]))) + + +def _router_exact_for_layer(cfg, layer_idx: int) -> bool: + """Whether ``TRTLLM_K3_ROUTER_EXACT`` serves the router of this layer's MoE block.""" + return _K3_ROUTER_EXACT_KDA and _is_kda_layer(cfg, layer_idx) + + +# TRTLLM_K3_ULL2_COLL_CP_TAIL=1 (BF16 shared expert): the MoE tail of a +# generation-only forward splits latent_up by output columns instead of by latent columns. After +# the latent all-reduce + RMSNorm (k3_latent_ar_norm's row kernel with its MXFP8 output, as +# TRTLLM_K3_ULL_TAIL_SANDWICH) every rank computes its whole 128-row blocks of latent_up's output +# over the full latent (the MXFP8 low-M GEMM with the replicated projection's tactic) and multicasts +# them into the all-gather rows of the next call on a Lamport workspace of the tail's own +# (k3_lamport_ar "ag"); the shared expert's BF16 partial is pushed into the same call's slots on +# the shared stream as soon as it is ready. The next layer's pre-attention mix finishes the call +# (k3_attn_res ``finish_ag``): bf16(latent_up + bf16(all-reduced shared)), the replicated tail's +# bf16(routed + shared), so the result is bitwise the replicated tail (latent_up on every rank + +# the shared expert's MNNVL all-reduce) while each rank reads 1 / W of latent_up as the +# row-parallel tail does. Compared with the row-parallel tail, the critical path after the latent +# all-reduce carries 1 / W of the bytes (W x M x 7168 x 2 B of slots land on every rank, but the +# shared expert's land while the experts still run) and no two-shot step above the one-shot size. +# The latent all-reduce runs one-shot on a workspace of the tail's own, up to the C++ one-shot size +# (the sizes where the C++ all-reduce is one-shot too, so the sums keep its order); larger sizes keep +# the row-parallel tail. The mix polls the call before its griddepcontrol.wait (``early_poll``). +# The tail's latent all-reduce row kernel reads its workspace's rotation state, counts itself and +# re-arms the previous call's buffer before its griddepcontrol.wait (k3_latent_ar_norm +# ``early_state``: the previous call on that workspace, the previous MoE layer's, completed long +# before the routed experts launched). Same bits. +# With the routed experts' output push, the tail of a generation-only forward of at most 8 rows is +# one kernel (k3_moe_tail1). The routed experts (CUTEDSL_FC12's small-batch kernel) push this +# rank's latent partial rows, combined across experts with BF16 vector reds (whose order, and so +# the rows' bits, may change from run to run), into the slots of the latent +# workspace (``output_push``); k3_moe_tail1 polls them, sums them in the latent all-reduce's order, +# applies the RMSNorm after the products, multiplies by this rank's HIDDEN / W rows of latent_up and +# multicasts them into the column-parallel tail call's all-gather rows, which the next +# pre-attention mix finishes as before (``finish_ag``). It replaces the latent all-reduce row kernel +# and the column-parallel GEMM. Not bitwise (BF16 activation instead of MXFP8, the RMS scale applied +# to the FP32 products). +# On the column-parallel tail path of at most 8 rows, the shared expert (gate/up, SiTU, down and +# the push of its partial into the tail call's slots) is one persistent kernel on +# _K3_ULL_SHARED_PERSIST_CTAS CTAs (k3_shared_persist) on the shared stream, beside the routed +# experts, whose backend leaves it that many SMs (``reserve_sms_for_concurrent_kernels``). +# The next pre-attention mix reads the slots as before. Not bitwise (BF16 GEMM products in another +# summation order, tanh / sigmoid with ex2.approx). Needs the BF16 GatedMlp shard with SiTU on the +# tail's TP group; the weights are copied once into the kernel's stream order (16.5 MB per MoE layer +# and rank at TP16). +_K3_ULL2_COLL_CP_TAIL = _F.ull2_coll_cp_tail +_K3_ULL_SHARED_PERSIST_CTAS = 16 + +# The shared expert's part of the forward was pushed by the persistent kernel (no partial). +_SHARED_PUSHED = object() + + +# The column-parallel tail's Lamport workspaces, by (use, all-reduce mapping), the same workspaces +# as the routed experts' ``output_push`` takes them, and the multicast handles that own their memory. +_K3_CP_TAIL_WORKSPACES: Dict[Tuple[str, Mapping], Any] = {} +_K3_CP_TAIL_PUSH_VIEWS: Dict[Tuple[str, Mapping], LamportWorkspaceView] = {} +_K3_CP_TAIL_HANDLES: Dict[Tuple[str, Mapping], Any] = {} + + +def _k3_cp_tail_workspace(mapping: Mapping, hidden_size: int, max_tokens: int, use: str = "tail"): + """``k3_lamport_ar.LamportWorkspace`` of the column-parallel MoE tail + (``TRTLLM_K3_ULL2_COLL_CP_TAIL``) and the rotation state. ``use="tail"``: three Lamport buffers + of ``max_tokens x (W + 1)`` BF16 rows of ``hidden_size`` (the one-shot slots plus the + all-gather rows); ``"latent"``: ``max_tokens x W`` one-shot rows. + Collective over the all-reduce group: every rank builds it at the same point of model + construction. ``_K3_CP_TAIL_HANDLES[(use, mapping)]`` holds the multicast buffer.""" + workspace = _K3_CP_TAIL_WORKSPACES.get((use, mapping)) + if workspace is not None: + return workspace + from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace + from tensorrt_llm.bindings.internal.runtime import McastGPUBuffer + + from .kernels.k3_lamport_ar import LamportWorkspace, required_bytes + + comm = _k3_mnnvl_mpi_comm(get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16)) + need = required_bytes(max_tokens, hidden_size, mapping.tp_size, True, ag=use == "tail") + buffer_bytes = -(-need // (1 << 20)) << 20 + use_fabric_handle = ( + os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node() + ) + handle = McastGPUBuffer( + 3 * buffer_bytes, + mapping.tp_size, + mapping.tp_rank, + mapping.local_rank, + use_fabric_handle, + comm.py2f(), + ) + buffer = handle.get_uc_buffer(mapping.tp_rank, (3 * buffer_bytes // 4,), torch.float32, 0) + buffer.fill_(-0.0) + torch.cuda.synchronize() + comm.Barrier() + flags = torch.tensor( + [0, 2, buffer_bytes, 0, 0, 0, 0, 0, 0], dtype=torch.uint32, device=buffer.device + ) + peers = [ + handle.get_uc_buffer(r, (1,), torch.uint8, 0).data_ptr() for r in range(mapping.tp_size) + ] + workspace = LamportWorkspace( + world=mapping.tp_size, + rank=mapping.tp_rank, + uc_local=buffer.data_ptr(), + mc_base=handle.get_mc_buffer((1,), torch.uint8, 0).data_ptr(), + peer_ptrs=torch.tensor(peers, dtype=torch.int64, device=buffer.device), + flags=flags, + bytes_per_buffer=buffer_bytes, + ) + _K3_CP_TAIL_WORKSPACES[(use, mapping)] = workspace + _K3_CP_TAIL_PUSH_VIEWS[(use, mapping)] = LamportWorkspaceView( + world=workspace.world, + rank=workspace.rank, + buffer=buffer, + mc_base=workspace.mc_base, + peer_ptrs=workspace.peer_ptrs, + flags=flags, + bytes_per_buffer=buffer_bytes, + ) + _K3_CP_TAIL_HANDLES[(use, mapping)] = handle + return workspace + + +# Row-parallel MoE tail for generation-only batches of up to _K3_RP_TAIL_MAX_M tokens (the low-M +# GEMM's range), where the routed latent is all-reduced by the model (MoE TP/EP over the all-reduce +# group, TP-sharded shared expert) and the layer's residual update is fused +# (TRTLLM_K3_FUSE_MOE_RESIDUAL). After the latent all-reduce + RMSNorm, every rank runs one low-M +# FP8 GEMM over its whole 128-column blocks of the normalized latent with the matching latent_up +# columns (FP8 codes and block scales read in place) and pushes its partial [M, hidden] plus its +# shared-expert partial into the next all-reduce on the latent all-reduce's MNNVL workspace: the +# shared expert keeps its BF16 down GEMM (row-parallel shard, no all-reduce) on its stream, and the +# tail GEMM's push epilogue adds that BF16 partial to its BF16 output. The next layer's +# pre-attention mix finishes that all-reduce (one-shot transport up to +# lamport._K3_FUSED_FINISH_ONE_SHOT_MAX_BYTES, two-shot above) and adds it to the residual +# (``_DeferredRpTail``). This replaces the shared expert's all-reduce and the replicated latent_up. +# Not bitwise: the latent_up sum is split over the ranks. +_K3_RP_TAIL_MAX_M = 32 +# TRTLLM_K3_ULL_FRONT=1: MoE forwards of 1 .. _K3_ULL_FRONT_MAX_M (<= 32) rows run the MoE +# front of cute_dsl_kernels/blackwell/k3_ull_front.py on the routed all-reduce's TP group, two kernels +# on the main stream: each rank computes the router logits of every expert (k3_router_exact's +# arithmetic) and the latent_down columns of its 128-column tiles (the low-M GEMM of +# k3_latent_down_fp4) into an MNNVL workspace (the latent is multicast to every rank); a routing-tail +# cluster polls the logits and writes the top-16 and this rank's moe_sort tables, and gather CTAs +# write the NVFP4 latent. This replaces the router GEMM, the routing tail, the replicated latent_down +# + quantization and the latent stream's fork / join. +# Routing is bitwise (per 16-row tile of the router); the latent is bitwise up to 16 rows (above, +# k3_latent_down_fp4 runs the mid-M kernel, split 4, and the front the low-M kernel, split 4). +# The routing reaches the routed experts as ``routing``: where they offer a +# routing word for the call (``published_routing_word``: the small-batch kernel on MoE-TP slices), +# the route kernel publishes each token's top-k into it instead of writing the moe_sort tables and +# the experts build their tables from it while the route kernel still runs +# (``PublishedSortTables``; same tables, same bits); otherwise the front's tile-128 moe_sort tables +# (``SortTables``) replace the experts' own moe_sort. +# The A_mid that finishes the pushed o_proj all-reduce releases the front's head at kernel entry +# (k3_attn_res ``entry_trigger``) where the head follows it, so the head's CTAs become resident +# while A_mid polls the o_proj all-reduce and stream the router and latent_down weights into +# shared memory before their griddepcontrol.wait; after the wait only the activation loads, the +# MMAs and the pushes remain. Before its wait the head reads only weights and writes nothing, and +# it releases the route kernel only after its wait. The early-polling A_mid otherwise releases it +# after its stores. Same bits. +# Forwards of up to _K3_ULL_FRONT_FORK_AT_HEAD_MAX_M rows fork the shared expert after the head +# (between the head and the route kernel): under CUDA graph capture the shared GEMMs become +# programmatic dependents of the kernel the fork follows, so a fork before the head lets the +# entry-releasing A_mid release them with the head, and their CTAs take the SMs the head's CTAs +# need while A_mid polls (helix16 n1 / n2: -3.4 / -2.6% per step forking at the head vs -0.3 / +# -0.8% before it; at 24 / 32 rows forking before the head stays ahead, -1.5 / -1.3% vs -1.0 / +# +0.1%, where the larger gate_up competes with the routed experts instead). +_K3_ULL_FRONT = _F.ull_front +_K3_ULL_FRONT_MAX_M = 32 +_K3_ULL_FRONT_FORK_AT_HEAD_MAX_M = 16 + + +def _ull_shared_fork(lead: bool, rows: int) -> Optional[str]: + """Where the ULL front forks the shared expert's GEMMs on ``rows`` rows (``None``: no fork, ``lead`` is + ``_shared_gemms_can_lead()``). + + * ``"head"`` (up to ``_K3_ULL_FRONT_FORK_AT_HEAD_MAX_M`` rows): the event is recorded after the head, so + the shared GEMMs depend on the head. + * ``"after_head"`` (more rows): the event is recorded before the head, so the shared GEMMs depend on what + precedes the head (the entry-releasing A_mid). + + The shared GEMMs themselves are issued after the front (``K3MoERuntime._ull_fork_shared``). + """ + if not lead: + return None + return "head" if rows <= _K3_ULL_FRONT_FORK_AT_HEAD_MAX_M else "after_head" + + +# --------------------------------------------------------------------------- +# Latent MoE block: router, routed experts (kernel chains of moe.py), shared expert. +# --------------------------------------------------------------------------- + +# Routed-expert key spellings that ModelOpt emits for Kimi K3. The NVFP4 +# checkpoint (``nvidia/Kimi-K3-NVFP4``) lists every prefix x module-name +# combination in ``quantized_layers``, so a lookup over this product finds it +# without needing the MiniMax-M3-style prefix normalization in ``ModelConfig``. +_K3_ROUTED_EXPERT_KEY_PREFIXES = ("language_model.model.", "model.", "") +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + +# The subset of the above that can be a real module path. ``exclude_modules`` +# matches with wildcards and walks ancestor prefixes, so an empty prefix would +# widen what matches instead of just missing, as it does in the dict lookup. +_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") + +# Routed-expert quantization used when the checkpoint declares nothing per +# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors +# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that +# checkpoint is what this default has always served. +_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 + + +# --------------------------------------------------------------------------- +# Routed-expert checkpoint layout. +# +# K3 streams routed experts one at a time (see ``load_weights``): the tensor names, the per-expert +# loader and the per-layer finalize it needs are decided by the checkpoint's routed-expert +# quantization. The expected-key plan, the loader and the file grouping read them off one spec. +# --------------------------------------------------------------------------- + + +def _load_nvfp4_expert( + backend, base, expert_idx, local_slot_id, get_tensor, get_optional_tensor +) -> None: + for w in ("w1", "w3"): + if get_optional_tensor(f"{base}.{expert_idx}.{w}.pre_quant_scale") is not None: + raise NotImplementedError( + "Kimi K3 routed experts take plain NVFP4 checkpoints; this one carries a " + "pre-quant activation scale (NVFP4_AWQ / NVFP4_ARC)" + ) + tensors = { + f"{w}_{kind}": get_tensor(f"{base}.{expert_idx}.{w}.{kind}") + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + } + if isinstance(backend, K3Nvfp4Experts): + backend.load_expert(global_expert_id=expert_idx, local_slot_id=local_slot_id, **tensors) + else: + backend.quant_method.load_streaming_nvfp4_expert( + backend, global_expert_id=expert_idx, local_slot_id=local_slot_id, **tensors + ) + + +def _k3_routed_backend_opt_ins(routed_experts, column_parallel_tail: bool) -> None: + """The routed experts' opt-ins of a MoE layer, made once at construction. + + SMs the FC12 kernels' persistent grids leave to the kernels the shared stream runs beside them + (the shared expert's GEMMs confined to ``_K3_SHARED_SM`` SMs; with the column-parallel tail, + the persistent shared expert's ``_K3_ULL_SHARED_PERSIST_CTAS`` CTAs). The CUTEDSL kernels take + no reserve. + """ + if _K3_FC12_DIE_MAJOR: + logger.info_once( + "TRTLLM_K3_FC12_DIE_LAYOUT: die-major expert weights are not available; the routed " + "experts keep the loaded layout", + key="k3_fc12_die_major_off", + ) + if getattr(routed_experts, "kernels", None) != "CUTEDSL_FC12": + return + reserve = _K3_SHARED_SM + if column_parallel_tail: + reserve = max(reserve, _K3_ULL_SHARED_PERSIST_CTAS) + if reserve > 0: + routed_experts.sm_reserve = reserve + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. Every + # combination is required: a missing one fails the key plan. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # The loader stages per-expert state that the layer's finalize consumes. + needs_layer_finalize: bool + # Whole per-expert key suffixes the loader reads WHEN PRESENT (to reject them). + optional_suffixes: Tuple[str, ...] = () + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.NVFP4: _K3ExpertCkptSpec( + kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"), + loader=_load_nvfp4_expert, + loaded_slots_attr="_streamed_expert_slots", + needs_layer_finalize=True, + optional_suffixes=("w1.pre_quant_scale", "w3.pre_quant_scale"), + ), +} + + +class _RpTailOperands(NamedTuple): + """One rank's operands of the row-parallel MoE tail GEMM: K is this + rank's ``k1`` latent columns from column ``k0`` (``k1 / 128`` whole blocks of latent_up, whose + FP8 codes the GEMM reads in place); the GEMM's push epilogue adds the shared expert's BF16 + partial.""" + + k0: int + k1: int + # DeepGEMM-packed UE8M0 scale of the [hidden, k1] tail weight. + scale: torch.Tensor + + +class _CpTailOperands(NamedTuple): + """One rank's operands of the column-parallel MoE tail (``TRTLLM_K3_ULL2_COLL_CP_TAIL``): the + ``n`` output rows of latent_up from row ``c0`` (whole 128-row blocks; their FP8 codes read in + place) and the DeepGEMM-packed UE8M0 scale of that ``[n, latent]`` slice.""" + + c0: int + n: int + weight: torch.Tensor + scale: torch.Tensor + + +def _rp_tail_latent_slice(latent_size: int, world: int, rank: int) -> Tuple[int, int]: + """``(first column, columns)`` of ``rank``'s 128-column blocks of the latent: whole blocks, + as even as possible (28 blocks over 8 ranks: 4 on ranks 0-3, 3 on ranks 4-7).""" + blocks = latent_size // 128 + base, extra = divmod(blocks, world) + first = rank * base + min(rank, extra) + return first * 128, (base + (rank < extra)) * 128 + + +def _k3_expert_ckpt_spec(quant_config: QuantConfig) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_config.quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_config.quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + group_size = getattr(quant_config, "group_size", None) + if group_size not in (None, 16): + raise NotImplementedError( + f"Kimi K3 NVFP4 routed experts take 16-element scale blocks, got group_size={group_size}" + ) + return spec + + +class _NoSmallBatch: + """The small-batch path of routed experts that have none.""" + + max_tokens = 0 + counters = None + + @staticmethod + def applies(num_tokens: int) -> bool: + return False + + @staticmethod + def published_routing_word(num_tokens: int) -> None: + return None + + @staticmethod + def post_load_weights() -> None: + pass + + +_NO_SMALL_BATCH = _NoSmallBatch() + + +class _K3SharedExpert(GatedMlp): + """The MoE layer's shared expert: :class:`GatedMlp` called through its projection modules, + so a projection ``_convert_moe_mlps_to_fp8_weight_read`` swapped for an FP8 weight read runs + its own kernel (a ``Bf16Linear`` without bias or all-reduce is ``F.linear`` either way).""" + + def forward( + self, x: torch.Tensor, final_all_reduce_params: Optional[AllReduceParams] = None + ) -> torch.Tensor: + h = self.gate_up_proj(x) + a = gated_activation(h, self.activation, self.beta, self.linear_beta) + return self.down_proj(a, all_reduce_params=final_all_reduce_params) + + +def _megamoe_routed_experts( + gate: "K3MoEGate", + routed_moe_model_config: ModelConfig, + routed_quant_config: QuantConfig, + num_experts: int, + hidden_size: int, + intermediate_size: int, + dtype: torch.dtype, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + situ_beta: float, + situ_linear_beta: float, +) -> nn.Module: + """The routed experts of a ``MEGAMOE_CUTEDSL`` layer: the MegaMoE module, which runs the + token exchange, both GEMMs and the combine in one kernel over symmetric memory.""" + from tensorrt_llm._torch.moe.fused_moe import ConfigurableMoE, SiTuActivation, create_moe + from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod + + routing_method = DeepSeekV3MoeRoutingMethod( + top_k=gate.top_k, + n_group=gate.num_expert_group, + topk_group=gate.topk_group, + routed_scaling_factor=gate.routed_scaling_factor, + callable_e_score_correction_bias=lambda: gate.e_score_correction_bias, + is_fused=True, + ) + moe = create_moe( + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + # Kimi owns the latent reduction so it can order that collective + # after the shared expert's auxiliary-stream reduction. + reduce_results=False, + # Lets the fused-communication kernel return its globally combined routed output. + accepts_global_routed_output=True, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict, + communication_method=None, + activation=SiTuActivation(gate_softcap=situ_beta, linear_softcap=situ_linear_beta), + allow_backend_degradation=False, + ) + if not isinstance(moe, ConfigurableMoE) or moe.layer_load_balancer is not None: + raise RuntimeError("Kimi K3 MegaMoE layers need ConfigurableMoE without EPLB") + return moe + + +class K3MoERuntime(nn.Module): + """Kimi K3 latent MoE block: the gate routes, the routed experts run the CuTe DSL kernel chain + of ``moe`` (MEGAMOE_CUTEDSL: the MegaMoE module), the shared expert runs beside + them.""" + + # TRTLLM_K3_LATENT_AR_OWN_WS: the routed latent all-reduce's own MNNVL workspace (``None``: it + # shares the one of the shared expert's all-reduce). + _latent_ar_ws: Optional[Dict[str, Any]] = None + # TRTLLM_K3_ULL2_COLL_CP_TAIL: the column-parallel tail's workspace and this rank's operands. + _cp_ws: Optional[Any] = None + _cp_latent_ws: Optional[Any] = None + # The latent workspace as the routed experts' ``output_push`` (k3_moe_tail1). + _cp_latent_push: Optional[LamportWorkspaceView] = None + _cp_tail: Optional[_CpTailOperands] = None + # k3_moe_tail1 (column-parallel tail): latent_up's UE8M0 block exponents + # [HIDDEN / 128, LATENT / 128] (None: not built). + _tail1_wexp: Optional[torch.Tensor] = None + # k3_shared_persist (column-parallel tail): the shared expert's weights in the kernel's stream + # order, its exchange workspace and SiTU's (beta, linear beta) (None: not built). + _persist_frags: Optional[torch.Tensor] = None + _persist_ws: Optional[Any] = None + _persist_act: Optional[Tuple[float, Optional[float]]] = None + # TRTLLM_K3_ROUTER_EXACT: whether this layer's router logits come from k3_router_exact (set + # per layer in ``__init__``; off unless the switch selects the layer). + _router_exact = False + # The row-parallel tail: this rank's tail GEMM operands, built after loading (``None``: not + # applicable to this layer). + _rp_tail: Optional[_RpTailOperands] = None + # TRTLLM_K3_ULL_FRONT: the front's MNNVL workspace of the routed all-reduce group (built at + # construction) and this layer's front operands on this rank (built after loading). + _ull_ws = None + _ull_plan = None + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + """Build the routed experts and the shared expert for one MoE layer. + + ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: + the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields + that ``ModelConfig`` does not carry. + + ``aux_stream_dict`` is shared across every layer of the model, so the + streams reached through it are borrowed and must not be synchronized + or reassigned here. + """ + super().__init__() + self.layer_idx = layer_idx + self._router_exact = _router_exact_for_layer(cfg, layer_idx) + self.hidden_size = cfg.hidden_size + self.num_experts = cfg.num_experts + self.top_k = cfg.num_experts_per_token + self.moe_hidden_size = cfg.routed_expert_hidden_size + # ValueError (not assert): these guard unsupported checkpoint + # configurations and must stay active under ``python -O``. + if self.moe_hidden_size is None: + raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") + if not getattr(cfg, "latent_moe_use_norm", False): + raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") + + situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) + dtype = torch.bfloat16 + + # Routing scores stay fp32; with attention-DP off the gate GEMM runs + # bf16xbf16 with fp32 accumulate/output (checkpoint stores the gate + # weight in bf16; saves a per-layer input cast + fp32 splitK pair on + # the bs1 decode path). Under attention-DP the legacy upcast-to-fp32 + # GEMM is kept: the bf16-input min-latency GEMM's different reduction + # order flips borderline top-16 picks (GSM8K 96.7 -> 96.1/96.4, + # 3-run bisect on 62b20dd868), and the bs1-latency win is irrelevant + # at DEP batch sizes. helix_a2a ranks share requests and keep bf16. + # KIMI_K3_ROUTER_BF16=1/0 forces either path. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = ( + _router_bf16_env == "1" + if _router_bf16_env is not None + else model_config.mapping.tp_group_shares_requests + ) + self.gate = K3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) + + routed_moe_model_config = self._routed_moe_model_config(model_config) + routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) + # Resolved here so ``load_weights`` reads the checkpoint layout off the + # module instead of re-deriving it at each of its three call sites. + self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config) + self._moe_backend = routed_moe_model_config.moe_backend + self._attention_dp = model_config.mapping.enable_attention_dp + routed_mapping = routed_moe_model_config.mapping + if self._moe_backend in _K3_ROUTED_KERNEL_SETS: + if routed_moe_model_config.moe_disable_finalize_fusion: + raise NotImplementedError( + "Kimi K3 routed experts run the fused-finalize kernels " + "(moe_disable_finalize_fusion / FORCE_DETERMINISTIC is not supported)" + ) + from tensorrt_llm._torch.locality_domain.policy import LocalityDomainExecutionPlanner + + locality_plan = LocalityDomainExecutionPlanner( + routed_moe_model_config.locality_domain_policy + ).plan_moe( + routed_quant_config, + moe_backend=self._moe_backend, + use_fused_finalize=True, + dtype_activation=dtype, + activation=ActivationType.SiTu.name, + ) + if locality_plan.enabled: + raise NotImplementedError( + "Kimi K3 routed experts do not run the locality-domain MoE kernels; " + "disable enable_locality_domains" + ) + self.routed_experts = K3RoutedExperts( + kernels=self._moe_backend, + num_experts=self.num_experts, + top_k=self.top_k, + hidden_size=self.moe_hidden_size, + intermediate_size=cfg.moe_intermediate_size, + mapping=routed_mapping, + layer_idx=layer_idx, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + aux_stream=aux_stream_dict.get(AuxStreamType.MoeOutputMemset), + max_chunk_tokens=routed_moe_model_config.moe_max_num_tokens, + ) + # The routed latent partial sums, reduced over the TP group (Kimi orders this + # collective after the shared expert's). + self.routed_experts.all_reduce = ( + TpAllReduce(routed_mapping, dtype, routed_moe_model_config.allreduce_strategy) + if not routed_mapping.enable_attention_dp and routed_mapping.tp_size > 1 + else None + ) + returns_global_routed_output = False + self._routed_fused_comm = False + else: + self.routed_experts = _megamoe_routed_experts( + self.gate, + routed_moe_model_config, + routed_quant_config, + self.num_experts, + self.moe_hidden_size, + cfg.moe_intermediate_size, + dtype, + layer_idx, + aux_stream_dict, + situ_beta, + situ_linear_beta, + ) + returns_global_routed_output = self.routed_experts.returns_global_routed_output + self._routed_fused_comm = True + self.small_experts = getattr(self.routed_experts, "small_batch", None) or _NO_SMALL_BATCH + local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) + if local_expert_ids != list( + range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) + ): + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming currently requires a " + "contiguous static expert partition." + ) + self.local_expert_ids = tuple(local_expert_ids) + self.experts_per_rank = len(local_expert_ids) + self.expert_lo = local_expert_ids[0] + self.expert_hi = self.expert_lo + self.experts_per_rank + + shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts + self._shared_intermediate = shared_intermediate + attention_dp = model_config.mapping.enable_attention_dp + # Under attention DP each rank owns different tokens, so the shared + # expert is replicated (TP size 1) and must not reduce across ranks. + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. + use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1 + # Fused-comm MegaMoE runs its own cross-rank exchange inside the + # kernel. Order the shared-expert all-reduce after the routed MoE + # instead of running the two collectives concurrently (startup + # deadlocks). CuteDSL is the validated scope, not a correctness + # dependency: DeepGEMM rejects TP without attention DP. + self._defer_shared_allreduce = use_shared_tp and self._routed_fused_comm + self._reduce_routed_output = use_shared_tp and not returns_global_routed_output + if self._reduce_routed_output and self.routed_experts.all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the routed all-reduce." + ) + if model_config.use_cute_dsl_bf16_gemm: + raise NotImplementedError( + "Kimi K3's shared expert runs the cuBLAS BF16 GEMMs " + "(use_cute_dsl_bf16_gemm is not supported)" + ) + self.shared_experts = _K3SharedExpert( + cfg.hidden_size, + shared_intermediate, + activation="situ", + beta=situ_beta, + linear_beta=situ_linear_beta, + dtype=dtype, + mapping=( + subgroup_mapping(model_config.mapping, 1) if attention_dp else model_config.mapping + ), + reduce_output=use_shared_tp, + allreduce_strategy=model_config.allreduce_strategy, + ) + if _K3_SHARED_SM > 0: + shared_forward = self.shared_experts.forward + + def _shared_forward_confined(*args, **kwargs): + with _k3_shared_sm_scope(): + return shared_forward(*args, **kwargs) + + self.shared_experts.forward = _shared_forward_confined + # Side stream (+ fork/join events) for overlapping shared-expert + # compute with the routed chain. Only engaged when multi-stream is + # active (CUDA graphs on); otherwise both run in order on the default + # stream. + self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.moe_main_event = torch.cuda.Event() + self.moe_shared_event = torch.cuda.Event() + self.moe_routed_event = torch.cuda.Event() if self._defer_shared_allreduce else None + # CuteDSL already quantizes the MoE input on the MoeOutputMemset stream; + # the latent down-projection runs there too, beside the routing chain. + self.latent_down_stream = ( + aux_stream_dict.get(AuxStreamType.MoeOutputMemset) + if self._moe_backend in ("CUTEDSL", "CUTEDSL_FC12") + else None + ) + has_latent_stream = self.latent_down_stream is not None + self.moe_latent_fork_event = torch.cuda.Event() if has_latent_stream else None + self.moe_latent_ready_event = torch.cuda.Event() if has_latent_stream else None + # The shared expert's GEMMs issued ahead of the routed experts are done. + self.moe_shared_gemm_event = torch.cuda.Event() if has_latent_stream else None + self.routed_expert_down_proj = nn.Linear( + cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype + ) + # FlashInfer RMSNorm on the reduced latent (or folded into the routed all-reduce). + self.routed_expert_norm = RMSNormWeight(self.moe_hidden_size, cfg.rms_norm_eps, dtype=dtype) + if _K3_LATENT_AR_OWN_WS and self._reduce_routed_output: + mnnvl = getattr(self.routed_experts.all_reduce, "mnnvl_allreduce", None) + if mnnvl is not None and mnnvl.mapping.tp_size > 1: + self._latent_ar_ws = _k3_latent_ar_workspace( + mnnvl.mapping, mnnvl.dtype, self.moe_hidden_size, _K3_LATENT_AR_OWN_WS_MAX_M + ) + if _K3_ULL_FRONT and self._reduce_routed_output: + from .kernels import k3_ull_front + + mnnvl = getattr(self.routed_experts.all_reduce, "mnnvl_allreduce", None) + if mnnvl is not None and k3_ull_front.supports_world(mnnvl.mapping.tp_size): + # Collective over the all-reduce group: every rank builds it here. + self._ull_ws = k3_ull_front.ull_front_workspace(mnnvl.mapping) + if _K3_ULL2_COLL_CP_TAIL and self._reduce_routed_output: + # TRTLLM_K3_ULL2_COLL_CP_TAIL: the column-parallel tail's workspaces (collective). + mnnvl = getattr(self.routed_experts.all_reduce, "mnnvl_allreduce", None) + if mnnvl is not None and mnnvl.mapping.tp_size > 1: + world = mnnvl.mapping.tp_size + row_bytes = self.moe_hidden_size * world * 2 + rows = min(_K3_RP_TAIL_MAX_M, self._cp_latent_oneshot_bytes() // row_bytes) + if rows >= 1: + self._cp_ws = _k3_cp_tail_workspace( + mnnvl.mapping, self.hidden_size, _K3_RP_TAIL_MAX_M + ) + self._cp_latent_ws = _k3_cp_tail_workspace( + mnnvl.mapping, self.moe_hidden_size, rows, use="latent" + ) + self._cp_latent_push = _K3_CP_TAIL_PUSH_VIEWS[("latent", mnnvl.mapping)] + _k3_routed_backend_opt_ins(self.routed_experts, self._cp_ws is not None) + + def _lu_mid_prequant_plan(self, latent: torch.Tensor): + """``(tactic, sf_blocks)`` of the mid-M path that + latent_up's :meth:`_Fp8BlockScaleWeightReadLinear.forward` runs on the BF16 rows ``latent`` + (SM107: the K3 low-M GEMM on ``lowm_weight_scale`` above the low-M cut-over), + or ``None`` where it runs another GEMM.""" + up = self.routed_expert_up_proj + if ( + not _K3_FP8W_LOWM_GEMM + or not isinstance(up, _Fp8BlockScaleWeightReadLinear) + or up.lowm_weight_scale is None + or up.next_ar_push is not None + ): + return None + if not up._lowm_gemm_applies(latent, up.lowm_weight_scale): + return None + from .kernels import k3_fp8w_lowm_gemm + + return k3_fp8w_lowm_gemm.mid_prequant_plan(latent.shape[0], up.out_features, up.in_features) + + def _latent_up_mid_prequantized( + self, routed_latent: torch.Tensor, xq: torch.Tensor, xs: torch.Tensor, tactic + ) -> torch.Tensor: + """latent_up on the MXFP8 form ``(xq, xs)`` of ``routed_latent`` in the mid-M GEMM's + layout: what :meth:`_routed_projection` computes, bit for + bit, without its quantization kernel.""" + up = self.routed_expert_up_proj + m = xq.shape[0] + out = up._take_next_output(m) + if out is None: + out = torch.empty((m, up.out_features), dtype=torch.bfloat16, device=xq.device) + k3_fp8w_mid_gemm_op(xq, xs, up.weight, up.lowm_weight_scale, out, tactic) + return out.reshape(routed_latent.shape[:-1] + (up.out_features,)) + + def _router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """FP32 router logits; ``TRTLLM_K3_ROUTER_EXACT`` serves ``M <= 16`` BF16 tokens with + ``k3_router_exact`` (bit-identical to the ``dsv3_router_gemm_op`` it replaces) in the + layers it selects.""" + if self._router_exact: + from .kernels import k3_router_exact + + x = hidden_states.reshape(-1, self.gate.gating_dim) + if k3_router_exact.can_run(x, self.gate.weight): + return k3_router_logits_op(x, self.gate.weight) + return self.gate.compute_logits(hidden_states) + + def _shared_gemms_can_lead(self) -> bool: + """Whether this forward can issue the shared expert's GEMMs ahead of the routed experts + (``_ull_fork_shared``) and its all-reduce separately.""" + return ( + self.shared_expert_stream is not None + and self.moe_shared_gemm_event is not None + and not self._defer_shared_allreduce + and not getattr(self.shared_experts.down_proj, "use_fused_gemm_allreduce", False) + and do_multi_stream() + ) + + def _shared_needs_all_reduce(self) -> bool: + """Whether the shared expert's output is a row-parallel partial to all-reduce.""" + return bool(getattr(self.shared_experts.down_proj, "reduce_output", False)) + + def _shared_all_reduce(self) -> TpAllReduce: + """The shared expert's row-parallel reduction.""" + return self.shared_experts.down_proj.all_reduce + + def post_load_weights(self) -> None: + self.small_experts.post_load_weights() + self.prepare_rp_tail() + self.prepare_ull_front() + + def cache_derived_state(self) -> None: + self.small_experts.post_load_weights() + self.prepare_rp_tail() + self.prepare_ull_front() + + def prepare_expert_load(self) -> None: + """Before this layer's first routed expert streams in.""" + backend = self.routed_experts.backend + if isinstance(backend, K3Nvfp4Experts): + backend.prepare_streaming_load() + else: + backend.quant_method.prepare_streaming_expert_load(backend) + + def finalize_expert_load(self) -> None: + """After this layer's last routed expert streamed in.""" + backend = self.routed_experts.backend + if isinstance(backend, K3Nvfp4Experts): + backend.finalize() + else: + backend.process_weights_after_loading() + + def _ull_front_unsupported(self) -> Optional[str]: + """Why this layer cannot run the sharded MoE front (``TRTLLM_K3_ULL_FRONT``), or + ``None``.""" + if self._ull_ws is None: + return "no MNNVL routed all-reduce over a supported TP group" + moe = self.routed_experts + if not isinstance(moe, K3RoutedExperts) or self._attention_dp: + return "the routed experts are not the CuTe DSL kernels of a TP group" + if moe.backend.num_slots != self.gate.num_experts: + return "the routed experts have no static contiguous expert slice" + down = self.routed_expert_down_proj + if ( + not isinstance(down, _Fp8BlockScaleWeightReadLinear) + or down._fused_projection is not None + ): + return "latent_down is not an FP8 weight-read projection" + if self.gate.weight.dtype != torch.bfloat16: + return "the router weight is not BF16" + return None + + def prepare_ull_front(self) -> bool: + """``TRTLLM_K3_ULL_FRONT``: build this layer's front operands on this rank (views of the + router and latent_down weights, the packed scale of the rank's latent_down tiles). Rebuilt + on every call (post-load), before CUDA graph capture.""" + self._ull_plan = None + if not _K3_ULL_FRONT: + return False + reason = self._ull_front_unsupported() + if reason is not None: + logger.info_once( + f"TRTLLM_K3_ULL_FRONT: MoE layers keep the replicated front: {reason}", + key=f"k3_ull_front_off_{reason}", + ) + return False + from .kernels import k3_ull_front + + down = self.routed_expert_down_proj + down.transform_weights() + backend = self.routed_experts.backend + ws = self._ull_ws + with torch.no_grad(): + self._ull_plan = k3_ull_front.make_plan( + down.weight, + down.weight_scale, + self.gate.weight, + self.gate.e_score_correction_bias, + float(self.gate.routed_scaling_factor), + backend.fc31_input_scale, + ws.world, + ws.rank, + backend.slot_start, + backend.expert_size_per_partition, + ) + logger.info_once( + f"TRTLLM_K3_ULL_FRONT: sharded MoE front over {ws.world} ranks, routed experts " + f"[{backend.slot_start}, {backend.slot_start + backend.expert_size_per_partition}) " + f"on rank {ws.rank}", + key="k3_ull_front_on", + ) + return True + + def _ull_front_applies(self, hidden_states: torch.Tensor) -> bool: + """Whether this forward runs the sharded MoE front.""" + plan = self._ull_plan + if plan is None: + return False + if not 1 <= hidden_states.shape[0] <= min(_K3_ULL_FRONT_MAX_M, 32): + return False + from .kernels import k3_ull_front + + x = hidden_states.reshape(-1, self.hidden_size) + if not k3_ull_front.can_run(x, plan): + return False + # The kernels compile on their first call, which must not fall into a CUDA graph capture. + return not torch.cuda.is_current_stream_capturing() or k3_ull_front.is_compiled( + x.shape[0], plan, self._ull_publish_word(x.shape[0]) is not None + ) + + def _ull_publish_word(self, m: int) -> Optional[torch.Tensor]: + """The routed experts' routing word into which the front of an ``m``-row forward publishes + the top-k (``K3SmallBatchExperts.published_routing_word``), or ``None`` where the front + writes the ``moe_sort`` tables.""" + if not getattr(self.routed_experts, "consumes_sort_tables", False): + return None + return self.small_experts.published_routing_word(m) + + def ull_front_releases_early(self, like: torch.Tensor) -> bool: + """Whether the mix that produces this forward's input may release the front's head at + kernel entry, i.e. whether the next forward on rows shaped like ``like`` (a contiguous BF16 + ``[M, hidden]`` tensor, e.g. the prefix sum the mix updates) runs the sharded front, whose + head is then the next launch on the stream. Conservative: ``False`` where ``like`` cannot + stand for the mix's output.""" + if like.dim() != 2 or not like.is_contiguous(): + return False + return self._ull_front_applies(like) + + def _ull_front(self, hidden_states: torch.Tensor, after_head=None): + """Run the sharded MoE front: returns ``(router_logits placeholder, NVFP4 routed input, + routing)``, ``routing`` the ``SortedRouting`` the routed experts take (with this rank's + ``moe_sort`` tables, the published top-k, or no tables); ``after_head`` runs between the + front's two launches.""" + from .kernels import k3_ull_front + + plan = self._ull_plan + x = hidden_states.reshape(-1, self.hidden_size) + if not x.is_contiguous(): + x = x.contiguous() + moe = self.routed_experts + backend = moe.backend + # The global scale can be replaced after the plan was built (calibration, reload). + plan.global_sf = backend.fc31_input_scale.reshape(-1)[:1] + # Where the routed experts offer a routing word, the route kernel publishes each token's + # top-k into it instead of writing the moe_sort tables; the experts build their tables. + word = self._ull_publish_word(x.shape[0]) + routing, codes, scales = k3_ull_front_op( + x, + plan, + self._ull_ws, + after_head=after_head, + publish=word.data_ptr() if word is not None else 0, + ) + if word is not None: + sorted_routing = SortedRouting(routing[0], routing[1], PublishedSortTables()) + elif moe.consumes_sort_tables and self.small_experts.applies(x.shape[0]): + sorted_routing = SortedRouting( + routing[0], routing[1], SortTables(k3_ull_front.TILE, tuple(routing[2:8])) + ) + else: + sorted_routing = SortedRouting(routing[0], routing[1], None) + # The routing does not read the logits; they only carry the shape. + logits = torch.empty( + (x.shape[0], self.gate.num_experts), dtype=torch.float32, device=x.device + ) + return logits, Fp4QuantizedTensor(codes, scales, is_sf_swizzled=False), sorted_routing + + def _ull_fork_shared( + self, identity: torch.Tensor, persist: bool, rows: int + ) -> Tuple[Optional[Callable[[], None]], Optional[Callable[[], None]], List[torch.Tensor]]: + """Fork the shared expert's GEMMs for a ULL front on ``rows`` rows where ``_ull_shared_fork`` puts it. + + Returns ``(after_head, after_front, forked)``: the callback the front runs between its head and route + launches, the one the caller runs once the front is issued (``None``: nothing to run) and the list that + holds the shared partial once the GEMMs are issued. + + The shared GEMMs are captured after the whole front, whichever kernel they depend on: each + early-releasing kernel's first dependent in the CUDA graph then is the one its trigger is meant for + (A_mid's the head, the head's the route kernel), and the executor keeps that one on the releasing + kernel's stream. A shared gate_up captured before it would take that place whenever no + other edge pins it to the shared stream, and the intended dependent would wait for the releasing kernel + to finish. + """ + forked: List[torch.Tensor] = [] + + def run_shared(): + forked.append(self._ull_shared_first(identity, persist)) + + fork = _ull_shared_fork(self._shared_gemms_can_lead(), rows) + if fork == "head": + return self.moe_main_event.record, run_shared, forked + if fork == "after_head": + self.moe_main_event.record() + return None, run_shared, forked + return None, None, forked + + def _ull_shared_first(self, identity: torch.Tensor, persist: bool = False) -> torch.Tensor: + """The shared expert's GEMMs on the shared stream from ``moe_main_event`` (recorded around + the sharded front's head, ``_ull_shared_fork``), without its all-reduce; + ``moe_shared_gemm_event`` marks them done. ``persist`` (the column-parallel tail): the + persistent kernel instead, which pushes the partial itself (returns ``_SHARED_PUSHED``).""" + with torch.cuda.stream(self.shared_expert_stream): + self.moe_main_event.wait() + if persist: + self._run_shared_persist(identity) + partial = _SHARED_PUSHED + else: + partial = self.shared_experts( + identity, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) + self.moe_shared_gemm_event.record() + return partial + + def _rp_tail_unsupported(self) -> Optional[str]: + """Why this layer cannot run the row-parallel MoE tail, or ``None``.""" + if not self._reduce_routed_output or self._defer_shared_allreduce: + return "the routed latent is not all-reduced by the model" + moe_all_reduce = self.routed_experts.all_reduce + if getattr(moe_all_reduce, "mnnvl_allreduce", None) is None: + return "the routed all-reduce is not the MNNVL all-reduce" + up = self.routed_expert_up_proj + if not isinstance(up, _Fp8BlockScaleWeightReadLinear) or up._fused_projection is not None: + return ( + "latent_up is not an FP8 weight-read projection (KIMI_K3_FP8_WEIGHT_READ_MOE_MLP)" + ) + mlp = self.shared_experts + if not isinstance(mlp, _K3SharedExpert) or not isinstance(mlp.down_proj, Bf16Linear): + return "the shared expert does not run the BF16 gate_up -> activation -> down path" + if not self._shared_needs_all_reduce(): + return "the shared expert is not row-parallel" + down = mlp.down_proj + if getattr(down, "bias", None) is not None: + return "the shared down projection has a bias" + world = moe_all_reduce.mapping.tp_size + shard = down.weight.shape[1] + if shard * world != self._shared_intermediate: + return f"the shared down shard ({shard}) is not a 1/{world} shard" + if ( + tuple(up.weight.shape) != (self.hidden_size, self.moe_hidden_size) + or self.moe_hidden_size % 128 + or self.hidden_size % 128 + or self.moe_hidden_size // 128 < world + ): + return "the latent is not split into whole 128-column blocks over the ranks" + if tuple(down.weight.shape) != (self.hidden_size, shard): + return "unexpected shared down weight shape" + return None + + def prepare_rp_tail(self) -> bool: + """Build this rank's row-parallel tail operands (``_RpTailOperands``): the packed UE8M0 + scale of this rank's latent_up blocks (whose FP8 codes stay where they are). Rebuilt on + every call (post-load), before CUDA graph capture.""" + self._rp_tail = None + reason = self._rp_tail_unsupported() + if reason is not None: + logger.info_once( + f"Kimi K3 MoE layers keep the replicated tail (no row-parallel tail): {reason}", + key=f"k3_rp_tail_off_{reason}", + ) + return False + from .kernels import k3_fp8w_lowm_gemm + + up = self.routed_expert_up_proj + up.transform_weights() + down = self.shared_experts.down_proj + mapping = self.routed_experts.all_reduce.mapping + k0, k1 = _rp_tail_latent_slice(self.moe_hidden_size, mapping.tp_size, mapping.tp_rank) + hidden, shard = down.weight.shape + with torch.no_grad(): + exponents = k3_fp8w_lowm_gemm.block_exponents_of( + up.weight_scale, hidden, self.moe_hidden_size + )[:, k0 // 128 : (k0 + k1) // 128] + scale = k3_fp8w_lowm_gemm.pack_ue8m0_exponents(exponents.contiguous(), hidden, k1) + self._rp_tail = _RpTailOperands(k0, k1, scale) + self._cp_tail = None + if self._cp_ws is not None: + # TRTLLM_K3_ULL2_COLL_CP_TAIL: this rank's whole 128-row blocks of latent_up's output. + c0, n = _rp_tail_latent_slice(hidden, mapping.tp_size, mapping.tp_rank) + with torch.no_grad(): + cp_exps = k3_fp8w_lowm_gemm.block_exponents_of( + up.weight_scale, hidden, self.moe_hidden_size + )[c0 // 128 : (c0 + n) // 128] + cp_scale = k3_fp8w_lowm_gemm.pack_ue8m0_exponents( + cp_exps.contiguous(), n, self.moe_hidden_size + ) + self._cp_tail = _CpTailOperands(c0, n, up.weight[c0 : c0 + n], cp_scale) + logger.info_once( + f"TRTLLM_K3_ULL2_COLL_CP_TAIL: column-parallel MoE tail, latent_up rows " + f"[{c0}, {c0 + n}) (rank {mapping.tp_rank} of {mapping.tp_size})", + key=f"k3_cp_tail_on_{n}", + ) + self._prepare_tail1(up, hidden) + self._prepare_shared_persist() + logger.info_once( + f"Kimi K3 row-parallel MoE tail, K = latent [{k0}, {k0 + k1}), BF16 shared " + f"partial added (rank {mapping.tp_rank} of {mapping.tp_size})", + key=f"k3_rp_tail_on_{k1}_{shard}", + ) + return True + + def _prepare_tail1(self, up, hidden: int) -> None: + """latent_up's block exponents for k3_moe_tail1 (the FP8 codes are read in place).""" + self._tail1_wexp = None + if self._cp_latent_push is None: + return + from .kernels import k3_fp8w_lowm_gemm, k3_moe_tail1 + + weight = up.weight + if ( + weight.dtype != torch.float8_e4m3fn + or tuple(weight.shape) != (k3_moe_tail1.HIDDEN, k3_moe_tail1.LATENT) + or hidden != k3_moe_tail1.HIDDEN + or self.moe_hidden_size != k3_moe_tail1.LATENT + or not weight.is_contiguous() + ): + logger.info_once( + "k3_moe_tail1: latent_up is not the [7168, 3584] FP8 weight, tail unchanged", + key="k3_tail1_off_weight", + ) + return + with torch.no_grad(): + wexp = k3_fp8w_lowm_gemm.block_exponents_of( + up.weight_scale, hidden, self.moe_hidden_size + ).contiguous() + if int(wexp.min().item()) < 1 or int(wexp.max().item()) > 254: + logger.info_once( + "k3_moe_tail1: latent_up block scale outside the kernel's range, tail unchanged", + key="k3_tail1_off_scale", + ) + return + self._tail1_wexp = wexp + world, rank = self._cp_latent_ws.world, self._cp_latent_ws.rank + c0, n = k3_moe_tail1.rank_rows(world, rank) + logger.info_once( + f"k3_moe_tail1: one-kernel MoE tail up to {k3_moe_tail1.MAX_M} rows, latent_up " + f"rows [{c0}, {c0 + n}) (rank {rank} of {world})", + key=f"k3_tail1_on_{n}", + ) + + def _prepare_shared_persist(self) -> None: + """This rank's shared expert weights in k3_shared_persist's stream order and its + workspace.""" + self._persist_frags = None + if self._cp_ws is None: + return + from .kernels import k3_shared_persist + + ws = self._cp_ws + mlp = self.shared_experts + sms = _K3_ULL_SHARED_PERSIST_CTAS + gate_up = getattr(getattr(mlp, "gate_up_proj", None), "weight", None) + down = getattr(getattr(mlp, "down_proj", None), "weight", None) + inter = k3_shared_persist.INTER // ws.world + reason = None + if getattr(mlp, "activation", None) != "situ": + reason = "the shared expert's activation is not SiTU" + elif ( + gate_up is None + or down is None + or gate_up.dtype != torch.bfloat16 + or down.dtype != torch.bfloat16 + or tuple(gate_up.shape) != (2 * inter, k3_shared_persist.HIDDEN) + or tuple(down.shape) != (k3_shared_persist.HIDDEN, inter) + ): + reason = "the shared expert is not the BF16 [2 I, 7168] / [7168, I] shard" + elif mlp.down_proj.tp_size != ws.world or mlp.down_proj.tp_rank != ws.rank: + reason = "the shared expert's TP group is not the tail's" + elif not k3_shared_persist.can_run(k3_shared_persist.MAX_M, ws.world, sms): + reason = f"{sms} CTAs do not split the {inter}-column shard" + if reason is not None: + logger.info_once( + f"k3_shared_persist: shared expert unchanged: {reason}", + key=f"k3_persist_off_{reason}", + ) + return + with torch.no_grad(): + self._persist_frags = k3_shared_persist.persist_fragments( + gate_up.data.contiguous(), down.data.contiguous(), sms + ) + self._persist_ws = k3_shared_persist.PersistWorkspace(ws.world, gate_up.device) + self._persist_act = (float(mlp.beta), mlp.linear_beta) + logger.info_once( + f"k3_shared_persist: the shared expert on {sms} CTAs beside the routed " + f"experts up to {k3_shared_persist.MAX_M} rows (rank {ws.rank} of {ws.world})", + key=f"k3_persist_on_{sms}", + ) + + def _persist_applies(self, num_tokens: int, target) -> bool: + """Whether this forward's shared expert is the persistent kernel (the column-parallel + tail).""" + if ( + self._persist_frags is None + or target is None + or self.shared_expert_stream is None + or not do_multi_stream() + or not _K3_ULL2_COLL_CP_TAIL + or not self._cp_tail_applies(num_tokens, target) + ): + return False + from .kernels import k3_shared_persist + + return k3_shared_persist.can_run(num_tokens, self._cp_ws.world, _K3_ULL_SHARED_PERSIST_CTAS) + + def _run_shared_persist(self, identity: torch.Tensor) -> None: + """The persistent shared expert of ``identity`` on the current (shared) stream, pushing + into the tail call's slots.""" + ws = self._cp_ws + x = identity if identity.is_contiguous() else identity.contiguous() + beta, linear_beta = self._persist_act + k3_shared_persist_op( + x=x, + frags=self._persist_frags, + world=ws.world, + rank=ws.rank, + sms=_K3_ULL_SHARED_PERSIST_CTAS, + ws=self._persist_ws, + out_base=ws.mc_base, + out_flags=ws.flags, + beta=beta, + linear_beta=linear_beta, + ) + + def _rp_tail_target( + self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor], rp_tail_ok: bool + ): + """The all-reduce push target of this forward's row-parallel MoE tail, or ``None`` where + the forward keeps the replicated tail: a generation-only batch of at most + ``_K3_RP_TAIL_MAX_M`` BF16 tokens with the fused residual update, served by the low-M + GEMM's push epilogue. The consumer finishes the call: one-shot transport up to + lamport._K3_FUSED_FINISH_ONE_SHOT_MAX_BYTES, two-shot above.""" + rp = self._rp_tail + num_tokens = hidden_states.shape[0] + if ( + rp is None + or not rp_tail_ok + or residual is None + or hidden_states.dtype != torch.bfloat16 + or not 1 <= num_tokens <= _K3_RP_TAIL_MAX_M + ): + return None + from .kernels import k3_fp8w_lowm_gemm + + if k3_fp8w_lowm_gemm.split_source_tactic(num_tokens, self.hidden_size, rp.k1) is None: + return None + return lamport_push_target( + self.routed_experts.all_reduce, + num_tokens, + self.hidden_size, + deferred=True, + ) + + def _rp_shared_output(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The shared expert's part of the row-parallel tail on this rank: its row-parallel BF16 + down-projection partial without the all-reduce.""" + return self.shared_experts( + hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) + + def _rp_tail_mx_applies(self, num_tokens: int) -> bool: + """Whether the row-parallel tail GEMM can run as the MXFP8 low-M GEMM on the latent + all-reduce's MXFP8 output (``TRTLLM_K3_ULL_TAIL_SANDWICH``; the tail GEMM has one source): + the same tactic as the split-source BF16 GEMM, so the same bits.""" + from .kernels import k3_fp8w_lowm_mxfp8 + + if num_tokens > k3_fp8w_lowm_mxfp8.MAX_M: + return False + try: + k3_fp8w_lowm_mxfp8.default_mx_tactic(num_tokens, self.hidden_size, self._rp_tail.k1) + except ValueError: + return False + return True + + def _rp_tail_latent_ar_early( + self, routed_out: torch.Tensor, target, mxfp8: bool = False + ) -> Optional[Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]]: + """``TRTLLM_K3_ULL_TAIL_SANDWICH``: the routed latent all-reduce + RMSNorm of the + row-parallel tail by ``k3_latent_ar_norm`` (its row kernel where it serves the size, else + the cluster kernel) on the MNNVL workspace the tail GEMM pushes into: bit-identical to the + C++ one-shot kernel, but it releases the tail GEMM right after its own + ``griddepcontrol.wait``, so the GEMM streams its weight slice while the all-reduce waits + for the other ranks. Returns ``(latent, (xq, xs) with mxfp8 else None)``, or ``None`` + where it does not apply (the C++ kernel's two-shot sizes). ``mxfp8``: the kernel also + stores the MXFP8 form of the latent, bitwise the low-M GEMM's own quantization of it.""" + from tensorrt_llm._torch.distributed.ops import _MNNVL_ONE_SHOT_THRESHOLD_BYTES + + from .kernels import k3_latent_ar_norm + + ws = target.workspace + latent = routed_out.reshape(-1, self.moe_hidden_size) + if not latent.is_contiguous(): + latent = latent.contiguous() + gamma = self.routed_expert_norm.weight + if latent.numel() * ws.world * latent.element_size() > _MNNVL_ONE_SHOT_THRESHOLD_BYTES: + return None + row = k3_latent_ar_norm.can_run_row(latent, gamma, ws.world, ws.bytes_per_buffer) + if not (row or k3_latent_ar_norm.can_run(latent, gamma, ws.world)): + return None + out = torch.empty_like(latent) + m, k = latent.shape + latent_mx = None + if mxfp8: + latent_mx = ( + torch.empty((m, k), dtype=torch.float8_e4m3fn, device=latent.device), + torch.empty((m, k // 32), dtype=torch.uint8, device=latent.device), + ) + k3_latent_ar_rmsnorm( + latent, + gamma, + float(self.routed_expert_norm.variance_epsilon), + out, + uc_addr=ws.uc_local, + mc_addr=ws.mc_base, + flags=ws.flags, + rank=ws.rank, + world_size=ws.world, + mxfp8_out=latent_mx, + row=row, + ) + return out, latent_mx + + @staticmethod + def _cp_latent_oneshot_bytes() -> int: + """Largest one-shot latent all-reduce of the column-parallel tail (bytes of slots): the C++ + one-shot size.""" + from tensorrt_llm._torch.distributed.ops import _MNNVL_ONE_SHOT_THRESHOLD_BYTES + + return _MNNVL_ONE_SHOT_THRESHOLD_BYTES + + def _tail1_applies(self, num_tokens: int) -> bool: + """Whether this forward's MoE tail runs k3_moe_tail1: its operands are built, the routed + experts push their latent rows at ``num_tokens`` (``K3SmallBatchExperts.applies``: the + small-batch kernel, no communication strategy) and the kernel serves the shape.""" + if ( + self._tail1_wexp is None + or self._cp_latent_push is None + or not self.small_experts.applies(num_tokens) + ): + return False + from .kernels import k3_moe_tail1 + + return k3_moe_tail1.can_run(num_tokens, self._cp_latent_push.world) + + def _cp_tail_applies(self, num_tokens: int, target) -> bool: + """Whether this forward's MoE tail runs column-parallel (``TRTLLM_K3_ULL2_COLL_CP_TAIL``): + the latent all-reduce runs one-shot by ``k3_latent_ar_norm``'s row kernel on the tail's + latent workspace (its MXFP8 output feeds the tail GEMM), the replicated projection's MXFP8 + tactic serves ``num_tokens`` and the tail's workspaces hold them.""" + cp = self._cp_tail + lws = self._cp_latent_ws + if cp is None or lws is None or num_tokens > _K3_RP_TAIL_MAX_M: + return False + from .kernels import k3_fp8w_lowm_mxfp8, k3_latent_ar_norm + + if ( + lws.world not in k3_latent_ar_norm.ROW_WORLD_SIZES + or num_tokens > k3_latent_ar_norm.ROW_MAX_M + or self.moe_hidden_size != k3_latent_ar_norm.HIDDEN + ): + return False + nbytes = num_tokens * self.moe_hidden_size * lws.world * 2 + if nbytes > self._cp_latent_oneshot_bytes() or nbytes > lws.bytes_per_buffer: + return False + if num_tokens > k3_fp8w_lowm_mxfp8.MAX_M: + return False + try: + k3_fp8w_lowm_mxfp8.default_mx_tactic(num_tokens, self.hidden_size, self.moe_hidden_size) + except ValueError: + return False + return True + + def _cp_tail_forward( + self, + routed_output: Callable[..., torch.Tensor], + identity: torch.Tensor, + shared_part: Optional[torch.Tensor], + residual: torch.Tensor, + target, + defer_residual: bool, + ): + """``TRTLLM_K3_ULL2_COLL_CP_TAIL``: the routed experts, the shared expert's BF16 partial + pushed into the tail's call on the shared stream, the latent all-reduce + RMSNorm (MXFP8 + output), then this rank's latent_up output columns multicast into the call's all-gather + rows; or, where k3_moe_tail1 applies, the routed experts push their latent rows + (``routed_output(output_push=...)``) and one kernel does the rest. Returns a + ``_DeferredRpTail`` whose pending call the next pre-attention mix finishes (``finish_ag``), + or the updated residual.""" + from .kernels.k3_fp8w_lowm_mxfp8 import default_mx_tactic + + cp = self._cp_tail + cp_ws = self._cp_ws + num_tokens = identity.shape[0] + tail1 = self._tail1_applies(num_tokens) + if tail1: + routed_output = functools.partial(routed_output, output_push=self._cp_latent_push) + + def push_shared(part: torch.Tensor) -> None: + part = part.reshape(num_tokens, -1) + if not part.is_contiguous(): + part = part.contiguous() + k3_lamport_push(cp_ws, part, True) + + side = self.shared_expert_stream is not None and do_multi_stream() + if shared_part is not None: + # The shared expert ran ahead on the shared stream (moe_shared_gemm_event); the + # persistent kernel pushed its partial itself. + routed_out = routed_output() + if side: + with torch.cuda.stream(self.shared_expert_stream): + self.moe_shared_gemm_event.wait() + if shared_part is not _SHARED_PUSHED: + push_shared(shared_part) + self.moe_shared_event.record() + else: + self.moe_shared_gemm_event.wait() + push_shared(shared_part) + elif side: + self.moe_main_event.record() + persist = self._persist_applies(num_tokens, target) + routed_out = routed_output() + with torch.cuda.stream(self.shared_expert_stream): + self.moe_main_event.wait() + if persist: + self._run_shared_persist(identity) + else: + push_shared(self._rp_shared_output(identity)) + self.moe_shared_event.record() + else: + routed_out = routed_output() + push_shared(self._rp_shared_output(identity)) + if tail1: + # The whole tail in one kernel (the experts pushed their latent rows), into the call's + # all-gather rows. + from .kernels import k3_moe_tail1 + + lws = self._cp_latent_ws + c0, _ = k3_moe_tail1.rank_rows(lws.world, lws.rank) + k3_moe_tail1_op( + m=num_tokens, + world=lws.world, + lat_uc=lws.uc_local, + lat_flags=lws.flags, + weight=self.routed_expert_up_proj.weight, + wexp=self._tail1_wexp, + gamma=self.routed_expert_norm.weight, + eps=float(self.routed_expert_norm.variance_epsilon), + ag_base=cp_ws.mc_base, + ag_flags=cp_ws.flags, + c0=c0, + ) + return self._cp_tail_finish(cp_ws, num_tokens, residual, side, defer_residual) + # The latent all-reduce + RMSNorm on the tail's latent workspace (one-shot, checked by + # _cp_tail_applies), releasing the tail GEMM at once, with its MXFP8 output. + latent = routed_out.reshape(-1, self.moe_hidden_size) + if not latent.is_contiguous(): + latent = latent.contiguous() + ws = self._cp_latent_ws + latent_mx = ( + torch.empty(latent.shape, dtype=torch.float8_e4m3fn, device=latent.device), + torch.empty( + (num_tokens, self.moe_hidden_size // 32), dtype=torch.uint8, device=latent.device + ), + ) + k3_latent_ar_rmsnorm( + latent, + self.routed_expert_norm.weight, + float(self.routed_expert_norm.variance_epsilon), + torch.empty_like(latent), + uc_addr=ws.uc_local, + mc_addr=ws.mc_base, + flags=ws.flags, + rank=ws.rank, + world_size=ws.world, + mxfp8_out=latent_mx, + early_state=True, + ) + # The replicated projection's tactic: every 128-row output tile is computed as there. + k3_fp8w_lowm_mxfp8_push_op( + *latent_mx, + cp.weight, + cp.scale, + cp_ws, + col0=cp.c0, + hidden=self.hidden_size, + tactic=default_mx_tactic(num_tokens, self.hidden_size, self.moe_hidden_size), + ) + return self._cp_tail_finish(cp_ws, num_tokens, residual, side, defer_residual) + + def _cp_tail_finish( + self, cp_ws, num_tokens: int, residual: torch.Tensor, side: bool, defer_residual: bool + ): + """The column-parallel tail call (all-gather rows pushed, the shared expert's slots pushed + on the shared stream, joined here) as a ``_DeferredRpTail`` the next pre-attention mix + finishes, or the updated residual.""" + from tensorrt_llm._torch.distributed.ops import _MNNVL_ONE_SHOT_THRESHOLD_BYTES + + from .layers import LamportPushTarget + + if side: + self.moe_shared_event.wait() + # The call reproduces the shared expert's C++ all-reduce (its transport by size). + cpp_one_shot = ( + num_tokens * self.hidden_size * cp_ws.world * 2 <= _MNNVL_ONE_SHOT_THRESHOLD_BYTES + ) + cp_target = LamportPushTarget( + cp_ws, True, cpp_one_shot, torch.empty_like(residual.reshape(num_tokens, -1)) + ) + cp_target.pushed = True + cp_target.deferred = True + cp_target.ag = True + deferred = _DeferredRpTail(residual.reshape(num_tokens, -1), LamportPending(cp_target)) + if defer_residual: + return deferred + return deferred.materialize().view(residual.shape) + + def _rp_tail_forward( + self, + routed_output: Callable[..., torch.Tensor], + identity: torch.Tensor, + shared_part: Optional[torch.Tensor], + residual: torch.Tensor, + target, + defer_residual: bool, + ): + """The row-parallel tail: the routed experts (``routed_output``) and the shared expert's + part (``_rp_shared_output``; ``shared_part`` if it already ran ahead of the routed experts), + then the latent all-reduce + RMSNorm, then this rank's tail GEMM pushing its partial into + ``target``. Returns a ``_DeferredRpTail`` (``defer_residual``) or the updated residual.""" + if _K3_ULL2_COLL_CP_TAIL and self._cp_tail_applies( + residual.numel() // self.hidden_size, target + ): + return self._cp_tail_forward( + routed_output, identity, shared_part, residual, target, defer_residual + ) + if shared_part is not None: + # The shared expert ran ahead on the shared stream and the routed experts waited for + # it (moe_shared_gemm_event). + routed_out = routed_output() + self.moe_shared_gemm_event.wait() + elif self.shared_expert_stream is not None and do_multi_stream(): + self.moe_main_event.record() + routed_out = routed_output() + with torch.cuda.stream(self.shared_expert_stream): + self.moe_main_event.wait() + shared_part = self._rp_shared_output(identity) + self.moe_shared_event.record() + self.moe_shared_event.wait() + else: + routed_out = routed_output() + shared_part = self._rp_shared_output(identity) + moe_all_reduce = self.routed_experts.all_reduce + routed_latent = None + latent_mx = None + if _K3_ULL_TAIL_SANDWICH: + num_rows = routed_out.numel() // self.moe_hidden_size + early = self._rp_tail_latent_ar_early( + routed_out, target, mxfp8=self._rp_tail_mx_applies(num_rows) + ) + if early is not None: + routed_latent, latent_mx = early + if routed_latent is None: + routed_latent = moe_all_reduce( + routed_out, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=self.routed_expert_norm.weight, + eps=float(self.routed_expert_norm.variance_epsilon), + ), + ) + latent = routed_latent.reshape(-1, self.moe_hidden_size) + num_tokens = latent.shape[0] + shared_part = shared_part.reshape(num_tokens, -1) + if not shared_part.is_contiguous(): + shared_part = shared_part.contiguous() + rp = self._rp_tail + lo, hi = rp.k0, rp.k0 + rp.k1 + if latent_mx is not None: + # TRTLLM_K3_ULL_TAIL_SANDWICH: the MXFP8 GEMM on the all-reduce's MXFP8 output (its + # slice read in place), no quantization in the tail GEMM; the same bits. + k3_fp8w_lowm_mxfp8_push_op( + *latent_mx, + self.routed_expert_up_proj.weight, + rp.scale, + target.workspace, + target.one_shot, + k0=lo, + k1=rp.k1, + addend=shared_part, + ) + else: + k3_fp8w_lowm_gemm_push_op( + latent[:, lo:hi], + self.routed_expert_up_proj.weight[:, lo:hi], + rp.scale, + target.workspace, + target.one_shot, + addend=shared_part, + strided=True, + ) + target.pushed = True + result = finish_lamport_push(moe_all_reduce, target, deferred=True) + deferred = _DeferredRpTail(residual.reshape(num_tokens, -1), result) + if defer_residual: + return deferred + return deferred.materialize().view(residual.shape) + + def _route_tail_target(self, num_tokens: int) -> Optional[Tuple[int, int]]: + """``(local_expert_offset, local_num_experts)`` if the CuTe routing tail serves this + forward, else ``None``.""" + if not isinstance(self.routed_experts, K3RoutedExperts) or self._attention_dp: + return None + backend = self.routed_experts.backend + slot_start = backend.slot_start + num_local = backend.expert_size_per_partition + if backend.num_slots != self.gate.num_experts: + return None + from .kernels import k3_route_tail + + device = self.gate.e_score_correction_bias.device + if not k3_route_tail.supports_device(device): + return None + if not 1 <= num_tokens <= min(_K3_ROUTE_TAIL_MAX_M, k3_route_tail.WIDE_MAX_M): + return None + if num_tokens > k3_route_tail.MAX_M: + if not k3_route_tail.wide_sort_fits(num_tokens, num_local): + return None + # The multi-CTA kernel's workspace must exist before graph capture. + if not torch.cuda.is_current_stream_capturing(): + k3_route_tail.wide_workspace(device) + elif not k3_route_tail.wide_workspace_ready(device): + return None + return slot_start, num_local + + def _router(self, tail_target: Optional[Tuple[int, int]]): + """The routing of this forward's router logits, for the experts + ``[local_expert_offset, local_expert_offset + local_num_experts)`` the routed experts ask + for: the CuTe routing tail where ``tail_target`` is set; else ``noaux_tc_route_sort`` + where it serves; else ``noaux_tc_op`` (the experts then run ``moe_sort``).""" + gate = self.gate + consumes_tables = self.routed_experts.consumes_sort_tables + + def route(logits: torch.Tensor, local_expert_offset: int, local_num_experts: int): + target = (local_expert_offset, local_num_experts) + if tail_target is not None: + routed = gate.route_tail(logits, *tail_target) + if consumes_tables and target == tail_target: + return routed + return SortedRouting(routed.token_selected_experts, routed.token_final_scales, None) + if consumes_tables: + routed = gate.route_with_sort(logits, *target) + if routed is not None: + return routed + return SortedRouting(*gate.route(logits), None) + + return route + + @staticmethod + def _fused_projection_applies(x: torch.Tensor, projection: nn.Module) -> bool: + """Whether the fused latent projection replaces ``projection(x)``.""" + if not _K3_LD_FUSED: + return False + if not isinstance(projection, _Fp8BlockScaleWeightReadLinear): + return False + packed_scale = projection.packed_weight_scale + if packed_scale is None or x.dtype != torch.bfloat16: + return False + if not _K3_LD_FUSED_MIN_M <= x.shape[0] <= _K3_LD_FUSED_MAX_M: + return False + from .kernels import k3_latent_fp8_gemm + + # Where the fused op itself runs the low-M GEMM (SM107), it only adds the epilogue. + if x.shape[0] <= k3_latent_fp8_gemm.lowm_route_max_m(*projection.weight.shape): + return True + if _K3_FP8W_LOWM_GEMM and _k3_lowm_gemm_can_run(x, projection.weight, packed_scale): + return False + return True + + def _fused_latent_input(self, hidden_states: torch.Tensor) -> Optional[Fp4QuantizedTensor]: + """NVFP4 routed-expert input straight from the hidden states, or ``None``. + + Runs the latent_down projection with the experts' input quantization fused in + (``ops.gemm.k3_latent_down_fp4``) when enabled, for the FC12 routed experts (which quantize + their input with ``fp4_quantize`` against ``fc31_input_scale``) of one launch chain. + """ + down = self.routed_expert_down_proj + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + if not self._fused_projection_applies(x, down): + return None + moe = self.routed_experts + if getattr(moe, "kernels", None) != "CUTEDSL_FC12" or x.shape[0] > moe.max_chunk_tokens: + return None + from .kernels import k3_latent_fp8_gemm + + global_sf = moe.backend.fc31_input_scale + w_scale = down.packed_weight_scale + if not k3_latent_fp8_gemm.can_run(x, down.weight, w_scale, global_sf): + return None + codes, scales = k3_latent_down_fp4_op(x, down.weight, w_scale, global_sf) + return Fp4QuantizedTensor(codes, scales, is_sf_swizzled=False) + + def _own_workspace_latent_tail( + self, + routed_out: torch.Tensor, + shared_out: torch.Tensor, + residual: Optional[torch.Tensor], + defer_residual: bool, + ): + """``TRTLLM_K3_LATENT_AR_OWN_WS``: the routed latent all-reduce + RMSNorm on its own MNNVL + workspace (``k3_latent_ar_norm`` up to ``_K3_LATENT_AR_ONESHOT_MAX_BYTES``, the C++ fused + all-reduce otherwise) and latent_up, issued on the main stream before it joins the shared-expert + stream, whose output only the residual update reads. Returns what ``forward`` returns.""" + workspace = self._latent_ar_ws + gamma = self.routed_expert_norm.weight + eps = float(self.routed_expert_norm.variance_epsilon) + latent = routed_out.reshape(-1, routed_out.shape[-1]).contiguous() + routed_latent = None + latent_mxfp8 = None + mid_plan = None + from .kernels import k3_latent_ar_norm + + world_size = workspace["world_size"] + one_shot = ( + latent.numel() * world_size * latent.element_size() <= _K3_LATENT_AR_ONESHOT_MAX_BYTES + ) + row = k3_latent_ar_norm.can_run_row(latent, gamma, world_size, workspace["buffer_bytes"]) + if one_shot and (row or k3_latent_ar_norm.can_run(latent, gamma, world_size)): + routed_latent = torch.empty_like(latent) + m, k = latent.shape + if row: + # The MXFP8 operand of latent_up's mid-M GEMM, in its scale layout. + mid_plan = self._lu_mid_prequant_plan(latent) + if mid_plan is not None: + latent_mxfp8 = ( + torch.empty((m, k), dtype=torch.float8_e4m3fn, device=latent.device), + torch.empty( + (k // 128) * mid_plan[1] * 512, + dtype=torch.uint8, + device=latent.device, + ), + ) + k3_latent_ar_rmsnorm( + latent, + gamma, + eps, + routed_latent, + uc_addr=workspace["buffer"].data_ptr(), + mc_addr=workspace["mc"].data_ptr(), + flags=workspace["flags"], + rank=workspace["rank"], + world_size=world_size, + mxfp8_out=latent_mxfp8, + sf_blocks=mid_plan[1] if mid_plan is not None else 0, + row=row, + ) + if routed_latent is None: + (routed_latent,) = torch.ops.trtllm.mnnvl_fusion_allreduce( + latent, + gamma, + None, + eps, + workspace["buffer"], + workspace["flags"], + True, + None, + int(AllReduceFusionOp.RMS_NORM), + ) + routed_latent = routed_latent.view(routed_out.shape) + if mid_plan is not None: + routed_out = self._latent_up_mid_prequantized(routed_latent, *latent_mxfp8, mid_plan[0]) + else: + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + self.moe_shared_event.wait() + if residual is not None: + if defer_residual: + return _DeferredMoEResidual(residual, routed_out, shared_out) + return k3_moe_residual_add_op(routed_out, shared_out, residual) + return routed_out + shared_out + + @staticmethod + def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): + return projection(hidden_states) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, projection.weight.t(), None, None + ) + + def _can_fuse_router_down(self, hidden_states: torch.Tensor) -> bool: + """Whether one fused GEMM can produce the router logits and the latent + down-projection (trtllm::dsv3_router_latent_gemm_op). + + Scoped to the validated configuration (TP MEGAMOE_CUTEDSL with fused + communication on SM107); the fusion itself does not depend on it. + """ + if ( + _K3_DISABLE_MIN_LATENCY_LATENT_PROJ + or not hidden_states.is_cuda + or self._attention_dp + or self._moe_backend != "MEGAMOE_CUTEDSL" + or not self._routed_fused_comm + or self.routed_experts.parallel_size <= 1 + or get_sm_version() != 107 + ): + return False + down = self.routed_expert_down_proj + if not isinstance(down, nn.Linear) or down.bias is not None: + return False + if ( + hidden_states.ndim != 2 + or not 1 <= hidden_states.shape[0] <= 16 + or hidden_states.shape[1] != 7168 + or self.gate.weight.shape != (896, 7168) + or down.weight.shape != (3584, 7168) + ): + return False + # Inspect the loaded weights: FP8 conversion can replace the down module. + return all( + tensor.dtype == torch.bfloat16 + and tensor.is_cuda + and tensor.device == hidden_states.device + and tensor.stride() == (7168, 1) + and tensor.data_ptr() % 16 == 0 + for tensor in (hidden_states, self.gate.weight, down.weight) + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + """Resolve the routed-expert ``(moe_tp, moe_ep)`` split. + + Precedence: + + 1. Explicit ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` + from the user config. Detected via + ``mapping.moe_tp_ep_user_specified`` so the auto-resolved mapping + default (``moe_tp=tp_size, moe_ep=1``) is NOT mistaken for a TP + request. + 2. Default: EP-only (``moe_tp=1, moe_ep=tp_size``), the historical + K3 layout. + """ + tp_size = mapping.tp_size + if getattr(mapping, "moe_tp_ep_user_specified", False): + return mapping.moe_tp_size, mapping.moe_ep_size + return 1, tp_size + + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Routed-expert quantization for ``layer_idx``, taken from the checkpoint. + + ``nvidia/Kimi-K3-NVFP4`` declares the routed experts per layer as + ``NVFP4`` with ``group_size=16``; the original ``moonshotai/Kimi-K3`` + declares nothing per layer and keeps the historical + ``W4A8_MXFP4_MXFP8`` default. Reading the checkpoint instead of + hardcoding is what lets one code path serve both. + + An exclusion outranks the per-layer entry and the default below: + ``create_weights`` treats an override as authoritative over anything + ``__post_init__`` wrote, so this return value stands in for both + quantization passes and exclusion is the one that runs second. It is + matched as a pattern, so it is asked only about real module names. + """ + quant_config = model_config.quant_config + if quant_config is not None and any( + quant_config.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + logger.debug( + "Kimi K3 layer %d routed experts: excluded from quantization, " + "keeping them unquantized", + layer_idx, + ) + return QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) + + per_layer = getattr(model_config, "quant_config_dict", None) + if per_layer: + for prefix in _K3_ROUTED_EXPERT_KEY_PREFIXES: + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES: + cfg = per_layer.get(f"{prefix}layers.{layer_idx}.{suffix}") + if cfg is not None and cfg.quant_algo is not None: + # Logged once per layer: the routed-expert format decides + # which MoE backends can serve this checkpoint at all. + logger.debug( + "Kimi K3 layer %d routed experts: %s (group_size=%s) " + "from the checkpoint", + layer_idx, + cfg.quant_algo, + cfg.group_size, + ) + return cfg + logger.debug( + "Kimi K3 layer %d routed experts: no per-layer quant config in the " + "checkpoint, defaulting to %s", + layer_idx, + _K3_DEFAULT_ROUTED_QUANT_ALGO, + ) + return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config. Default split is EP-only; see ``_select_moe_tp_ep``.""" + # The CuTe DSL kernel sets K3 runs (context: CUTEDSL; generation: CUTEDSL_FC12 or the + # fused-communication MEGAMOE_CUTEDSL). + supported_backends = (*_K3_ROUTED_KERNEL_SETS, "MEGAMOE_CUTEDSL") + if model_config.moe_backend not in supported_backends: + raise ValueError( + f"Kimi K3 SiTU routed experts run on {supported_backends}; " + f"got {model_config.moe_backend!r}." + ) + if model_config.moe_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "EPLB or replicated expert slots." + ) + mapping = model_config.mapping + if getattr(mapping, "_dwdp_size", 0) > 1: + # DWDP does not hook the load path at all: setup_dwdp() scans the + # model *after* every weight is loaded, reads w3_w1_weight / + # w2_weight and their block scales straight off the backend + # (dwdp/setup.py::collect_moe_params), copies them into fabric + # memory, frees the originals, and rebinds param.data to a + # composite VA view. K3's packed-checkpoint streaming produces + # exactly those tensors, so the two are compatible -- but only on + # the NVFP4 CuteDSL leaves that declare supports_dwdp (CUTEDSL, + # CUTEDSL_FC12). + # + # MEGAMOE_CUTEDSL is excluded deliberately rather than by + # oversight: its process_weights_after_loading packs into the mega + # buffers and shrinks the raw source params back to placeholders + # (see maybe_finalize_layer), which has never been checked against + # DWDP's copy-then-free of those same originals. + if model_config.moe_backend not in ("CUTEDSL", "CUTEDSL_FC12"): + raise NotImplementedError( + "Kimi K3 with DWDP requires moe_config.backend CUTEDSL or " + f"CUTEDSL_FC12; got {model_config.moe_backend!r}." + ) + + moe_tp, moe_ep = K3MoERuntime._select_moe_tp_ep(mapping) + if moe_tp < 1 or moe_ep < 1 or moe_tp * moe_ep != mapping.tp_size: + raise ValueError( + f"Kimi K3 routed MoE split moe_tp={moe_tp} x moe_ep={moe_ep} " + f"must multiply to tp_size={mapping.tp_size}." + ) + if moe_tp > 1 and mapping.enable_attention_dp: + raise NotImplementedError( + "Kimi K3 MoE tensor parallelism requires " + "enable_attention_dp=false (the attention-DP dispatch/combine " + "path is validated for EP-only splits)." + ) + logger.info_once( + f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " + f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", + key="kimi_k3_moe_tp_ep_split", + ) + + mapping_dict = mapping.to_dict() + mapping_dict["moe_cluster_size"] = 1 + mapping_dict["moe_tp_size"] = moe_tp + mapping_dict["moe_ep_size"] = moe_ep + routed_mapping = Mapping.from_dict(mapping_dict) + + routed_model_config = copy.copy(model_config) + routed_model_config._frozen = False + routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) + routed_model_config.mapping = routed_mapping + routed_model_config.moe_backend = model_config.moe_backend + # MegaMoE uses this value as global DP SymmBuffer capacity, then + # divides it by EP size for the per-rank allocation. Other backends + # keep the user-configured value as their MoE chunking bound. + # Preserve an explicitly larger capacity. + if routed_model_config.moe_backend == "MEGAMOE_CUTEDSL": + default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size + configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) + if configured_moe_max_num_tokens < default_moe_max_num_tokens: + logger.info_once( + "Kimi K3 MegaMoE raises moe_max_num_tokens from " + f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " + "because the global DP SymmBuffer requires capacity for " + "max_num_tokens * dp_size.", + key=( + "kimi_k3_megamoe_capacity_override_" + f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" + ), + ) + routed_model_config.moe_max_num_tokens = max( + configured_moe_max_num_tokens, + default_moe_max_num_tokens, + ) + routed_model_config._frozen = True + return routed_model_config + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + *, + residual: Optional[torch.Tensor] = None, + defer_residual: bool = False, + rp_tail_ok: bool = False, + ) -> torch.Tensor: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. + + With ``residual``, returns ``residual + moe_output``; with ``defer_residual`` as well, a + ``_DeferredMoEResidual`` instead where that update would be the separate BF16 add, or a + ``_DeferredRpTail`` where the row-parallel tail runs (the caller allows it with + ``rp_tail_ok`` for generation-only batches). + """ + identity = hidden_states + rp_target = self._rp_tail_target(hidden_states, residual, rp_tail_ok) + # The shared expert as the persistent kernel (forked with the ULL front's shared fork, + # else in the column-parallel tail). + persist = self._persist_applies(hidden_states.shape[0], rp_target) + projected_input = None + early_routed_in = None + shared_partial = None + front_routing = None + route_target = self._route_tail_target(hidden_states.shape[0]) + if self._ull_front_applies(hidden_states): + # TRTLLM_K3_ULL_FRONT: the shared expert forks where _ull_fork_shared puts it; the sharded + # front routes (route_target is the front's) and produces the NVFP4 routed input on this stream. + route_target = None + after_head, after_front, forked = self._ull_fork_shared( + identity, persist, hidden_states.shape[0] + ) + router_logits, projected_input, front_routing = self._ull_front( + hidden_states, after_head=after_head + ) + if after_front is not None: + after_front() + if forked: + shared_partial = forked[0] + elif self._can_fuse_router_down(hidden_states): + router_logits, projected_input = torch.ops.trtllm.dsv3_router_latent_gemm_op( + hidden_states, self.gate.weight, self.routed_expert_down_proj.weight + ) + else: + router_issued = False + if self.latent_down_stream is not None and do_multi_stream(): + # Forked ahead of the router GEMM: the latent chain needs only + # hidden_states. + self.moe_latent_fork_event.record() + if hidden_states.shape[0] <= _K3_ROUTER_FIRST_MAX_M: + # Router before latent_down in creation order (see _K3_ROUTER_FIRST_MAX_M). + router_logits = self._router_logits(hidden_states) + router_issued = True + with torch.cuda.stream(self.latent_down_stream): + self.moe_latent_fork_event.wait() + early_routed_in = self._fused_latent_input(hidden_states) + if early_routed_in is None: + early_routed_in = self._routed_projection( + hidden_states, self.routed_expert_down_proj + ) + self.moe_latent_ready_event.record() + # No record_stream (under graph capture it holds every layer's + # blocks until capture ends): both tensors are read on main only + # after the join, and later work on this stream forks from main. + if not router_issued: + router_logits = self._router_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None + overlap_shared_allreduce = ( + self._defer_shared_allreduce + and do_multi_stream() + and self.shared_expert_stream is not None + ) + num_tokens = hidden_states.shape[0] + + def _routed_output(output_push: Optional[LamportWorkspaceView] = None): + # ``output_push``: also push the routed rows into that one-shot all-reduce's slots + # (the column-parallel tail's k3_moe_tail1). + # Latent down/up projections via the min-latency fused GEMM op: + # at <=16 tokens (decode graphs) it runs a single pipelined + # bf16 kernel per projection instead of cuBLAS's split-K GEMV + + # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); + # for larger token counts the op falls back to cuBLAS internally. + # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear + # (A/B escape hatch). When the FP8 weight-read conversion has + # replaced the projection module, call it directly: its weight is + # an e4m3 buffer the bf16 dsv3 op must not read, and its forward + # is already a single fused GEMM (fp8_swap_ab_gemm). + logits = router_logits + moe_kwargs = {} + if projected_input is not None: + routed_in = projected_input + elif early_routed_in is not None: + routed_in = early_routed_in + # The routed experts join the event before their first read of x. + moe_kwargs = {"input_ready_event": self.moe_latent_ready_event} + else: + routed_in = self._fused_latent_input(hidden_states) + if routed_in is None: + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + if isinstance(routed_in, Fp4QuantizedTensor): + moe_kwargs["output_dtype"] = hidden_states.dtype + if isinstance(self.routed_experts, K3RoutedExperts): + # Routed by the CuTe tail (route_target), the sorted routing kernel or + # noaux_tc_op, unless the front routed already; then the kernels. + y = self.routed_experts( + routed_in, + logits, + route=self._router(route_target), + routing=front_routing, + input_ready_event=moe_kwargs.get("input_ready_event"), + output_push=output_push, + all_rank_num_tokens=all_rank_num_tokens, + ) + else: + if front_routing is not None or output_push is not None: + raise RuntimeError( + "caller routing or an output push needs the CuTe DSL experts" + ) + y = self.routed_experts( + routed_in, + logits, + all_rank_num_tokens=all_rank_num_tokens, + **moe_kwargs, + ) + if overlap_shared_allreduce: + self.moe_routed_event.record() + if self._reduce_routed_output: + return y + # Communication-backed paths return a complete routed result. + y = self.routed_expert_norm(y) + return self._routed_projection(y, self.routed_expert_up_proj) + + def _shared_output(): + """The shared expert's (reduced) output.""" + final_params = ( + AllReduceParams(enable_allreduce=False) if self._defer_shared_allreduce else None + ) + return _shared_reduce( + self.shared_experts(identity, final_all_reduce_params=final_params) + ) + + def _shared_reduce(shared_out: torch.Tensor) -> torch.Tensor: + """Reduce the shared expert's down-projection output.""" + if overlap_shared_allreduce: + # Wait for all routed chunks, but overlap the norm/up tail. + self.moe_routed_event.wait() + shared_out = self._shared_all_reduce()(shared_out) + return shared_out + + if rp_target is not None: + return self._rp_tail_forward( + _routed_output, identity, shared_partial, residual, rp_target, defer_residual + ) + # The final join includes the shared all-reduce, keeping it ordered before any outer + # routed reduction or the next layer's collectives. + if shared_partial is not None: + routed_out = _routed_output() + with torch.cuda.stream(self.shared_expert_stream): + shared_out = shared_partial + if self._shared_needs_all_reduce(): + shared_out = self._shared_all_reduce()(shared_partial) + self.moe_shared_event.record() + self.moe_shared_event.wait() + elif ( + self._latent_ar_ws is not None + and moe_all_reduce is not None + and not self._defer_shared_allreduce + and num_tokens <= self._latent_ar_ws["max_tokens"] + and self.shared_expert_stream is not None + and do_multi_stream() + ): + # The fork / join of fork_join, with the join left open until it is + # known whether the latent all-reduce runs on its own workspace. + self.moe_main_event.record() + routed_out = _routed_output() + with torch.cuda.stream(self.shared_expert_stream): + self.moe_main_event.wait() + shared_out = _shared_output() + self.moe_shared_event.record() + if routed_out.dtype == torch.bfloat16: + return self._own_workspace_latent_tail( + routed_out, shared_out, residual, defer_residual + ) + self.moe_shared_event.wait() + else: + # Shared experts depend only on the block input, so overlap their GEMMs + # with the routed dispatch/expert/combine chain. Multi-stream engages + # only under CUDA graphs; otherwise both branches run in order on the + # default stream. The shared expert includes its output all-reduce on + # the auxiliary stream (deferred past the routed MoE for fused-comm + # MegaMoE). The final join includes the shared all-reduce, so it + # precedes the routed all-reduce: concurrent collectives on different + # streams can corrupt SYMM_MEM all-reduce state. + routed_out, shared_out = fork_join( + _routed_output, + _shared_output, + self.moe_main_event, + self.moe_shared_event, + self.shared_expert_stream, + ) + if self._defer_shared_allreduce and not overlap_shared_allreduce: + shared_out = self._shared_all_reduce()(shared_out) + if self._reduce_routed_output: + # RMSNorm folded into the routed all-reduce (RMS_NORM fusion): the + # MNNVL kernel normalizes straight out of the lamport buffer, so + # the reduced row never round-trips through HBM for a norm kernel. + # Strategies without the fusion fall back inside AllReduce. + routed_latent = moe_all_reduce( + routed_out, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=self.routed_expert_norm.weight, + eps=float(self.routed_expert_norm.variance_epsilon), + ), + ) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + if residual is not None: + if defer_residual: + return _DeferredMoEResidual(residual, routed_out, shared_out) + return k3_moe_residual_add_op(routed_out, shared_out, residual) + return routed_out + shared_out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/_fp4_mla_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/_fp4_mla_decode.py new file mode 100644 index 000000000000..d97ddab4438a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/_fp4_mla_decode.py @@ -0,0 +1,980 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""K3's packed FP4 MLA generation decode: kernel, task plan and split reduction. + +The FP4 MLA backend (``attention/backends/fp4_mla``) owns the cache layout, the generation cache +update (KV append, HP ring, FP4 query rows) and the step metadata buffers; this module runs the +attention over them with K3's packed CuTe DSL decode kernel (``kernels/fp4_mla_mufu16``), whose +task plan and split reduction are ``kernels/fp4_mla_packed_reduce``. The kernel stores partials +by plan slot and optionally as normalized FP16, so the plan it reads is its own: written into +its own workspace with its own key, never the backend's (whose cache update writes the backend's +plan for the backend's decode). + +Per step, the first call that needs it writes the generation lengths (shared with the backend: +whichever runs first writes them, the other finds them current) fused with this plan, then the +device page table. :func:`begin_step` forgets the previous step's plan; the K3 model calls it at +the start of every forward, and an invalidation of the backend's step metadata (``prepare`` / +``on_update_kv_lens``) forgets it too. + +Settings: the split count (8 task slots per request), the makespan (LPT) plan and the wave plan +are fixed; the page plan, the reduction's programmatic launch, FP16 partials and the decode's +programmatic launch after the fused query are ``features`` fields. A late-released fold consumer +(``DeferredPackedReduce(release_early=False)``) is released from the decode's epilogue. +""" + +from __future__ import annotations + +import dataclasses +from typing import Any, NamedTuple, Optional + +import torch + +from tensorrt_llm._torch.attention.backends.fp4_mla.config import ( + _FP4_MLA_CUTEDSL_BACKEND, + _FP4_MLA_PACKED_MAX_SPLITS, + FP4_MLA_K_RESIDUAL_DIM, + FP4_MLA_Q_LOGICAL_DIM, + FP4_MLA_Q_PACKED_DIM, + FP4_MLA_Q_RESIDUAL_DIM, + FP4_MLA_SCALE_ROW_GROUP, + FP4_MLA_TOKENS_PER_BLOCK, + _cutedsl_backend_available, + _fp4_mla_attention_backend, + _fp4_mla_packed_rows_per_seq, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.layout import ( + _ensure_workspace_tensor, + _get_fp4_mla_global_scale, + _get_fp4_mla_kv_cache_tensors, + _get_fp4_mla_q_global_scale, + _get_fp4_mla_swizzled_scale_size, + _validate_fp4_mla_attention_q_shape, + _validate_fp4_mla_cache_shape, + _validate_fp4_mla_kv_storage_shape, + get_fp4_mla_v_scale_pool_view, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.metadata import ( + _fp4_mla_generation_page_ids, + _get_linear_mtp_query_len_per_seq, + _max_generation_pages, + materialize_fp4_mla_device_page_table, + populate_fp4_mla_generation_lengths, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.v_cache import ( + _get_cutedsl_persistent_v_packed_cache, + _get_fp4_mla_v_packed_pool_base, + _get_fp4_mla_v_scale_pool_base, +) +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger + +from ...features import active as _features + + +def _ceil_div(lhs: int, rhs: int) -> int: + return (lhs + rhs - 1) // rhs + + +# Task slots per request of the packed decode (grid.z = requests * splits). +_SPLITS = 8 +assert 1 < _SPLITS <= _FP4_MLA_PACKED_MAX_SPLITS + +# Workspace names (``fp4_mla_state.workspaces``); the backend's own buffers have other names. +_PLAN_WS = "_k3_fp4_mla_packed_plan" +_PARTIAL_WS = "_k3_fp4_mla_split_partial" +_PARTIAL_F16_WS = "_k3_fp4_mla_split_partial_f16" +_MAXIMA_WS = "_k3_fp4_mla_split_maxima" +_SUMS_WS = "_k3_fp4_mla_split_sums" + + +def _kernel_module() -> Any: + from ...kernels import fp4_mla_mufu16 + + return fp4_mla_mufu16 + + +def _reduce_module() -> Any: + from ...kernels import fp4_mla_packed_reduce + + return fp4_mla_packed_reduce + + +def applies(attn: Any) -> bool: + """Whether ``attn`` (a ``TrtllmAttention``) runs its FP4 MLA generation decode here: an FP4 MLA + layer on the CuTe DSL backend (SM107 with a CuTe DSL that builds the kernels). Elsewhere the + backend's own decode (Triton) runs.""" + return bool( + getattr(attn, "uses_fp4_mla_attention", False) + and _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND + ) + + +# --------------------------------------------------------------------------- +# Per-step state. +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class _Step: + """This module's state of one metadata object (kept on its ``fp4_mla_state``).""" + + # Key of the plan in _PLAN_WS for the current step; None: to be written. + plan_key: Optional[tuple] = None + # Kernels this module launched that write decode metadata (plan, lengths, page table), and + # that count when the kernel writing the query rows was launched (:func:`mark_query_launch`). + metadata_launches: int = 0 + query_launch_mark: Optional[int] = None + + +def _step(metadata: Any) -> _Step: + state = metadata.fp4_mla_state + step = getattr(state, "_k3_fp4_mla_decode", None) + if step is None: + step = _Step() + state._k3_fp4_mla_decode = step + return step + + +def begin_step(metadata: Any) -> None: + """Forget the previous step's plan (the K3 model calls this at the start of every forward).""" + if getattr(metadata, "fp4_mla_state", None) is not None: + _step(metadata).plan_key = None + + +def _count_metadata_launch(metadata: Any) -> None: + """Count a kernel that writes decode metadata on the current stream; a decode launched right + after one must not read that metadata before its wait.""" + _step(metadata).metadata_launches += 1 + + +def mark_query_launch(metadata: Any) -> None: + """Record that the kernel writing this step's query rows was just launched (the fused query, + ``features.ull2_mla_decode_pdl``): a decode launched as its programmatic dependent requires that + no metadata kernel was launched since.""" + step = _step(metadata) + step.query_launch_mark = step.metadata_launches + + +# --------------------------------------------------------------------------- +# Plan geometry and step metadata. +# --------------------------------------------------------------------------- + +_SM_COUNT_CACHE: dict[int, int] = {} + + +def _get_sm_count(device: torch.device) -> int: + """Return the SM (multiprocessor) count for ``device``, cached per index.""" + index = device.index if device.index is not None else torch.cuda.current_device() + count = _SM_COUNT_CACHE.get(index) + if count is None: + count = torch.cuda.get_device_properties(index).multi_processor_count + _SM_COUNT_CACHE[index] = count + return count + + +def _packed_wave_tasks(device: torch.device, m_tiles: int) -> int: + """Tasks that fill one wave of 2-CTA clusters.""" + clusters_per_wave = max(1, _get_sm_count(device) // 2) + return max(1, clusters_per_wave // max(1, m_tiles)) + + +class PackedPlanGeometry(NamedTuple): + """Per-step launch geometry of the packed split-KV decode.""" + + num_requests: int + query_len: int + splits: int + wave_tasks: int + profile_tiles: int + task_capacity: int + split_kv_len: int + helix_kv_bounds: torch.Tensor | None + plan: torch.Tensor + key: tuple + # Decode clusters per task and the wave's clusters of the makespan plan. + lpt_m_tiles: int + lpt_wave_clusters: int + # The wave plan (task length by simulated waves, full tasks first). + wave_plan: bool + + +def packed_plan_geometry( + metadata: Any, + *, + num_heads: int, + query_len: int, + num_requests: int, + device: torch.device, + helix_kv_bounds: torch.Tensor | None, +) -> PackedPlanGeometry | None: + """Slots, task-length cap and plan buffer of this step's packed decode. + + None when the packed Q layout is off or there are no generation pages. The geometry depends + only on the batch shape and the page-table capacity, so the prelaunch and the decode derive + the same one. + """ + rows_per_seq = _fp4_mla_packed_rows_per_seq(num_heads, query_len) + if rows_per_seq == 0 or num_requests <= 0: + return None + if not metadata.fp4_mla_state.device_page_table: + return None + # The fixed-stride generation page-table capacity (see _max_generation_pages). + max_pages = int(metadata.fp4_mla_state.page_table_stride) + if max_pages <= 0: + return None + max_kv_len = max_pages * metadata.page_size + is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) + m_tiles = rows_per_seq // FP4_MLA_SCALE_ROW_GROUP + splits = _SPLITS + # Tasks that fill one wave of 2-CTA clusters: the plan's task-length reference. + wave_tasks = _packed_wave_tasks(device, m_tiles) + # The runtime KV profile only has to cover one task: the kernel clamps the + # plan's task length to it, and ceil(longest / (splits - 1)) tiles per task + # always fits the splits * requests slots. A batch that fills the machine + # with one task per request keeps the full profile so even batches never + # have to split. + kv_tiles = _ceil_div(max(max_kv_len, 1), 256) + if wave_tasks <= 2 * num_requests: + task_tiles = kv_tiles + else: + task_tiles = min(kv_tiles, max(_ceil_div(kv_tiles, splits - 1), 2)) + split_kv_len = max(task_tiles * 256, 256) + # The kernels clamp the plan to the compiled profile (a bucket >= split_kv_len). + module = _kernel_module() + if not _features().fp4_mla_page_plan: + split_kv_len = max(split_kv_len, module.SMEM_P4_PAGE_PLAN_PROFILE_KV + 1) + profile_tiles = module._select_runtime_kv_profile(split_kv_len) // 256 + reduce = _reduce_module() + tasks = splits * num_requests + # Graph buckets hold exactly their tasks; eager metadata grows on demand. + task_capacity = tasks if is_cuda_graph else max(tasks, 32) + plan = _ensure_workspace_tensor( + metadata, + _PLAN_WS, + (reduce.packed_plan_length(num_requests, task_capacity),), + dtype=torch.int32, + device=device, + ) + kv_lens = metadata.fp4_mla_state.generation_kv_lens + lpt_m_tiles = m_tiles + lpt_wave_clusters = max(1, _get_sm_count(device) // 2) + wave_plan = True + key = ( + num_requests, + splits, + wave_tasks, + profile_tiles, + query_len, + task_capacity, + kv_lens.data_ptr() if isinstance(kv_lens, torch.Tensor) else 0, + helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else 0, + plan.data_ptr(), + lpt_m_tiles, + wave_plan, + ) + return PackedPlanGeometry( + num_requests, + query_len, + splits, + wave_tasks, + profile_tiles, + task_capacity, + split_kv_len, + helix_kv_bounds, + plan, + key, + lpt_m_tiles, + lpt_wave_clusters, + wave_plan, + ) + + +def helix_plan_bounds(metadata: Any, num_tokens: int, *, use_helix: bool) -> torch.Tensor | None: + """Per-token Helix KV bounds the decode plans on, or None (validated at decode).""" + if not use_helix or not bool(getattr(metadata, "_helix_spec_tokens_valid", False)): + return None + bounds = getattr(metadata, "helix_kv_bounds", None) + if ( + not isinstance(bounds, torch.Tensor) + or bounds.dtype != torch.int32 + or bounds.ndim != 1 + or bounds.numel() < num_tokens + or not bounds.is_contiguous() + ): + return None + return bounds[:num_tokens] + + +def _launch_packed_plan( + geometry: PackedPlanGeometry, + kv_lens: torch.Tensor, + metadata: Any, + *, + generation_lengths: tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None = None, +) -> None: + """Write the step's packed split plan, fused with the generation lengths when given.""" + _count_metadata_launch(metadata) + _reduce_module().launch_packed_plan( + kv_lens, + geometry.plan, + num_requests=geometry.num_requests, + num_splits=geometry.splits, + wave_tasks=geometry.wave_tasks, + max_tiles_per_task=geometry.profile_tiles, + query_len=geometry.query_len, + helix_kv_bounds=geometry.helix_kv_bounds, + generation_lengths=generation_lengths, + lpt_m_tiles=geometry.lpt_m_tiles, + lpt_wave_clusters=geometry.lpt_wave_clusters, + wave_plan=geometry.wave_plan, + ) + _step(metadata).plan_key = geometry.key + + +def uniform_generation_lengths( + metadata: Any, + num_gen_tokens: int, + num_gen: int, + plan_geometry: PackedPlanGeometry | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return the preallocated CUDA generation KV and append lengths of the backend's state. + + With ``plan_geometry`` the launch that writes the lengths also writes this module's plan of + the step (one CuTe DSL kernel instead of two launches), and a stale plan is rewritten even + when the lengths are already current. + """ + if num_gen <= 0 or num_gen_tokens % num_gen != 0: + raise RuntimeError("FP4 MLA generation requires a non-empty uniform request batch.") + if plan_geometry is not None and plan_geometry.num_requests != num_gen: + raise RuntimeError("FP4 MLA packed plan geometry does not match the generation batch.") + + state = metadata.fp4_mla_state + step = _step(metadata) + if state.generation_lengths_num_tokens == -1: + # The backend invalidated its step metadata (prepare / on_update_kv_lens): a new step. + step.plan_key = None + num_contexts = metadata.num_contexts + num_seqs = metadata.num_seqs + kv_lens_gen = metadata.kv_lens_cuda_runtime[num_contexts:num_seqs] + prompt_lens_gen = metadata.prompt_lens_cuda_runtime[num_contexts:num_seqs] + corrected_kv_lens = getattr(state, "generation_kv_lens", None) + generation_lens = getattr(state, "generation_append_lens", None) + tensors = ( + kv_lens_gen, + prompt_lens_gen, + corrected_kv_lens, + generation_lens, + ) + if ( + not all(isinstance(tensor, torch.Tensor) for tensor in tensors) + or corrected_kv_lens.numel() < num_gen + or generation_lens.numel() < num_gen + or not all(tensor.is_cuda for tensor in tensors) + ): + raise RuntimeError("FP4 MLA generation lengths require preallocated CUDA buffers.") + + record_for_capture = bool( + getattr(metadata, "is_cuda_graph", False) + and torch.cuda.is_current_stream_capturing() + and not getattr(state, "generation_lengths_capture_recorded", False) + ) + precomputed = ( + not record_for_capture + and state.generation_lengths_num_tokens == num_gen_tokens + and state.generation_lengths_num_seqs == num_gen + and state.generation_lengths_num_contexts == num_contexts + ) + plan_pending = plan_geometry is not None and step.plan_key != plan_geometry.key + if not precomputed: + if plan_pending: + _launch_packed_plan( + plan_geometry, + kv_lens_gen, + metadata, + generation_lengths=( + prompt_lens_gen, + corrected_kv_lens[:num_gen], + generation_lens[:num_gen], + num_gen_tokens // num_gen, + ), + ) + else: + _count_metadata_launch(metadata) + populate_fp4_mla_generation_lengths( + kv_lens_gen, + prompt_lens_gen, + corrected_kv_lens[:num_gen], + generation_lens[:num_gen], + num_gen_tokens=num_gen_tokens, + num_gen=num_gen, + ) + state.generation_lengths_num_tokens = num_gen_tokens + state.generation_lengths_num_seqs = num_gen + state.generation_lengths_num_contexts = num_contexts + if record_for_capture: + state.generation_lengths_capture_recorded = True + elif plan_pending: + _launch_packed_plan(plan_geometry, corrected_kv_lens[:num_gen], metadata) + return corrected_kv_lens[:num_gen], generation_lens[:num_gen] + + +def materialize_page_table( + metadata: Any, generation_kv_lens: Optional[torch.Tensor] = None +) -> None: + """The backend's fixed-stride device page table of this forward, from the final device + lengths (once per forward; counted as a metadata launch when it launches).""" + state = metadata.fp4_mla_state + if not bool(getattr(state, "device_page_table", False)): + raise RuntimeError("FP4 MLA cache update requires fixed-stride device page metadata.") + if bool(getattr(state, "device_page_table_valid", False)): + return + num_contexts = int(metadata.num_contexts) + num_sequences = int(metadata.num_seqs) + if generation_kv_lens is None: + if num_sequences > num_contexts: + num_generation_tokens = int(metadata.num_tokens) - int(metadata.num_ctx_tokens) + generation_kv_lens, _ = uniform_generation_lengths( + metadata, num_generation_tokens, num_sequences - num_contexts + ) + else: + generation_kv_lens = metadata.kv_lens_cuda_runtime[num_contexts:num_sequences] + _count_metadata_launch(metadata) + materialize_fp4_mla_device_page_table( + metadata, + metadata.kv_lens_cuda_runtime[:num_sequences], + generation_kv_lens, + ) + + +def prelaunch_generation_metadata( + metadata: Any, *, num_heads: int, use_helix: bool, device: torch.device +) -> bool: + """``features.ull4_mla_plan_early``: launch on the current stream what the first FP4 MLA + generation layer of a step launches for the step's metadata: the generation lengths fused + with this module's plan of ``num_heads`` query heads, then the device page table. Later calls + of the step find them current and launch nothing. Returns whether it launched (a + generation-only step with the device page table); the caller orders the step's first FP4 MLA + kernel after this stream.""" + state = getattr(metadata, "fp4_mla_state", None) + if state is None or not state.device_page_table: + return False + num_contexts = int(metadata.num_contexts) + num_gen = int(metadata.num_seqs) - num_contexts + num_tokens = int(metadata.num_tokens) + if num_contexts != 0 or num_gen <= 0 or num_tokens % num_gen != 0 or num_heads <= 0: + return False + # The plan workspace is reallocated on a device mismatch, and tensors carry an index. + device = torch.device(device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + plan_geometry = packed_plan_geometry( + metadata, + num_heads=num_heads, + query_len=num_tokens // num_gen, + num_requests=num_gen, + device=device, + helix_kv_bounds=helix_plan_bounds(metadata, num_tokens, use_helix=use_helix), + ) + kv_lens_gen, _ = uniform_generation_lengths( + metadata, num_tokens, num_gen, plan_geometry=plan_geometry + ) + materialize_page_table(metadata, kv_lens_gen) + return True + + +# --------------------------------------------------------------------------- +# Decode. +# --------------------------------------------------------------------------- + + +def _run_packed_decode( + module: Any, + metadata: Any, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_packed: torch.Tensor, + v_sf: torch.Tensor, + kv_global_scale: torch.Tensor, + q_global_scale: torch.Tensor, + src_page_ids: torch.Tensor, + page_indptr: torch.Tensor, + kv_lens: torch.Tensor, + output: torch.Tensor, + *, + num_heads: int, + query_len: int, + num_requests: int, + rows_per_seq: int, + sm_scale: float, + v_page_offset: int, + helix_kv_bounds: torch.Tensor | None, + softmax_stats_tensor: torch.Tensor | None, + plan_geometry: PackedPlanGeometry, + metadata_launches: int, + q_rows_external: bool, +) -> None: + """Packed CuTe DSL decode: one kernel over (split, request) tasks. + + The splits write partials (FP32, or normalized FP16 with ``features.fp4_mla_f16_partials``) + that one reduction kernel folds into the output and stats (or a ``DeferredPackedReduce`` + consumer folds). Nothing pads, replicates, or copies Q or the output. + + ``metadata_launches`` is the metadata-launch count when the decode call began; with + ``features.ull2_mla_decode_pdl`` the decode is launched programmatically only if no metadata + kernel was launched since and the query rows come from another kernel (``q_rows_external``). + """ + valid_m, physical_m = module.packed_query_rows(num_heads, query_len) + if physical_m != rows_per_seq: + raise RuntimeError( + f"FP4 MLA packed Q rows {rows_per_seq} do not match the kernel geometry {physical_m}." + ) + if v_packed is None: + raise RuntimeError("FP4 MLA packed decode requires the persistent packed-V cache.") + if not output.is_contiguous() or output.numel() != num_requests * valid_m * 512: + raise RuntimeError( + "FP4 MLA packed decode requires a contiguous [tokens, heads, 512] output." + ) + device = output.device + splits = plan_geometry.splits + step = _step(metadata) + use_pdl_wait = ( + _features().ull2_mla_decode_pdl + and q_rows_external + and (step.query_launch_mark is None or step.query_launch_mark == metadata_launches) + ) + output_rows = output.view(num_requests * valid_m, 512) + stats_rows = None + if softmax_stats_tensor is not None: + stats_rows = softmax_stats_tensor.view(num_requests * valid_m, 2) + common = dict( + max_kv_len=plan_geometry.split_kv_len, + sm_scale=sm_scale, + num_heads=num_heads, + query_len=query_len, + num_requests=num_requests, + num_splits=splits, + q_global_scale=q_global_scale, + v_page_offset=v_page_offset, + helix_kv_bounds=helix_kv_bounds, + ) + reduce = _reduce_module() + reduce_pdl = reduce.reduce_pdl_enabled() + deferred = reduce.DeferredPackedReduce.active() + launch_dependents = reduce_pdl and (deferred is None or deferred.release_early) + if deferred is not None and not deferred.release_early: + launch_dependents = module.LAUNCH_DEPENDENTS_EPILOGUE + f16_partials = reduce.f16_partials_enabled() + tasks = splits * num_requests + task_capacity = plan_geometry.task_capacity + plan = plan_geometry.plan + # The first call of the step wrote the plan fused with the generation lengths; a stale key + # means this call plans on other bounds. + if step.plan_key != plan_geometry.key: + _launch_packed_plan(plan_geometry, kv_lens, metadata) + partial = _ensure_workspace_tensor( + metadata, + _PARTIAL_F16_WS if f16_partials else _PARTIAL_WS, + (task_capacity, physical_m, 512), + dtype=torch.float16 if f16_partials else torch.float32, + device=device, + )[:tasks] + maxima = _ensure_workspace_tensor( + metadata, + _MAXIMA_WS, + (task_capacity, physical_m), + dtype=torch.float32, + device=device, + )[:tasks] + sums = _ensure_workspace_tensor( + metadata, + _SUMS_WS, + (task_capacity, physical_m), + dtype=torch.float32, + device=device, + )[:tasks] + module.run_trtllm_fp4_mla_decode_packed( + q_fp4, + q_sf, + kv_cache, + sf_cache, + v_packed, + v_sf, + kv_global_scale, + src_page_ids, + page_indptr, + kv_lens, + output_rows, + split_kv_buffers=(partial, maxima, sums), + plan=plan, + **common, + launch_dependents=launch_dependents, + pdl_wait=use_pdl_wait and step.metadata_launches == metadata_launches, + ) + reduce.reduce_packed_splits( + partial, + maxima, + sums, + output_rows, + q_global_scale, + kv_global_scale, + num_requests=num_requests, + num_splits=splits, + valid_m=valid_m, + physical_m=physical_m, + sm_scale=sm_scale, + plan=plan, + # FP16 partials are already in output units. + output_scale=1.0 if f16_partials else 1.0 / module.FP4_MLA_P_GLOBAL_SCALE, + stats=stats_rows, + sm_count=_get_sm_count(device), + use_pdl=reduce_pdl, + ) + + +def run_decode( + metadata: Any, + layer_idx: int, + local_layer: int, + q: torch.Tensor, + output: torch.Tensor, + *, + sm_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, + prequantized_q: torch.Tensor, + prequantized_q_sf: torch.Tensor, + q_batch_capacity: int, + softmax_stats_tensor: torch.Tensor | None = None, + q_rows_external: bool = False, +) -> None: + """MLA decode with FP4 QK and FP4 PV tensor-core matmuls on the packed query rows the + backend's generation cache update prepared (and wrote, or ``q_rows_external``: another + kernel wrote them, the fused query). + + ``q`` (``[tokens, heads, latent + RoPE]``) provides shapes; ``output`` is ``[tokens, heads, + 512]`` and ``softmax_stats_tensor`` the Helix ``[tokens, heads, 2]`` (max, sum) statistics. + """ + head_dim = kv_lora_rank + qk_rope_head_dim + metadata_launches = _step(metadata).metadata_launches + if qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM: + raise ValueError( + "FP4 MLA K residual attention requires " + f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got {qk_rope_head_dim}." + ) + _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) + if metadata.page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + f"FP4 MLA attention decode requires page_size={FP4_MLA_TOKENS_PER_BLOCK}, " + f"got {metadata.page_size}." + ) + if q.ndim != 3 or q.shape[-1] != head_dim: + raise ValueError( + "FP4 MLA attention Q must have shape " + f"[tokens, heads, {head_dim}], got {tuple(q.shape)}." + ) + if not q.is_contiguous(): + raise ValueError("FP4 MLA attention Q must be contiguous.") + + num_queries = q.shape[0] + if num_queries == 0: + raise ValueError("FP4 MLA attention decode requires at least one query token.") + num_gen_seqs = metadata.num_seqs - metadata.num_contexts + query_len_per_seq = _get_linear_mtp_query_len_per_seq( + metadata, + num_queries=num_queries, + num_gen_seqs=num_gen_seqs, + ) + num_heads = q.shape[1] + if output.shape[:2] != (num_queries, num_heads): + raise ValueError("FP4 MLA attention output batch dimensions do not match.") + + if get_sm_version() != 107: + raise RuntimeError( + "K3's FP4 MLA decode requires Rubin SM107; " + f"current architecture is SM{get_sm_version()}." + ) + if not _cutedsl_backend_available(): + raise RuntimeError( + "K3's FP4 MLA decode requires a CuTe DSL build that provides the tcgen05 " + "descriptor K-segment offset, and cuda-python." + ) + if not 0 < num_heads <= 128 or kv_lora_rank != 512: + raise ValueError( + "K3's FP4 MLA decode requires 1-128 local heads and kv_lora_rank=512, got " + f"num_heads={num_heads}, kv_lora_rank={kv_lora_rank}." + ) + + helix_spec_tokens_valid = bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + helix_kv_bounds = None + if softmax_stats_tensor is not None: + if query_len_per_seq != 1 and not helix_spec_tokens_valid: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + expected_stats_shape = (num_queries, num_heads, 2) + if ( + softmax_stats_tensor.shape != expected_stats_shape + or softmax_stats_tensor.dtype != torch.float32 + or softmax_stats_tensor.device != q.device + or not softmax_stats_tensor.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix requires contiguous same-device float32 softmax " + f"stats with shape {expected_stats_shape}." + ) + if helix_spec_tokens_valid: + helix_kv_bounds = getattr(metadata, "helix_kv_bounds", None) + if ( + not isinstance(helix_kv_bounds, torch.Tensor) + or helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.device != q.device + or helix_kv_bounds.ndim != 1 + or helix_kv_bounds.numel() < num_queries + or not helix_kv_bounds.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix KV bounds must be a contiguous " + "same-device int32 tensor covering every query token." + ) + helix_kv_bounds = helix_kv_bounds[:num_queries] + state = metadata.fp4_mla_state + v_sf_pool = getattr(state, "v_scale_pool", None) + if v_sf_pool is None: + raise RuntimeError( + "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." + ) + + global_scale = _get_fp4_mla_global_scale(metadata, q.device) + _validate_fp4_mla_attention_q_shape(head_dim, FP4_MLA_Q_RESIDUAL_DIM) + if prequantized_q is None or prequantized_q_sf is None or q_batch_capacity is None: + raise RuntimeError( + "FP4 MLA decode requires Q prequantized by the fused generation cache update." + ) + capacity = int(q_batch_capacity) + # The generation cache update records the Q layout it produced. + rows_per_seq = int(getattr(state, "prequantized_q_rows_per_seq", None) or 0) + if rows_per_seq == 0: + raise NotImplementedError( + "K3's FP4 MLA decode reads the packed per-request Q layout; the generation cache " + "update prepared the unpacked one (TRTLLM_FP4_MLA_PACKED_Q=0, or more than " + "1024 rows per request)." + ) + if rows_per_seq != _fp4_mla_packed_rows_per_seq(num_heads, query_len_per_seq): + raise RuntimeError( + "FP4 MLA prequantized Q layout does not match the packed decode " + f"configuration (rows per request {rows_per_seq})." + ) + active_q_rows = num_gen_seqs * rows_per_seq + active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size(active_q_rows, FP4_MLA_Q_LOGICAL_DIM) + if ( + q.dtype != torch.bfloat16 + or capacity <= 0 + or num_queries > capacity + or prequantized_q.ndim != 2 + or prequantized_q.shape[0] < active_q_rows + or prequantized_q.shape[1] != FP4_MLA_Q_PACKED_DIM + or prequantized_q_sf.numel() < active_q_sf_bytes + or prequantized_q.dtype != torch.uint8 + or not prequantized_q.is_contiguous() + or prequantized_q_sf.dtype != torch.float8_e4m3fn + or not prequantized_q_sf.is_contiguous() + ): + raise ValueError("FP4 MLA prequantized Q does not satisfy the fused-Q contract.") + q_fp4 = prequantized_q[:active_q_rows] + q_sf = prequantized_q_sf[:active_q_sf_bytes] + q_global_scale = _get_fp4_mla_q_global_scale(metadata, q.device) + + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=_FP4_MLA_CUTEDSL_BACKEND, + ) + sf_cache = sf_cache.view(torch.float8_e4m3fn) + v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=kv_lora_rank)[local_layer].view( + torch.float8_e4m3fn + ) + # The kv_lens runtime alias can lag at the decode anchor (seq_lens == 1) under CUDA graph / + # one-engine MTP; the generation lengths recover the true total per sequence so the + # per-query causal masking sees the full 1 + draft_len window. + plan_geometry = packed_plan_geometry( + metadata, + num_heads=num_heads, + query_len=query_len_per_seq, + num_requests=num_gen_seqs, + device=q.device, + helix_kv_bounds=helix_kv_bounds, + ) + if plan_geometry is None: + raise RuntimeError("FP4 MLA packed decode requires generation cache pages.") + kv_lens, _ = uniform_generation_lengths( + metadata, num_queries, num_gen_seqs, plan_geometry=plan_geometry + ) + materialize_page_table(metadata, kv_lens) + src_page_ids = _fp4_mla_generation_page_ids(metadata, num_gen_seqs) + if _max_generation_pages(metadata) == 0: + raise RuntimeError("FP4 MLA attention decode requires generation cache pages.") + + module = _kernel_module() + v_packed = _get_cutedsl_persistent_v_packed_cache( + metadata, + local_layer, + kv_cache, + v_head_dim=kv_lora_rank, + page_size=metadata.page_size, + block_v=module.SMEM_P4_V_N_PER_CTA, + ) + core_v_packed = _get_fp4_mla_v_packed_pool_base(metadata) + if core_v_packed is None: + raise RuntimeError("Persistent FP4 MLA V packing requires a stable full-pool base.") + get_v_page_offset = getattr(metadata.kv_cache_manager, "get_mla_v_packed_page_offset", None) + v_page_offset = ( + int(get_v_page_offset(local_layer)) + if callable(get_v_page_offset) + else local_layer * kv_cache.shape[0] + ) + page_bytes = kv_lora_rank * (metadata.page_size // 2) + if core_v_packed.data_ptr() + v_page_offset * page_bytes != v_packed.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-packed layer view does not match its full-pool base and page " + "offset." + ) + v_sf_pool_base = _get_fp4_mla_v_scale_pool_base(metadata) + if v_sf_pool_base is None: + v_sf_pool_base = v_sf_pool.flatten(0, 1).view(torch.uint8) + if v_sf_pool_base.data_ptr() != v_sf_pool.data_ptr(): + raise RuntimeError("Persistent FP4 MLA V-scale pool must flatten without a copy.") + if ( + not isinstance(v_sf_pool_base, torch.Tensor) + or v_sf_pool_base.dtype != torch.uint8 + or v_sf_pool_base.device != v_sf.device + or v_sf_pool_base.ndim != 2 + or v_sf_pool_base.shape[1] != v_sf_pool.shape[-1] + or not v_sf_pool_base.is_contiguous() + ): + raise RuntimeError( + "Persistent FP4 MLA V-scale pool base must be a contiguous two-dimensional uint8 " + "tensor with the configured page stride." + ) + core_v_sf = v_sf_pool_base.view(torch.float8_e4m3fn) + get_v_sf_page_offset = getattr(metadata.kv_cache_manager, "get_mla_v_scale_page_offset", None) + v_sf_page_offset = ( + int(get_v_sf_page_offset(local_layer)) + if callable(get_v_sf_page_offset) + else local_layer * kv_cache.shape[0] + ) + if v_sf_page_offset != v_page_offset: + raise RuntimeError( + "Persistent FP4 MLA V-packed and V-scale pools require matching encoded layer offsets." + ) + expected_v_sf_ptr = ( + core_v_sf.data_ptr() + v_sf_page_offset * v_sf_pool.stride(1) * v_sf_pool.element_size() + ) + if expected_v_sf_ptr != v_sf.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-scale layer view does not match its full-pool base and page " + "offset." + ) + _run_packed_decode( + module, + metadata, + q_fp4, + q_sf, + kv_cache, + sf_cache, + core_v_packed, + core_v_sf, + global_scale, + q_global_scale, + src_page_ids, + state.paged_kv_indptr_decode[: num_gen_seqs + 1], + kv_lens, + output, + num_heads=num_heads, + query_len=query_len_per_seq, + num_requests=num_gen_seqs, + rows_per_seq=rows_per_seq, + sm_scale=float(sm_scale), + v_page_offset=v_page_offset, + helix_kv_bounds=helix_kv_bounds, + softmax_stats_tensor=softmax_stats_tensor, + plan_geometry=plan_geometry, + metadata_launches=metadata_launches, + q_rows_external=q_rows_external, + ) + + +def run_generation( + attn: Any, + q: torch.Tensor, + metadata: Any, + *, + softmax_stats_tensor: Optional[torch.Tensor] = None, + q_rows_external: bool = False, +) -> torch.Tensor: + """The absorbed generation attention of the FP4 MLA layer ``attn`` (a ``TrtllmAttention``, + :func:`applies`): what its ``forward`` returns for a generation-only call with the fused query + ``q`` (``[tokens, heads * (latent + RoPE)]``), after the backend's generation cache update. + Returns ``[tokens, heads * kv_lora_rank]`` in ``q``'s dtype; ``softmax_stats_tensor`` receives + the Helix statistics.""" + kv_lora_rank = attn.kv_lora_rank or 0 + qk_rope_head_dim = attn.qk_rope_head_dim or 0 + num_heads = attn.num_heads + num_tokens = q.shape[0] + if metadata.num_generations <= 0: + raise RuntimeError("FP4 MLA generation requires generation requests.") + kv_cache_manager = metadata.kv_cache_manager + if kv_cache_manager is None or kv_cache_manager.kv_factor != 1: + raise RuntimeError("FP4 MLA requires a SELF-K-only KV cache manager.") + if metadata.beam_width != 1: + raise NotImplementedError("FP4 MLA does not support beam search.") + state = metadata.fp4_mla_state + if state is None or state.hp_pool is None: + raise RuntimeError("FP4 MLA requires prepared high-precision KV state.") + if ( + getattr(metadata, "helix_position_offsets", None) is not None + and softmax_stats_tensor is None + ): + # Helix combines the per-rank partial outputs with these stats. + raise RuntimeError("FP4 MLA Helix requires softmax_stats_tensor.") + if not bool(getattr(state, "generation_cache_scattered", False)): + raise RuntimeError( + "FP4 MLA generation requires fused RoPE/Q quantization and " + "cache/HP-pool update before attention." + ) + state.generation_cache_scattered = False + features = _features() + logger.info_once( + f"Kimi K3 FP4 MLA generation decode: k3_rubin packed kernel, {_SPLITS} KV splits, " + f"page_plan={features.fp4_mla_page_plan}, reduce_pdl={features.fp4_mla_reduce_pdl}, " + f"f16_partials={features.fp4_mla_f16_partials}, " + f"helix={softmax_stats_tensor is not None}", + key="k3_fp4_mla_generation_decode", + ) + output = q.new_empty((num_tokens, num_heads * kv_lora_rank)) + sm_scale = 1.0 / (attn.q_scaling * ((attn.qk_nope_head_dim or 0) + qk_rope_head_dim) ** 0.5) + try: + run_decode( + metadata, + attn.layer_idx, + attn.get_fp4_mla_local_layer_idx(metadata), + q.view(num_tokens, num_heads, kv_lora_rank + qk_rope_head_dim), + output.view(num_tokens, num_heads, kv_lora_rank), + sm_scale=sm_scale, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + prequantized_q=state.prequantized_q, + prequantized_q_sf=state.prequantized_q_sf, + q_batch_capacity=state.q_batch_capacity, + softmax_stats_tensor=softmax_stats_tensor, + q_rows_external=q_rows_external, + ) + finally: + state.prequantized_q = None + state.prequantized_q_sf = None + state.q_batch_capacity = None + state.prequantized_q_rows_per_seq = None + return output diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/fp4_mla_gather_mxfp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/fp4_mla_gather_mxfp8.py new file mode 100644 index 000000000000..9ac4abb9fbac --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/fp4_mla_gather_mxfp8.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gather one cached-prefix partition of the FP4 MLA cache as MXFP8 kv_b input and FP8 K-pe.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import fp4_mla_context_gather as _gather +from .k3_mla_q_fused import e2m1_values, swizzled_sf_offset, unpack_e2m1 + +_ROW_BYTES = 320 +_SF_PER_ROW = 40 + + +def _round_up_ue8m0(x: torch.Tensor) -> torch.Tensor: + """Biased exponent of the smallest power of two >= ``x`` (positive FP32).""" + bits = x.view(torch.int32) + exponent = (bits >> 23) & 0xFF + return (exponent + ((bits & 0x7FFFFF) != 0).to(torch.int32)).clamp(0, 254) + + +class Fp4MlaGatherMxfp8(OpWrapper): + """``gather_mxfp8``: per partition token, the BF16-rounded latent of its FP4 cache row as + MXFP8 into ``a`` (E4M3 ``[M, 512]``) / ``asc`` (UE8M0, K32 R128c4, ``pad_up(M, 128) * 16`` + bytes), and the BF16-rounded RoPE part plus its residual as unit-scale E4M3 into ``pe`` + ``[M, 64]``. Returns None.""" + + ARCHS = frozenset({Arch.SM_107}) + # Spec: the cached-prefix chunk schedule of `ctx_split_k`'s tests: per-request cached + # `prefix` and new tokens, the chunk size and buffer batch; each loop of the schedule is one + # call. The cache is the FP4 MLA V2 cache with the backend's KV global scale. + CELLS: tuple[Cell, ...] = ( + Cell( + why="tp1 context: one request with a 1000-token prefix, 256-token chunks", + spec=dict(prefix=[1000], new=[64], chunk_size=256, buffer_batch=1), + ), + Cell( + why="tp1 context: three requests, prefixes across page boundaries, one partition each", + spec=dict(prefix=[300, 129, 5], new=[17, 3, 40], chunk_size=4096, buffer_batch=3), + ), + ) + + note = """ + Reads dense page-128 FP4 MLA V2 storage `kv` (320-byte rows: 512 latent + 64 RoPE + 64 RoPE + residual E2M1 values) with its swizzled E4M3 scales `sf` (40 per row) and FP32 + `global_scale`; request r's tokens are rows `offsets[r] .. offsets[r] + cu[r + 1] - cu[r]` of + its pages `pages[r * page_table_stride ..]`, written at rows `cu[r] ..` of the outputs. The + bytes are those of the dense BF16 gather followed by `fp8_quantize_1x128_packed_ue8m0` and + `static_quantize_e4m3_per_tensor` with a unit scale. Grid: ceil(max_chunk / warps) + 1 by + the requests; prefix lengths and page strides are runtime values. A process cache holds one + compiled kernel per (device, warps). + """ + + def __call__( + self, + a: torch.Tensor, + asc: torch.Tensor, + pe: torch.Tensor, + kv: torch.Tensor, + sf: torch.Tensor, + pages: torch.Tensor, + cu: torch.Tensor, + offsets: torch.Tensor, + global_scale: torch.Tensor, + max_chunk: int, + page_size: int, + page_table_stride: int, + num_contexts: int, + ) -> None: + _gather.gather_mxfp8( + a, + asc, + pe, + kv, + sf, + pages, + cu, + offsets, + global_scale, + max_chunk, + page_size, + page_table_stride=page_table_stride, + num_contexts=num_contexts, + ) + + def reference( + self, + a: torch.Tensor, + asc: torch.Tensor, + pe: torch.Tensor, + kv: torch.Tensor, + sf: torch.Tensor, + pages: torch.Tensor, + cu: torch.Tensor, + offsets: torch.Tensor, + global_scale: torch.Tensor, + max_chunk: int, + page_size: int, + page_table_stride: int, + num_contexts: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The post-state ``(a, asc, pe)``. Each cache value is ``e2m1 * e4m3 scale / + global_scale`` rounded to BF16 (exact division where the kernel uses ``div.full``); a + 128-value block of the latent gets the UE8M0 exponent of the smallest power of two >= + amax / 448 (amax floored at 1e-10), stored in its four K32 slots, and codes + ``e4m3(value / 2^(e - 127))``; K-pe is ``e4m3(bf16(main + residual))``. Scale rows past + ``M`` up to the next multiple of 128 become 0; rows of invalid pages are left alone.""" + out_a, out_asc, out_pe = a.clone(), asc.clone(), pe.clone() + gs = global_scale.float().view(()) + num_pages = kv.shape[0] + rows_view = kv.as_strided( + (num_pages, page_size, _ROW_BYTES), (kv.stride(0), kv.stride(2), 1), kv.storage_offset() + ) + sf_view = sf.view(torch.uint8).as_strided( + (num_pages, page_size * _SF_PER_ROW), (sf.stride(0), 1), sf.storage_offset() + ) + m = a.shape[0] + padded = (m + 127) // 128 * 128 + if padded > m: + pad_rows = torch.arange(m, padded, device=a.device) + pad_offsets = swizzled_sf_offset( + pad_rows[:, None], torch.arange(16, device=a.device), 16 + ) + out_asc.view(-1)[pad_offsets.flatten()] = 0 + cu_list, offset_list = cu.tolist(), offsets.tolist() + table = pages.tolist() + for r in range(num_contexts): + begin, end = cu_list[r], cu_list[r + 1] + for local in range(min(end - begin, max_chunk)): + position = offset_list[r] + local + logical = position // page_size + if position < 0 or logical >= page_table_stride: + continue + physical = table[r * page_table_stride + logical] + if not 0 <= physical < num_pages: + continue + row = position % page_size + values = self._row_values(rows_view[physical, row], sf_view[physical], row, gs) + out = begin + local + latent = values[:512].view(4, 128) + amax = latent.abs().amax(-1).clamp_min(1e-10) + exponent = _round_up_ue8m0((amax.double() / 448.0).float()) + inv = torch.exp2(127.0 - exponent.double()) + codes = (latent.double() * inv[:, None]).clamp(-448, 448).float() + out_a[out] = codes.to(torch.float8_e4m3fn).view(512) + sf_cols = torch.arange(16, device=a.device) + sf_offsets = swizzled_sf_offset(torch.tensor(out, device=a.device), sf_cols, 16) + out_asc.view(-1)[sf_offsets] = exponent.repeat_interleave(4).to(torch.uint8) + out_pe[out] = values[512:].clamp(-448, 448).to(torch.float8_e4m3fn) + return out_a, out_asc, out_pe + + @staticmethod + def _row_values(row_bytes: torch.Tensor, sf_page: torch.Tensor, row: int, gs) -> torch.Tensor: + """The BF16-rounded 512 latent and 64 K-pe (main + residual) values of one cache row.""" + cols = torch.arange(_SF_PER_ROW, device=row_bytes.device) + offsets = swizzled_sf_offset(torch.tensor(row, device=row_bytes.device), cols, _SF_PER_ROW) + scales = sf_page[offsets].view(torch.float8_e4m3fn).float() + codes = unpack_e2m1(row_bytes) + values = e2m1_values(codes).view(_SF_PER_ROW, 16).double() * scales.double()[:, None] + values = (values / gs.double()).float() + latent = values[:32].flatten().to(torch.bfloat16).float() + rope = (values[32:36].flatten() + values[36:40].flatten()).to(torch.bfloat16).float() + return torch.cat((latent, rope)) + + +fp4_mla_gather_mxfp8 = Fp4MlaGatherMxfp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/fp4_mla_kv_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/fp4_mla_kv_append.py new file mode 100644 index 000000000000..79b3a2bca360 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/fp4_mla_kv_append.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KV-only FP4 MLA generation cache update of the fused query (Triton).""" + +from __future__ import annotations + +from typing import Any + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import fp4_mla_kv_append as _append + + +class Fp4MlaKvAppend(OpWrapper): + """``fp4_mla_kv_append_generation``: append this step's latent rows (BF16 ``[T, 576]``) of the + FP4 MLA layer behind ``mqa`` to its FP4 cache, HP ring, V scales and packed V, and prepare + (not write) the packed-Q buffers of ``num_q_heads`` heads on ``metadata.fp4_mla_state``.""" + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = () + + note = """ + The FP4 MLA backend's fused generation cache update with its query rows switched off + (`WRITE_Q_ROWS`) and a grid of the KV programs only: the cache bytes are those of the full + update. The cache pools, page tables, lengths and Q buffers are the backend's, reached through + `metadata` (the KV cache manager owns their layout); `mqa` supplies the layer index, the RoPE + table and the local layer of the pools. The fused query (`k3_mla_q_fused`) then writes the + query rows into the buffers prepared here, on another stream. + """ + + def __call__( + self, mqa: Any, latent_cache: torch.Tensor, metadata: Any, num_q_heads: int + ) -> None: + _append.fp4_mla_kv_append_generation(mqa, latent_cache, metadata, num_q_heads) + + def reference( + self, mqa: Any, latent_cache: torch.Tensor, metadata: Any, num_q_heads: int + ) -> None: + """Not modelled yet: the post-state of the FP4 cache page, HP ring, V scale and packed V + pools the KV cache manager owns, in their swizzled layouts. Until it is, the entry has no + cells; the kernel is checked bitwise against the backend's full generation update + (``test_k3_mla_q_fused.py``).""" + raise NotImplementedError("Fp4MlaKvAppend.reference: no plain-torch model yet") + + +fp4_mla_kv_append = Fp4MlaKvAppend() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_context_fmha.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_context_fmha.py new file mode 100644 index 000000000000..8c8c39462980 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_context_fmha.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 Rubin FP8 split-K context FMHA: attend one KV partition and merge it into the output.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ...kernels import k3_context_fmha_op as _fmha + +# Heads per reference block: bounds the FP32 score tensor of an 8k x 8k causal chunk. +_REF_HEADS = 4 + + +def _merge(old_out, old_max, old_sum, cur_out, cur_max, cur_sum): + """Merge two softmax-normalized partial outputs by their (row max, row sum) statistics.""" + merged_max = torch.maximum(old_max, cur_max) + old_weight = old_sum * (old_max - merged_max).exp() + cur_weight = cur_sum * (cur_max - merged_max).exp() + merged_sum = old_weight + cur_weight + merged = (old_out * old_weight[..., None] + cur_out * cur_weight[..., None]) / merged_sum[ + ..., None + ] + return merged, merged_max, merged_sum + + +class K3ContextFmha(OpWrapper): + """``run_k3_context_fmha``: FP8 Q ``[T_q, H, 192]`` against one FP8 KV partition (K-nope + ``[T_kv, H, 128]``, the shared K-pe ``[T_kv, 1, 64]``, V ``[T_kv, H, 128]``), merged into the + BF16 ``output`` ``[T_q, H, 128]`` and FP32 ``merge_stats`` ``[T_q, H, 2]`` (row max, row sum, + natural log) per request by ``merge_op``: 2 writes, 1 merges, 0 leaves both unchanged.""" + + ARCHS = frozenset({Arch.SM_107}) + # Spec: per-request `q_lens` / `kv_lens` of one partition, `causal`, the per-request + # `merge_ops`, and the `cu_kv` dtype (int64: a cached-prefix partition's cu_chunked_seq_len; + # int32: the current chunk, whose cu_kv is cu_q). 96 heads: the tp1 context server. + CELLS: tuple[Cell, ...] = ( + Cell( + why="tp1 cached prefix, first partition: two requests write output and stats (op 2)", + spec=dict( + q_lens=[257, 17], kv_lens=[384, 128], causal=False, merge_ops=[2, 2], cu_kv="int64" + ), + ), # fmt: skip + Cell( + why="tp1 cached prefix, later partition: merge (op 1) beside a request without one (0)", + spec=dict( + q_lens=[257, 17], kv_lens=[384, 0], causal=False, merge_ops=[1, 0], cu_kv="int64" + ), + ), # fmt: skip + Cell( + why="tp1 current chunk after a cached prefix: causal, merged into the prefix (op 1)", + spec=dict( + q_lens=[257, 17], kv_lens=[257, 17], causal=True, merge_ops=[1, 1], cu_kv="int32" + ), + ), # fmt: skip + Cell( + why="tp1 first chunk of an 8k prompt: causal over itself, writes (op 2)", + spec=dict(q_lens=[8192], kv_lens=[8192], causal=True, merge_ops=[2], cu_kv="int32"), + ), + Cell( + why="tp1 8k chunk against an 8k cached-prefix partition (op 1)", + spec=dict(q_lens=[8192], kv_lens=[8192], causal=False, merge_ops=[1], cu_kv="int64"), + ), + ) + + note = """ + The chunked prefill of an NVFP4-KV MLA layer runs this once per cached-prefix partition + (non-causal, the partition's `merge_op_tensor` row) and once for the current chunk (causal, + top-left aligned mask: query i of a request sees keys 0..i of its partition). Writes `output` + and `merge_stats` in place and returns None. + + Inputs are unit-scale E4M3; the kernel rounds the softmax probabilities to E4M3 (prescale 256) + before the PV product, so a causal row over few keys differs from the FP32 reference by up to + half an E4M3 step of P times |V|. The stats are (row max, row sum) of exp(score - max) in + natural-log units; only max + log(sum) is determined (the kernel may keep another max). + + A process cache holds one compiled kernel per (Q dtype, heads, causal, cu_kv_seqlens dtype); + the first call of a key compiles (not capturable). `cu_q_seqlens` int32, `cu_kv_seqlens` int32 + or int64, both `[B + 1]`; `merge_op` int64 `[B]`. K-nope / V need unit head stride and a + 128-element head pitch (views of the kv_b output are fine). SM107 only; the launcher raises + ValueError on any other device or layout. + """ + + def __call__( + self, + q: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + merge_stats: torch.Tensor, + merge_op: torch.Tensor, + max_q_len: int, + max_kv_len: int, + softmax_scale: float, + is_causal: bool, + ) -> None: + _fmha.run_k3_context_fmha( + q, + k_nope, + k_pe, + v, + output, + cu_q_seqlens, + cu_kv_seqlens, + merge_stats, + merge_op, + max_q_len, + max_kv_len, + softmax_scale, + is_causal, + ) + + def reference( + self, + q: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + merge_stats: torch.Tensor, + merge_op: torch.Tensor, + max_q_len: int, + max_kv_len: int, + softmax_scale: float, + is_causal: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The post-state ``(output, merge_stats)``: FP32 attention of every request over its + partition, merged into the pre-call state per ``merge_op``. Reads ``output`` and + ``merge_stats`` (the running state) and writes neither; ``max_q_len`` / ``max_kv_len`` + are launch bounds only.""" + heads = q.shape[1] + out = output.float().clone() + stats = merge_stats.clone() + cu_q = cu_q_seqlens.tolist() + cu_kv = cu_kv_seqlens.tolist() + ops = merge_op.tolist() + with true_fp32_matmul(): + for b, op in enumerate(ops): + if op == 0: + continue + q0, q1, k0, k1 = cu_q[b], cu_q[b + 1], cu_kv[b], cu_kv[b + 1] + if q1 == q0: + continue + for h0 in range(0, heads, _REF_HEADS): + hs = slice(h0, min(h0 + _REF_HEADS, heads)) + qb = q[q0:q1, hs].float() + scores = torch.einsum( + "qhd,khd->hqk", qb[..., :128], k_nope[k0:k1, hs].float() + ) + torch.einsum("qhd,kd->hqk", qb[..., 128:], k_pe[k0:k1, 0].float()) + scores *= softmax_scale + if is_causal: + rows = torch.arange(q1 - q0, device=q.device)[:, None] + cols = torch.arange(k1 - k0, device=q.device)[None, :] + scores.masked_fill_((cols > rows)[None], -torch.inf) + cur_max = scores.amax(-1) + weights = (scores - cur_max[..., None]).exp() + cur_sum = weights.sum(-1) + cur_out = torch.einsum("hqk,khd->hqd", weights, v[k0:k1, hs].float()) + cur_out = (cur_out / cur_sum[..., None]).transpose(0, 1) + cur_max, cur_sum = cur_max.transpose(0, 1), cur_sum.transpose(0, 1) + if op == 1: + cur_out, cur_max, cur_sum = _merge( + out[q0:q1, hs], + stats[q0:q1, hs, 0], + stats[q0:q1, hs, 1], + cur_out, + cur_max, + cur_sum, + ) + out[q0:q1, hs] = cur_out + stats[q0:q1, hs, 0] = cur_max + stats[q0:q1, hs, 1] = cur_sum + return out.to(output.dtype), stats + + def is_valid( + self, + q: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + merge_stats: torch.Tensor, + merge_op: torch.Tensor, + max_q_len: int, + max_kv_len: int, + softmax_scale: float, + is_causal: bool, + ) -> None: + # The kernel tiles queries and keys up to these bounds: a request longer than them + # is attended only partly, with no error. + q_lens = cu_q_seqlens.diff() + kv_lens = cu_kv_seqlens.diff() + assert int(q_lens.max()) <= max_q_len, "max_q_len below a request's query count" + assert int(kv_lens.max()) <= max_kv_len, "max_kv_len below a request's key count" + # A merge code outside {0, 1, 2} selects no defined epilogue. + assert bool(((merge_op >= 0) & (merge_op <= 2)).all()), "merge_op codes are 0, 1, 2" + + +k3_context_fmha = K3ContextFmha() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_fp4_mla_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_fp4_mla_decode.py new file mode 100644 index 000000000000..1ec6ec471fd4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_fp4_mla_decode.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""K3's packed FP4 MLA generation decode (CuTe DSL decode kernel and its split reduction).""" + +from __future__ import annotations + +from typing import Any + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from . import _fp4_mla_decode as _decode + + +class K3Fp4MlaDecode(OpWrapper): + """``run_generation``: the absorbed generation attention of the FP4 MLA layer behind ``attn`` + (its ``TrtllmAttention``) over the FP4 cache, for the fused query rows the generation cache + update (or ``k3_mla_q_fused``, ``q_rows_external``) wrote; ``q`` ``[T, H * 576]`` provides + shapes. Returns ``[T, H * 512]`` in ``q``'s dtype; ``softmax_stats_tensor`` (Helix) receives + the FP32 ``[T, H, 2]`` (max, sum) statistics.""" + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = () + + note = """ + Per call: the packed split plan (once per step, fused with the generation lengths), the + device page table (once per step), the packed decode kernel over (split, request) tasks and + the split reduction. Inside a `DeferredPackedReduce` context the reduction is recorded instead + of launched (the Helix tail folds the partials); the rows of single-task requests are then the + only ones written. Plan, partials and maxima live in workspaces on `metadata.fp4_mla_state`; + `begin_step` forgets the plan of the previous step and `prelaunch_generation_metadata` may + launch the plan and page table early on another stream (the result does not change). + `features` fields select FP16 partials, the reduction's PDL launch, the page plan and the + decode's PDL wait after the fused query. Consumes the prequantized query on the state. + """ + + def __call__( + self, + attn: Any, + q: torch.Tensor, + metadata: Any, + softmax_stats_tensor: torch.Tensor | None = None, + q_rows_external: bool = False, + ) -> torch.Tensor: + return _decode.run_generation( + attn, + q, + metadata, + softmax_stats_tensor=softmax_stats_tensor, + q_rows_external=q_rows_external, + ) + + def reference( + self, + attn: Any, + q: torch.Tensor, + metadata: Any, + softmax_stats_tensor: torch.Tensor | None = None, + q_rows_external: bool = False, + ) -> torch.Tensor: + """Not modelled yet: attention of the dequantized prequantized query rows over each + request's dequantized FP4 cache rows, the HP ring tail and the packed FP4 V copy, all in + layouts the KV cache manager and ``metadata.fp4_mla_state`` own. Until it is, the entry + has no cells; the kernel is checked against the FP4 MLA backend's decode + (``test_k3_fp4_mla_decode.py``).""" + raise NotImplementedError("K3Fp4MlaDecode.reference: no plain-torch model yet") + + +k3_fp4_mla_decode = K3Fp4MlaDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_helix_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_helix_tail.py new file mode 100644 index 000000000000..f08b2795bf0f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_helix_tail.py @@ -0,0 +1,337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 Helix generation tail in one launch: split-KV fold, V up-projection, exchange, combine.""" + +from __future__ import annotations + +import math +from typing import Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_helix_tail as _tail + +_LOG2_E = math.log2(math.e) + + +def packed_reduce_reference( + fold: _tail.FoldInputs, rows: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """FP64 fold of the packed decode's split state: ``(latent BF16 [requests * valid_m, 512], + stats FP32 [requests * valid_m, 2])``, natural-log ``(max, sum)``. + + A request planned on one task keeps the row the decode stored in ``rows`` and that task's + statistics. FP32 partials are max-relative numerators (``output_scale`` converts them), FP16 + partials normalized split outputs.""" + plan = fold.plan.tolist() + begins = plan[2 : 3 + fold.num_requests] + valid_m = fold.valid_m + scale = fold.sm_scale * _LOG2_E / (fold.q_scale.double().item() * fold.kv_scale.double().item()) + normalized = fold.partial.dtype == torch.float16 + latent, maxima, sums = [], [], [] + for r, (begin, end) in enumerate(zip(begins, begins[1:])): + row_sum = fold.sums[begin:end, :valid_m].double() + active = row_sum != 0 + row_max = fold.maxima[begin:end, :valid_m].double() * scale + global_max = row_max.masked_fill(~active, -torch.inf).amax(0) + weights = torch.where(active, torch.exp2(row_max - global_max), 0.0) + values = fold.partial[begin:end, :valid_m].double() + if normalized: + values = values * row_sum.unsqueeze(-1) + numerator = (torch.where(active.unsqueeze(-1), values, 0.0) * weights.unsqueeze(-1)).sum(0) + denominator = (row_sum * weights).sum(0) + out = torch.where( + denominator.unsqueeze(-1) != 0, + numerator / denominator.unsqueeze(-1) * fold.output_scale, + 0.0, + ) + if end - begin == 1: + out = rows[r * valid_m : (r + 1) * valid_m].double() + latent.append(out) + maxima.append(torch.where(denominator != 0, global_max / _LOG2_E, -torch.inf)) + sums.append(denominator) + stats = torch.stack((torch.cat(maxima), torch.cat(sums)), -1).float() + return torch.cat(latent).to(rows.dtype), stats + + +def _combine(values: Sequence[torch.Tensor], stats: Sequence[torch.Tensor]) -> torch.Tensor: + """``helix_post_process_native`` in FP64: the stats-weighted mean of the sources' rows.""" + maxima = torch.stack([s[..., 0].double() for s in stats]) + sums = torch.stack([s[..., 1].double() for s in stats]) + global_max = maxima.amax(0) + weights = torch.where(sums != 0, sums * (maxima - global_max).exp(), 0.0) + num = sum(w.unsqueeze(-1) * v.double() for w, v in zip(weights, values)) + return num / weights.sum(0).unsqueeze(-1) + + +def _mxfp8_reference(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + from ...kernels.k3_fp8w_lowm_mxfp8 import mxfp8_quantize_reference + + return mxfp8_quantize_reference(x) + + +class K3HelixTailOp(OpWrapper): + """``k3_helix_tail``: this rank's BF16 ``out`` ``[T, H / P * 128]`` from the decode's latent + rows of all ``H`` heads (``rows`` ``[T * H, 512]`` with ``stats`` ``[T * H, 2]``, or the split + state ``fold``), the per-head V up-projection ``w_uv`` ``[H, 128, 512]``, the exchange over the + ``P`` CP ranks and the stats-weighted combine. Writes ``out`` (and ``mx_out``'s ``xq`` / + ``xs``) in place and returns None.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # Spec: `cp_size` (P; > 1 runs on one GPU with `loopback`), the decode's split state + # (`kv_tiles` per request, `splits` task slots, `wave_tasks`, `query_len`; T = requests * + # query_len), `partial` ("fp32", "fp16" normalized, or None: rows stored directly), `ull2`, + # `mx_out` (o_proj's MXFP8 operand). H = 96 heads per rank, latent 512, V head 128. + CELLS: tuple[Cell, ...] = ( + Cell( + why="helix8 decode: P = 8 (12 heads out), FP32 split partials folded, 8 requests x 4", + spec=dict( + cp_size=8, + kv_tiles=[40] * 8, + splits=8, + wave_tasks=106, + query_len=4, + partial="fp32", + ull2=False, + mx_out=False, + ), + ), # fmt: skip + Cell( + why="helix8 decode: one 21.9k-token request split over 8 tasks", + spec=dict( + cp_size=8, + kv_tiles=[86], + splits=8, + wave_tasks=106, + query_len=4, + partial="fp32", + ull2=False, + mx_out=False, + ), + ), # fmt: skip + Cell( + why="helix16 ULL: P = 16 (6 heads out), FP16 normalized partials, staged (ull2) fold", + spec=dict( + cp_size=16, + kv_tiles=[86], + splits=8, + wave_tasks=106, + query_len=4, + partial="fp16", + ull2=True, + mx_out=False, + ), + ), # fmt: skip + Cell( + why="helix16 ULL above 16 tokens: the gate joins first, o_proj's MXFP8 operand stored", + spec=dict( + cp_size=16, + kv_tiles=[60, 9, 33, 1, 86, 20, 5, 40], + splits=8, + wave_tasks=106, + query_len=4, + partial="fp16", + ull2=True, + mx_out=True, + ), + ), # fmt: skip + Cell( + why="helix16 ULL DSpark drafter: rows the attention stored directly, no fold", + spec=dict( + cp_size=16, + kv_tiles=[1] * 2, + splits=1, + wave_tasks=106, + query_len=8, + partial=None, + ull2=True, + mx_out=False, + ), + ), # fmt: skip + Cell( + why="one rank (P = 1): single-task requests between multi-task ones, rows copied", + spec=dict( + cp_size=1, + kv_tiles=[1, 79, 3, 176], + splits=6, + wave_tasks=17, + query_len=8, + partial="fp32", + ull2=False, + mx_out=False, + ), + ), # fmt: skip + ) + + note = """ + Collective over the CP group whose receive buffers are the MNNVL workspace at `ws_addr` + (rank r's at `ws_addr + r * rank_stride`, two halves of `half_bytes`, armed with the sentinel); + `state` is the launch counter of that workspace. Launches sharing a workspace must run in the + same order on one stream of every rank. `loopback`: one GPU stands for the group (records of + head group d are sent to this rank as source d), for single-GPU tests. + + Numerics: the fold and the combine are bitwise the packed split reduction and + `helix_post_process_native`; the V up-projection rounds FP32 dot products in another summation + order than the BMM (within one BF16 ulp of the FP64 product). Rows of tokens a rank owns no KV + for (`zero_kv_mask`) contribute nothing. `mx_out = (gate, xq, xs)` also stores the MXFP8 form + of `bf16(out * gate)` (o_proj's operand, bitwise the gated low-M GEMM's quantization). + + A process cache holds compiled kernels per feature combination; one CTA per SM at most + (`max_ctas`). `ull2` selects the staged fold (bitwise the same; always for FP16 partials). + `timestamps` (microbenchmarks) stores %globaltimer phase stamps. The receive buffer and `state` + are internal: the reference does not model them. The cells run one GPU (`loopback`), so no + cell masks a token (every source of a looped-back token is the same rank). + """ + + def __call__( + self, + rows: torch.Tensor, + stats: torch.Tensor | None, + w_uv: torch.Tensor, + out: torch.Tensor, + fold: _tail.FoldInputs | None, + zero_kv_mask: torch.Tensor | None, + cp_rank: int, + cp_size: int, + ws_addr: int, + rank_stride: int, + half_bytes: int, + state: torch.Tensor, + max_ctas: int, + use_pdl: bool, + mx_out: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None, + timestamps: torch.Tensor | None, + loopback: bool, + mx_atom: bool, + ull2: bool, + ) -> None: + _tail.k3_helix_tail( + rows, + stats, + w_uv, + out, + fold=fold, + zero_kv_mask=zero_kv_mask, + cp_rank=cp_rank, + cp_size=cp_size, + ws_addr=ws_addr, + rank_stride=rank_stride, + half_bytes=half_bytes, + state=state, + max_ctas=max_ctas, + use_pdl=use_pdl, + mx_out=mx_out, + timestamps=timestamps, + loopback=loopback, + mx_atom=mx_atom, + ull2=ull2, + ) + + def reference( + self, + rows: torch.Tensor | Sequence[torch.Tensor], + stats: torch.Tensor | Sequence[torch.Tensor] | None, + w_uv: torch.Tensor, + out: torch.Tensor, + fold: _tail.FoldInputs | Sequence[_tail.FoldInputs] | None, + zero_kv_mask: torch.Tensor | Sequence[torch.Tensor] | None, + cp_rank: int, + cp_size: int, + ws_addr: int, + rank_stride: int, + half_bytes: int, + state: torch.Tensor, + max_ctas: int, + use_pdl: bool, + mx_out: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None, + timestamps: torch.Tensor | None, + loopback: bool, + mx_atom: bool, + ull2: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The post-state ``out`` (and ``(out, xq, xs)`` with ``mx_out``), FP64-accumulated. + + A collective: with ``cp_size > 1`` and not ``loopback``, ``rows``, ``stats``, ``fold`` and + ``zero_kv_mask`` are the group's inputs in rank order (``w_uv`` is replicated). With + ``loopback`` (or ``P = 1``) they are this rank's, and source ``d`` is this rank's head + group ``d``. The workspace arguments, ``state``, ``max_ctas``, ``use_pdl``, ``timestamps`` + and ``ull2`` do not change the result. ``mx_atom`` (another scale layout) is not + modelled: its cells are none.""" + heads = w_uv.shape[0] + group = heads // cp_size + per_rank = cp_size > 1 and not loopback + + def as_list(value): + if per_rank: + return list(value) if value is not None else [None] * cp_size + return [value] + + sources = [] + for r_rows, r_stats, r_fold, r_mask in zip( + as_list(rows), as_list(stats), as_list(fold), as_list(zero_kv_mask) + ): + if r_fold is not None: + r_rows, r_stats = packed_reduce_reference(r_fold, r_rows) + tokens = r_rows.shape[0] // heads + latent = r_rows.view(tokens, heads, -1).double() + values = torch.einsum("tjk,jnk->tjn", latent, w_uv.double()).to(torch.bfloat16) + r_stats = r_stats.view(tokens, heads, 2).clone() + if r_mask is not None: + dead = r_mask[:tokens].bool() + values[dead] = 0 + r_stats[dead] = torch.tensor([-torch.inf, 0.0], device=r_stats.device) + sources.append((values, r_stats)) + if per_rank: + heads_of = slice(cp_rank * group, (cp_rank + 1) * group) + values = [v[:, heads_of] for v, _ in sources] + stats_by_source = [s[:, heads_of] for _, s in sources] + else: + v, s = sources[0] + values = [v[:, d * group : (d + 1) * group] for d in range(cp_size)] + stats_by_source = [s[:, d * group : (d + 1) * group] for d in range(cp_size)] + combined = _combine(values, stats_by_source) + result = combined.reshape(out.shape).to(out.dtype) + if mx_out is None: + return result + gate = mx_out[0] + xq, xs = _mxfp8_reference(result * gate) + return result, xq, xs + + def is_valid( + self, + rows: torch.Tensor, + stats: torch.Tensor | None, + w_uv: torch.Tensor, + out: torch.Tensor, + fold: _tail.FoldInputs | None, + zero_kv_mask: torch.Tensor | None, + cp_rank: int, + cp_size: int, + ws_addr: int, + rank_stride: int, + half_bytes: int, + state: torch.Tensor, + max_ctas: int, + use_pdl: bool, + mx_out: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None, + timestamps: torch.Tensor | None, + loopback: bool, + mx_atom: bool, + ull2: bool, + ) -> None: + # Every CTA polls after its own pushes: more CTAs than SMs can leave a sender unscheduled + # behind a spinning receiver. + sms = torch.cuda.get_device_properties(out.device).multi_processor_count + assert max_ctas <= sms, f"max_ctas {max_ctas} exceeds the {sms} SMs (all CTAs co-resident)" + # Records that overrun a receive half land in the other half or past the workspace. + tokens = out.shape[0] + need = tokens * cp_size * _tail.record_bytes(w_uv.shape[0] // cp_size, w_uv.shape[1]) + assert need <= half_bytes, f"{need} bytes of records exceed a {half_bytes}-byte half" + # A mask shorter than the batch is read past its end. + if zero_kv_mask is not None: + assert zero_kv_mask.numel() >= tokens and zero_kv_mask.is_contiguous() + + +k3_helix_tail = K3HelixTailOp() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_mla_latent_row.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_mla_latent_row.py new file mode 100644 index 000000000000..b01a918918ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_mla_latent_row.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The latent row of the MLA KV half: kv_a RMSNorm of ``c_kv`` and the ``k_pe`` copy, one launch.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_mla_q_fused as _kernels + + +class K3MlaLatentRow(OpWrapper): + """``torch.cat([rmsnorm(c_kv, weight, eps), k_pe], -1)``: BF16 ``c_kv`` ``[T, 512]`` and + ``k_pe`` ``[T, 64]`` (row strides multiples of 8, unit column stride) into a new contiguous + BF16 ``[T, 576]``.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # Spec: `tokens` T <= 16 (the fused query's batches: generation-only Helix, FP4 KV); c_kv and + # k_pe are the kv_a and K-pe thirds of the projection output [T, 2112]; eps 1e-6. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=why, spec=dict(tokens=tokens, row_stride=2112, eps=1e-6)) + for why, tokens in ( + ("helix16 ULL c1: one token, one warp of one CTA", 1), + ("helix16 ULL c2 with draft tokens: a partial last CTA", 6), + ("helix8 / helix16 fused-query cap: 16 tokens, four full CTAs", 16), + ) + ) + + note = """ + One warp per token with the arithmetic of FlashInfer's CuTe DSL `RMSNormKernel` for a + 512-wide row, so the output is bitwise FlashInfer's norm followed by the copy. The caller + takes it only where `latent_row_supported` holds and FlashInfer's norm is its CuTe DSL kernel. + PDL: the weight is read before `griddepcontrol.wait` and dependents are released right after + it. A process cache holds one compiled kernel per `use_pdl`; compile before graph capture. + """ + + def __call__( + self, + c_kv: torch.Tensor, + k_pe: torch.Tensor, + weight: torch.Tensor, + eps: float, + use_pdl: bool = True, + ) -> torch.Tensor: + return _kernels.k3_mla_latent_row(c_kv, k_pe, weight, eps, use_pdl=use_pdl) + + def reference( + self, + c_kv: torch.Tensor, + k_pe: torch.Tensor, + weight: torch.Tensor, + eps: float, + use_pdl: bool = True, + ) -> torch.Tensor: + """FP32 RMSNorm rounded once, then ``k_pe``; ``use_pdl`` does not change the result.""" + x = c_kv.float() + normed = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * weight.float() + return torch.cat([normed.to(c_kv.dtype), k_pe], dim=-1) + + def is_valid( + self, + c_kv: torch.Tensor, + k_pe: torch.Tensor, + weight: torch.Tensor, + eps: float, + use_pdl: bool = True, + ) -> None: + # The launcher checks nothing: it reads 512 / 64 columns at the given row strides with + # 16-byte loads whatever the shapes, so an operand outside these is read wrongly. + assert _kernels.latent_row_supported(c_kv, k_pe, weight), ( + "c_kv [T, 512] / k_pe [T, 64] / weight [512] BF16, row strides % 8, 16-byte aligned" + ) + + +k3_mla_latent_row = K3MlaLatentRow() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_mla_q_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_mla_q_fused.py new file mode 100644 index 000000000000..d15cddd2a518 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/k3_mla_q_fused.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 FP4 MLA decode query in one kernel: q_a RMSNorm, q_b, W_UK, RoPE, FP4 quantization.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) + +from ...kernels import k3_mla_q_fused as _kernels + +#: E2M1 magnitudes by code (bit 3 is the sign). +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +#: Smallest FP4 group scale the query quantization uses (2^-9). +_SCALE_FLOOR = 2.0**-9 + + +def swizzled_sf_offset(row: torch.Tensor, col: torch.Tensor, cols: int) -> torch.Tensor: + """Byte offset of scale ``(row, col)`` in the 128x4-swizzled layout of ``cols`` columns.""" + padded = (cols + 3) // 4 * 4 + return ( + col % 4 + + col // 4 * (4 * 128) + + row % 32 * 16 + + row % 128 // 32 * 4 + + row // 128 * (128 * padded) + ) + + +def e2m1_values(codes: torch.Tensor) -> torch.Tensor: + """FP32 values of 4-bit E2M1 codes.""" + table = torch.tensor(_E2M1, dtype=torch.float32, device=codes.device) + mag = table[(codes & 7).long()] + return torch.where((codes & 8) != 0, -mag, mag) + + +def e2m1_round(v: torch.Tensor) -> torch.Tensor: + """E2M1 codes of FP32 ``v``: round to nearest, ties to the even code, saturating at 6.""" + table = torch.tensor(_E2M1, dtype=torch.float32, device=v.device) + mag = v.abs().clamp_max(6.0) + lower = (torch.bucketize(mag, table, right=True) - 1).clamp(0, 7) + upper = (lower + 1).clamp_max(7) + d_lo = mag - table[lower] + d_hi = table[upper] - mag + take_hi = (d_hi < d_lo) | ((d_hi == d_lo) & (upper % 2 == 0) & (upper != lower)) + code = torch.where(take_hi, upper, lower) + return (code | torch.where(v < 0, 8, 0)).to(torch.uint8) + + +def fp4_quantize_groups(v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """FP4 groups of 16 (last dim 16): ``(codes uint8 [..., 16], scale byte uint8 [...], scale + FP32 [...])`` with the scale ``e4m3(max(amax / 6, 2^-9))`` and codes ``e2m1(v / scale)``.""" + amax = v.abs().amax(-1) + scale = (amax.double() / 6.0).clamp(_SCALE_FLOOR, 448.0).to(torch.float8_e4m3fn) + scale_f32 = scale.float() + codes = e2m1_round((v.double() / scale_f32.double().unsqueeze(-1)).float()) + return codes, scale.view(torch.uint8), scale_f32 + + +def pack_e2m1(codes: torch.Tensor) -> torch.Tensor: + """Pack codes (last dim even) two per byte, element 2k in the low nibble of byte k.""" + pairs = codes.view(*codes.shape[:-1], -1, 2).to(torch.int32) + return (pairs[..., 0] | (pairs[..., 1] << 4)).to(torch.uint8) + + +def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: + """Inverse of :func:`pack_e2m1`.""" + lo = packed & 15 + hi = packed >> 4 + return torch.stack((lo, hi), -1).flatten(-2) + + +class K3MlaQFused(OpWrapper): + """``k3_mla_q_fused``: the packed FP4 decode query of ``T <= 16`` tokens of 96 heads, written + into the decode's Q buffers ``q_fp4`` (uint8 ``[rows, 320]``) and ``q_sf`` (swizzled E4M3 + scales). Returns None.""" + + ARCHS = frozenset({Arch.SM_107}) + # Spec: `requests` x `gen_len` tokens (T <= 16), `norm` (the q_a RMSNorm fused: FlashInfer's + # CuTe DSL norm is the one replaced), `early_trigger` (features.ull2_mla_decode_pdl). The + # rotary table is K3's identity table and the query global scale the FP4 MLA backend's. + CELLS: tuple[Cell, ...] = ( + Cell( + why="helix16 ULL c1: one request of 4 tokens, fused norm, decode released at entry", + spec=dict(requests=1, gen_len=4, norm=True, early_trigger=True), + ), + Cell( + why="helix16 ULL c4: four requests of 4 tokens (T = 16, the 16-token compile)", + spec=dict(requests=4, gen_len=4, norm=True, early_trigger=True), + ), + Cell( + why="helix8: two requests of 8 tokens, decode released after the query stores", + spec=dict(requests=2, gen_len=8, norm=True, early_trigger=False), + ), + Cell( + why="one token, q_a already normalized (FlashInfer's CUDA norm ran before)", + spec=dict(requests=1, gen_len=1, norm=False, early_trigger=False), + ), + ) + + note = """ + x: BF16 [T, 1536] q_a rows (raw with `norm = (weight, eps)`: the kernel applies the q_a + RMSNorm with FlashInfer's arithmetic; else already normalized). w_qb: FP8 [96 * 192, 1536] + with block exponents `w_qb_exp` (`block_exponents`); w_uk: BF16 [96, 512, 128]. + Token t of request t // gen_len owns row (t // gen_len) * rows_per_seq + (t % gen_len) * 96 + + head: 256 bytes of E2M1 absorbed values (global scale `q_global_scale`), then 32 bytes of RoPE + values and 32 of their residual; one E4M3 scale per 16 values in the 128x4 swizzled layout. + A negative position skips the token's RoPE part. + + The q_b product follows the low-M FP8 GEMM (MXFP8 activation, 128x128 UE8M0 weight blocks, + FP32 per 32-value group, the two K halves added) up to the MMA's internal order; W_UK is a + BF16 MMA with FP32 accumulation; RoPE and the FP4 quantization repeat the FP4 MLA cache + update's operations (`div.full` / `rcp.approx`). Bitwise the production chain in the kernel's + own tests. Writes only the rows of the given tokens. A process cache holds one compiled + kernel per (T padded to 8 / 16, use_pdl, norm, early_trigger). `early_trigger` releases the + dependent decode at entry (it reads only the query after its own wait). + """ + + def __call__( + self, + x: torch.Tensor, + w_qb: torch.Tensor, + w_qb_exp: torch.Tensor, + w_uk: torch.Tensor, + q_global_scale: torch.Tensor, + cos_sin: torch.Tensor, + positions: torch.Tensor, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + gen_len: int, + rows_per_seq: int, + use_pdl: bool, + norm: tuple[torch.Tensor, float] | None, + early_trigger: bool, + ) -> None: + _kernels.k3_mla_q_fused( + x, + w_qb, + w_qb_exp, + w_uk, + q_global_scale, + cos_sin, + positions, + q_fp4, + q_sf, + gen_len=gen_len, + rows_per_seq=rows_per_seq, + use_pdl=use_pdl, + norm=norm, + early_trigger=early_trigger, + ) + + def reference( + self, + x: torch.Tensor, + w_qb: torch.Tensor, + w_qb_exp: torch.Tensor, + w_uk: torch.Tensor, + q_global_scale: torch.Tensor, + cos_sin: torch.Tensor, + positions: torch.Tensor, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + gen_len: int, + rows_per_seq: int, + use_pdl: bool, + norm: tuple[torch.Tensor, float] | None, + early_trigger: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The post-state ``(q_fp4, q_sf)`` (``q_sf`` as uint8): the rows of the given tokens + rewritten, every other byte as passed. FP64 q_b product of the MXFP8 activation + (``mxfp8_quantize_reference``) with the block-scaled weight, rounded to BF16; FP64 W_UK + absorption rounded to BF16; GPT-J RoPE in FP32 with the pair-``p`` table values + ``cos_sin[pos, 2p]``, ``[2p + 1]``, rounded to BF16; then FP4 groups of 16 of the values + times ``q_global_scale`` (exact division where the kernel uses ``div.full`` / + ``rcp.approx``), the RoPE groups with their residual. ``use_pdl`` / ``early_trigger`` + do not change the result; the cells use positions >= 0.""" + from ...kernels.k3_fp8w_lowm_mxfp8 import mxfp8_quantize_reference + + heads, nope, rope, lora = ( + _kernels.NUM_HEADS, + _kernels.NOPE_DIM, + _kernels.ROPE_DIM, + _kernels.KV_LORA_RANK, + ) + tokens = x.shape[0] + out_q, out_sf = q_fp4.clone(), q_sf.clone().view(torch.uint8) + if tokens == 0: + return out_q, out_sf + rows = x if norm is None else flashinfer_rmsnorm.reference(x, norm[0], norm[1]) + xq, xs = mxfp8_quantize_reference(rows.contiguous()) + a = xq.double() * torch.exp2(xs.double() - 127).repeat_interleave(32, -1) + n, k = w_qb.shape + w_scale = torch.exp2(w_qb_exp.double() - 127).repeat_interleave(128, 0)[:n] + w = w_qb.double() * w_scale.repeat_interleave(128, 1)[:, :k] + q = (a @ w.T).to(torch.bfloat16).view(tokens, heads, nope + rope) + absorbed = torch.einsum("tjd,jcd->tjc", q[..., :nope].double(), w_uk.double()) + absorbed = absorbed.to(torch.bfloat16) + table = cos_sin.view(-1, 2 * rope)[positions[:tokens].long()].view(tokens, 1, -1, 2) + cos, sin = table[:, :, : rope // 2, 0], table[:, :, : rope // 2, 1] + x0, x1 = q[..., nope:].float().view(tokens, heads, rope // 2, 2).unbind(-1) + even = (x0 * cos).sub(x1 * sin) + odd = (x1 * cos).add(x0 * sin) + rotated = torch.stack((even, odd), -1).view(tokens, heads, rope).to(torch.bfloat16) + gs = q_global_scale.float().view(()) + + prefix = (absorbed.float() * gs).view(tokens, heads, lora // 16, 16) + p_codes, p_sf, _ = fp4_quantize_groups(prefix) + pe = (rotated.float() * gs).view(tokens, heads, rope // 16, 16) + m_codes, m_sf, m_sf_f32 = fp4_quantize_groups(pe) + main = e2m1_values(m_codes).double() * m_sf_f32.double().unsqueeze(-1) + residual = (pe.double() - main).float() + r_codes, r_sf, _ = fp4_quantize_groups(residual) + + rope_groups = rope // 16 + packed = torch.cat( + ( + pack_e2m1(p_codes.flatten(-2)), + torch.cat((pack_e2m1(m_codes), pack_e2m1(r_codes)), -1).flatten(-2), + ), + -1, + ) + sf_bytes = torch.cat( + (p_sf, torch.stack((m_sf, r_sf), -1).view(tokens, heads, 2 * rope_groups)), -1 + ) + t = torch.arange(tokens, device=x.device) + seq = t // gen_len + token_rows = seq * rows_per_seq + (t - seq * gen_len) * heads + q_rows = (token_rows[:, None] + torch.arange(heads, device=x.device)[None, :]).flatten() + out_q[q_rows] = packed.view(-1, packed.shape[-1]) + cols = sf_bytes.shape[-1] + col = torch.arange(cols, device=x.device) + offsets = swizzled_sf_offset(q_rows[:, None], col[None, :], cols) + out_sf.view(-1)[offsets.flatten()] = sf_bytes.reshape(-1) + return out_q, out_sf + + def is_valid( + self, + x: torch.Tensor, + w_qb: torch.Tensor, + w_qb_exp: torch.Tensor, + w_uk: torch.Tensor, + q_global_scale: torch.Tensor, + cos_sin: torch.Tensor, + positions: torch.Tensor, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + gen_len: int, + rows_per_seq: int, + use_pdl: bool, + norm: tuple[torch.Tensor, float] | None, + early_trigger: bool, + ) -> None: + # The launcher checks only positions, q_fp4 and cos_sin; the weights and x are read by + # TMA / 16-byte loads in the geometry the kernel is built for, whatever their shapes. + assert _kernels.can_run(x, w_qb, w_uk), "x [T <= 16, 1536] BF16, K3 q_b / W_UK weights" + tokens = x.shape[0] + assert tokens % gen_len == 0, "every request has gen_len tokens" + assert rows_per_seq >= gen_len * _kernels.NUM_HEADS, "rows_per_seq below gen_len * 96" + assert q_fp4.shape[0] >= tokens // gen_len * rows_per_seq, "q_fp4 lacks rows" + + +k3_mla_q_fused = K3MlaQFused() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/rmsnorm_release_at_launch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/rmsnorm_release_at_launch.py new file mode 100644 index 000000000000..ac1bf29e1718 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/rmsnorm_release_at_launch.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RMSNorm of the MLA q_a latent that releases its programmatic dependent (q_b) at launch.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) + +from ...kernels.mla import q_norm_pdl as _norm + + +class RmsnormReleaseAtLaunch(OpWrapper): + """``rmsnorm(x, weight, eps)`` of BF16 ``x`` ``[M, H]`` (unit column stride), a new contiguous + BF16 tensor; FlashInfer's CuTe DSL ``RMSNormKernel`` with its PDL trigger at launch.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # Spec: `tokens` T; `x` is the q_a third of the fused q_a / kv_a projection output + # [T, 1536 + 512 + 64] (row stride 2112; one row counts as contiguous), weight [1536], eps 1e-6. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=why, spec=dict(tokens=tokens, width=1536, row_stride=2112, eps=1e-6)) + for why, tokens in ( + ("helix16 ULL c1: one decode token (the contiguous compile)", 1), + ("helix16 ULL c2 / c4 with draft tokens", 8), + ("helix8 decode batch", 64), + ("helix8 c240-sized decode batch", 960), + ) + ) + + note = """ + Bitwise `flashinfer.norm.rmsnorm(x, weight, eps, enable_pdl=True)` (the same kernel code); + only `griddepcontrol.launch_dependents` moves to the CTA's start, so the dependent q_b GEMM + streams its weights while q_a and this norm run. The dependent must read nothing but its + weights before its own `griddepcontrol.wait`. Called only where `q_norm_pdl.can_use` holds + (`features.mla_q_norm_pdl` and PDL on). A process cache holds one compiled kernel per + (dtype, H, SM, x contiguous); compile before graph capture. + """ + + def __call__(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + return _norm.rmsnorm_release_at_launch(x, weight, eps) + + def reference(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """The catalog RMSNorm's FP32 reference (the kernel reduces in FP32, rounds once).""" + return flashinfer_rmsnorm.reference(x, weight, eps) + + +rmsnorm_release_at_launch = RmsnormReleaseAtLaunch() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/static_quantize_e4m3.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/static_quantize_e4m3.py new file mode 100644 index 000000000000..c75eade3eb28 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/static_quantize_e4m3.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 -> FP8 E4M3 per-tensor static quantization of the K3 FP8 context FMHA inputs (Triton).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import fp8_static_quant as _quant + +_E4M3_MAX = 448.0 + + +class StaticQuantizeE4m3(OpWrapper): + """``x * (1 / scale)`` in E4M3 (round to nearest even, saturating), a new FP8 tensor. + + ``rows=False``: ``x`` contiguous BF16 of even size, the output has ``x``'s shape. + ``rows=True``: ``x`` a row-strided BF16 ``[rows, width]`` view (unit column stride, even row + stride and storage offset, power-of-two width >= 2), the output is contiguous ``[rows, width]``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # Spec: `rows` (the op's flag), `tokens`, and the input: `q` (the context Q rows, 96 heads x + # 192, contiguous), `latent_k_pe` (the K-pe columns of latent_cache [T, 576], row stride 576) + # or `k_pe` (a gathered BF16 K-pe partition [T, 64], contiguous). Unit scale, as K3 runs. + CELLS: tuple[Cell, ...] = ( + Cell( + why="tp1 context: Q of a full 8k chunk, 96 heads", + spec=dict(rows=False, tokens=8192, source="q"), + ), + Cell( + why="tp1 context: Q of a short last chunk", + spec=dict(rows=False, tokens=17, source="q"), + ), + Cell( + why="tp1 context: current chunk's K-pe read in place from latent_cache", + spec=dict(rows=True, tokens=8192, source="latent_k_pe"), + ), + Cell( + why="tp1 context: K-pe of a few current tokens, one program", + spec=dict(rows=True, tokens=17, source="latent_k_pe"), + ), + Cell( + why="tp1 context: a BF16-gathered cached-prefix K-pe partition", + spec=dict(rows=True, tokens=3000, source="k_pe"), + ), + ) + + note = """ + The bytes of `torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, scale)` (the C++ + quantizer the attention op's context copy uses): `cvt.rn.satfinite.e4m3` of + `float(x) * (1 / scale)` with an IEEE FP32 reciprocal. `scale` is a one-element FP32 tensor on + `x`'s device, read on the device (no sync). The contiguous form converts 8192 BF16 pairs per + program; the row form 64 rows per program. Triton; needs SM89+ for the E4M3 conversion. + """ + + def __call__(self, x: torch.Tensor, scale: torch.Tensor, rows: bool = False) -> torch.Tensor: + if rows: + return _quant.static_quantize_e4m3_rows(x, scale) + return _quant.static_quantize_e4m3(x, scale) + + def reference(self, x: torch.Tensor, scale: torch.Tensor, rows: bool = False) -> torch.Tensor: + """``x * (1 / scale)`` with the reciprocal rounded once to FP32, then saturated and + rounded to E4M3 (``rows`` only decides the input layout the op reads).""" + inv_scale = (1.0 / scale.double()).float() + y = (x.float() * inv_scale).clamp(-_E4M3_MAX, _E4M3_MAX) + return y.to(torch.float8_e4m3fn).contiguous() + + +static_quantize_e4m3 = StaticQuantizeE4m3() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_attn_res_lamport.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_attn_res_lamport.py new file mode 100644 index 000000000000..182890f97322 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_attn_res_lamport.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 decode attention-residual mix + RMSNorm whose addend is a pushed MNNVL Lamport +all-reduce that the kernel finishes itself (CuTe DSL, ``kernels/k3_attn_res.py``).""" + +from __future__ import annotations + +from typing import Optional, Sequence, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_attn_res as _kernel +from ..norm.k3_attn_res_norm import attn_res_rmsnorm_reference, mxfp8_rows_reference +from .k3_lamport_finish import lamport_sum_reference + + +class K3AttnResLamport(OpWrapper): + """``(updated, output)`` with ``updated = bf16(prefix + addend)`` and ``output`` the RMSNorm of + the attention-residual mix of ``block`` and ``updated``; with ``mxfp8`` also ``(xq, xs)``, the + MXFP8 form of ``output``. ``addend`` is the all-reduce of the call every rank pushed on the + Lamport workspace at ``uc_local`` (this rank's unicast buffer) / ``flags`` (the rotation + state's address) / ``mc_base`` (multicast base): ``finish_world`` ranks, ``finish_chunk`` + ranks per FP32 partial sum, one-shot or ``finish_two_shot`` transport, this ``rank``. + ``updated`` given: the kernel stores the updated prefix there (contiguous, ``prefix``'s shape + and dtype) instead of a new tensor. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The group is the test harness's (2 and 4 GPUs of one node); K3 runs 8 (helix8) and 16 (helix16 ULL). + # cpp_one_shot: the C++ all-reduce whose sums the call reproduces (finish_chunk follows from it). + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=why, + spec=dict( + rows=rows, + candidates=candidates, + hidden=7168, + two_shot=two_shot, + cpp_one_shot=cpp_one_shot, + add_rn=add_rn, + early_poll=early_poll, + entry_trigger=entry_trigger, + mxfp8=mxfp8, + into=into, + rms_eps=1e-6, + out_eps=1e-6, + ), + ) + for why, rows, candidates, two_shot, cpp_one_shot, add_rn, early_poll, entry_trigger, mxfp8, into in ( + ( + "helix8 A_mid finishing the o_proj all-reduce: 16 rows, N=5, one-shot", + 16, + 5, + False, + True, + False, + False, + False, + False, + False, + ), + ( + "helix8 b16 mtp5 verify A_mid: 96 rows, N=9, two-shot", + 96, + 9, + True, + False, + False, + False, + False, + False, + False, + ), + ( + "two-shot transport with the one-shot C++ sums: 24 rows, N=7", + 24, + 7, + True, + True, + False, + False, + False, + False, + False, + ), + ( + "helix16 ULL o_proj sandwich (early poll, entry trigger): 8 rows, N=6", + 8, + 6, + False, + True, + False, + True, + True, + False, + False, + ), + ( + "helix16 ULL row-parallel MoE tail: add_rn, MXFP8, early poll, 32 rows, N=8", + 32, + 8, + False, + True, + True, + True, + True, + True, + False, + ), + ( + "helix16 ULL tail into the snapshot slot: add_rn, MXFP8, 8 rows, N=4", + 8, + 4, + False, + True, + True, + True, + False, + True, + True, + ), + ) + ) + + note = """ + Finishes the pushed call as k3_lamport_finish does (same sums, rotation + advance and re-arm of the previous buffer), so no finish kernel runs before + it; every rank must make this call at the same position of its sequence on + the workspace. The producer must trigger only after its own + griddepcontrol.wait; two-shot needs every rank's CTAs co-resident (grid of + M * cluster CTAs). The addend has the bits of the finished all-reduce. + add_rn: round-to-nearest FP32 add (the MoE residual update), else the C++ + kernels' fast-math add. early_poll: finish the call and compute the mix + before griddepcontrol.wait, store after it (same bits). entry_trigger: + release the dependent at kernel entry. finish_ag (one-shot): the call + carries all-gather rows, addend = bf16(y + bf16(slot sum)); reached only + through the column-parallel MoE tail and not in this entry's cells. + Process-wide compile cache as k3_attn_res_norm (keyed also by world, chunk, + transport, add_rn, early_poll, finish_ag); compile each key before capture. + Applies where kernels/k3_attn_res.can_use_k3_attn_res holds for the + one-addend mix (H 7168, 2 <= N <= 12, contiguous BF16). + """ + + def __call__( + self, + prefix: torch.Tensor, + uc_local: int, + flags: int, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + finish_world: int, + finish_chunk: int, + finish_two_shot: bool, + mc_base: int, + rank: int, + add_rn: bool, + early_poll: bool = False, + entry_trigger: bool = False, + finish_ag: bool = False, + mxfp8: bool = False, + updated: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, ...]: + if updated is not None: + return _kernel.k3_attn_res_lamport_add_rmsnorm_into( + updated, + mxfp8, + prefix, + uc_local, + flags, + block, + res_weight, + rms_weight, + out_weight, + rms_eps, + out_eps, + finish_world, + finish_chunk, + finish_two_shot, + mc_base, + rank, + add_rn, + early_poll, + entry_trigger, + finish_ag, + ) + op = ( + _kernel.k3_attn_res_lamport_add_rmsnorm_mx + if mxfp8 + else _kernel.k3_attn_res_lamport_add_rmsnorm + ) + return op( + prefix, + uc_local, + flags, + block, + res_weight, + rms_weight, + out_weight, + rms_eps, + out_eps, + finish_world, + finish_chunk, + finish_two_shot, + mc_base, + rank, + add_rn, + early_poll, + entry_trigger, + finish_ag, + ) + + def reference( + self, + prefix: torch.Tensor, + inputs_in_rank_order: Sequence[torch.Tensor], + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + finish_world: int, + finish_chunk: int, + finish_two_shot: bool, + mc_base: int, + rank: int, + add_rn: bool, + early_poll: bool = False, + entry_trigger: bool = False, + finish_ag: bool = False, + mxfp8: bool = False, + updated: Optional[torch.Tensor] = None, + ag_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, ...]: + """The collective exception: every rank's pushed rows in rank order replace the workspace + addresses ``uc_local`` / ``flags``; ``ag_rows`` is the all-gathered ``y`` with + ``finish_ag``. The addend is the rank-order sum of ``k3_lamport_finish``'s reference, the + add is FP32 rounded to BF16 (``add_rn`` and the fast-math add agree away from + denormals), then ``k3_attn_res_norm``'s mix and norm. ``updated`` is not read; the + scheduling flags, ``finish_two_shot``, ``mc_base`` and ``rank`` do not change values.""" + addend = lamport_sum_reference(inputs_in_rank_order, finish_chunk) + if finish_ag: + addend = (ag_rows.float() + addend.float()).to(torch.bfloat16) + row = (prefix.float() + addend.float()).to(torch.bfloat16) + output = attn_res_rmsnorm_reference( + row, block, res_weight, rms_weight, out_weight, rms_eps, out_eps + ) + if mxfp8: + return (row, output, *mxfp8_rows_reference(output)) + return row, output + + def is_valid( + self, + prefix: torch.Tensor, + uc_local: int, + flags: int, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + finish_world: int, + finish_chunk: int, + finish_two_shot: bool, + mc_base: int, + rank: int, + add_rn: bool, + early_poll: bool = False, + entry_trigger: bool = False, + finish_ag: bool = False, + mxfp8: bool = False, + updated: Optional[torch.Tensor] = None, + ) -> None: + # The launcher does not re-validate: operands outside the kernel's shape are read as + # dense rows of 7168 and give a plausible wrong mix. + assert _kernel.can_use_k3_attn_res( + prefix, (prefix,), block, res_weight, rms_weight, out_weight + ), "outside can_use_k3_attn_res (H 7168, 2 <= N <= 12, M <= 1024, contiguous BF16)" + # A rank or chunk that does not match the workspace polls slots nobody writes, or sums in + # another order than the all-reduce the result must reproduce. + assert 0 <= rank < finish_world, f"rank {rank} outside a group of {finish_world}" + assert finish_world % finish_chunk == 0, "finish_chunk must divide finish_world" + + +k3_attn_res_lamport = K3AttnResLamport() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_lamport_finish.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_lamport_finish.py new file mode 100644 index 000000000000..cda5a3286031 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_lamport_finish.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Finish an MNNVL Lamport all-reduce whose input every rank pushed (``k3_lamport_push`` or a +GEMM's push epilogue): poll the slots, sum them in the C++ kernel's order, re-arm the previous +call's buffer and advance the workspace (CuTe DSL, ``kernels/k3_lamport_ar.py``).""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.k3_lamport_ar import LamportWorkspace, reduce_chunk, run_lamport_finish + + +def sanitized(x: torch.Tensor) -> torch.Tensor: + """BF16 ``x`` with ``-0`` replaced by ``+0``: what a push stores into a Lamport slot (the + sentinel word is a ``-0`` pair).""" + return torch.where(x.view(torch.int16) == -32768, torch.zeros_like(x), x) + + +def lamport_sum_reference(inputs_in_rank_order: Sequence[torch.Tensor], chunk: int) -> torch.Tensor: + """BF16 sum of the ranks' pushed rows: FP32 partial sums of ``chunk`` consecutive ranks, each + from zero, added in rank order and rounded once (the C++ MNNVL kernels' order).""" + total: Optional[torch.Tensor] = None + for base in range(0, len(inputs_in_rank_order), chunk): + acc = torch.zeros_like(inputs_in_rank_order[0], dtype=torch.float32) + for part in inputs_in_rank_order[base : base + chunk]: + acc = acc + sanitized(part).float() + total = acc if total is None else total + acc + return total.to(torch.bfloat16) + + +class K3LamportFinish(OpWrapper): + """Write the all-reduce of the call every rank pushed on ``ws`` into ``out`` (contiguous BF16 + ``[M, hidden]``) and return ``out``. ``one_shot``: the transport the producers used; + ``cpp_one_shot``: the C++ kernel whose bits to reproduce (default: the same transport); + ``ag``: the call carries all-gather rows, ``out = bf16(y + bf16(sum))``.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The group is the test harness's (2 and 4 GPUs of one node); K3 runs 8 (helix8) and 16 (helix16 ULL). + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=why, + spec=dict(rows=rows, hidden=7168, one_shot=one_shot, cpp_one_shot=cpp_one_shot), + ) + for why, rows, one_shot, cpp_one_shot in ( + ("one token, one-shot", 1, True, True), + ("helix16 ULL 1-request step: 8 rows, one-shot", 8, True, True), + ("helix16 ULL 4-request step: 32 rows, two-shot", 32, False, False), + ("two-shot transport reproducing the one-shot C++ sums: 24 rows", 24, False, True), + ("one-shot transport reproducing the two-shot C++ result: 16 rows", 16, True, False), + ("helix8 b16 mtp5 verify step: 96 rows, two-shot", 96, False, False), + ) + ) + + note = """ + The second half of a pushed all-reduce. ws (kernels/k3_lamport_ar + .LamportWorkspace) holds the workspace's device addresses: this rank's + unicast buffer, the multicast base, every rank's unicast mapping and the + uint32[9] rotation state the C++ MNNVL kernels share, so pushed and C++ + calls alternate on one workspace. Writes out, re-arms the previous call's + buffer with the sentinel and advances the rotation state. Every rank makes + the same sequence of calls on a workspace, ordered by streams. + One-shot: every rank sums the W slots of each token; two-shot: the owner of + token t (t % W) reduces it and multicasts the row, the others copy it. + Worker CTAs read the state and poll without griddepcontrol.wait (the + producer releases this grid after its own wait); the updater CTA waits for + the producer before it advances the state. Sums in FP32 in the C++ order + (chunks of eight ranks above eight ranks for the one-shot result, up to 16 + in one chunk for the two-shot one); bitwise the C++ all-reduce of the same + inputs. A process-wide cache (`_compiled_finish`) holds one compiled kernel + per (W, hidden, transport, PDL, wait_first, cpp_one_shot, ag); compile each + key eagerly before graph capture. The all-gather variant is reached only + through the column-parallel MoE tail and is not in this entry's cells. + """ + + def __call__( + self, + ws: LamportWorkspace, + out: torch.Tensor, + one_shot: bool, + use_pdl: bool = True, + wait_first: bool = False, + cpp_one_shot: Optional[bool] = None, + ag: bool = False, + ) -> torch.Tensor: + return run_lamport_finish( + ws, + out, + one_shot, + use_pdl=use_pdl, + wait_first=wait_first, + cpp_one_shot=cpp_one_shot, + ag=ag, + ) + + def reference( + self, + inputs_in_rank_order: Sequence[torch.Tensor], + out: torch.Tensor, + one_shot: bool, + use_pdl: bool = True, + wait_first: bool = False, + cpp_one_shot: Optional[bool] = None, + ag: bool = False, + ag_rows: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The collective exception: every rank's pushed ``[M, hidden]`` rows in rank order instead + of ``ws``; ``ag_rows`` the all-gathered ``y`` with ``ag``. ``out`` is not read; + ``use_pdl`` and ``wait_first`` only change scheduling. Returns the expected ``out``.""" + if cpp_one_shot is None: + cpp_one_shot = one_shot + chunk = reduce_chunk(len(inputs_in_rank_order), cpp_one_shot) + total = lamport_sum_reference(inputs_in_rank_order, chunk) + if ag: + total = (ag_rows.float() + total.float()).to(torch.bfloat16) + if not cpp_one_shot: + total = sanitized(total) + return total + + def is_valid( + self, + ws: LamportWorkspace, + out: torch.Tensor, + one_shot: bool, + use_pdl: bool = True, + wait_first: bool = False, + cpp_one_shot: Optional[bool] = None, + ag: bool = False, + ) -> None: + # The kernel stores dense rows of `hidden` BF16 values at out's address. + assert out.dtype == torch.bfloat16 and out.dim() == 2 and out.is_contiguous(), ( + "out must be a contiguous BF16 [M, hidden] tensor" + ) + assert out.device == ws.flags.device, "out must be on the workspace's device" + + +k3_lamport_finish = K3LamportFinish() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_lamport_push.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_lamport_push.py new file mode 100644 index 000000000000..c8846f4fff32 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_lamport_push.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Push a BF16 tensor into the Lamport slots of the next MNNVL all-reduce on a workspace (the +C++ kernels' first phase as a kernel of its own; CuTe DSL, ``kernels/k3_lamport_ar.py``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.k3_lamport_ar import LamportWorkspace, run_lamport_push +from .k3_lamport_finish import sanitized + + +class K3LamportPush(OpWrapper): + """Store this rank's contiguous BF16 ``x [M, hidden]`` into its slots of the next call on + ``ws``: slot ``(t, rank)`` of every rank's buffer (``one_shot``, multicast) or slot + ``(t // W, rank)`` of the scatter stage of token ``t``'s owner ``t % W`` (two-shot). Returns + None; ``k3_lamport_finish`` or a consumer that finishes the call (``k3_attn_res_lamport``) + completes it.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The group is the test harness's (2 and 4 GPUs of one node); K3 runs 8 (helix8) and 16 (helix16 ULL). + CELLS: tuple[Cell, ...] = tuple( + Cell(why=why, spec=dict(rows=rows, hidden=7168, one_shot=one_shot)) + for why, rows, one_shot in ( + ("one token, one-shot", 1, True), + ("helix16 ULL 1-request step (o_proj / shared-expert push): 8 rows, one-shot", 8, True), + ("helix16 ULL 4-request step: 32 rows, one-shot", 32, True), + ("32 rows, two-shot (the transport above the one-shot size)", 32, False), + ("helix8 b16 mtp5 verify step: 96 rows, two-shot", 96, False), + ("helix8 b16 mtp5 verify step: 96 rows, one-shot", 96, True), + ) + ) + + note = """ + For a producer without a push epilogue. ws (kernels/k3_lamport_ar + .LamportWorkspace) holds the device addresses of the workspace; the kernel + reads the rotation state after griddepcontrol.wait, releases its dependents + right after, and writes only the current buffer's slots (values sanitized: + -0 stored as +0, the sentinel being a -0 pair). x is read-only. Grid + (M, ceil(hidden / 1024)) CTAs of 128 threads, one 16-byte vector each. + A process-wide cache (`_compiled_push`) holds one compiled kernel per + (W, hidden, transport, PDL); compile each key before graph capture. + """ + + def __call__( + self, ws: LamportWorkspace, x: torch.Tensor, one_shot: bool, use_pdl: bool = True + ) -> None: + run_lamport_push(ws, x, one_shot, use_pdl=use_pdl) + + def reference( + self, ws: LamportWorkspace, x: torch.Tensor, one_shot: bool, use_pdl: bool = True + ) -> torch.Tensor: + """The rows this rank's slots hold once the push has landed: ``x`` with ``-0`` stored as + ``+0``. ``ws``, ``one_shot`` and ``use_pdl`` decide where they land and when, not what.""" + return sanitized(x) + + def is_valid( + self, ws: LamportWorkspace, x: torch.Tensor, one_shot: bool, use_pdl: bool = True + ) -> None: + # The kernel loads dense rows of `hidden` BF16 values from x's address. + assert x.dtype == torch.bfloat16 and x.dim() == 2 and x.is_contiguous(), ( + "x must be a contiguous BF16 [M, hidden] tensor" + ) + + +k3_lamport_push = K3LamportPush() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_latent_ar_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_latent_ar_rmsnorm.py new file mode 100644 index 000000000000..1571f46c85ab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/comm/k3_latent_ar_rmsnorm.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 routed latent all-reduce + RMSNorm (MNNVL one-shot) that releases its programmatic +dependents right after its own ``griddepcontrol.wait`` (CuTe DSL, ``kernels/k3_latent_ar_norm.py``).""" + +from __future__ import annotations + +from typing import Optional, Sequence, Tuple, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.k3_latent_ar_norm import run_k3_latent_ar_norm, run_k3_latent_ar_norm_row +from ..norm.k3_attn_res_norm import mxfp8_rows_reference +from .k3_lamport_finish import lamport_sum_reference + +#: Ranks per FP32 partial sum of the C++ one-shot kernel the result reproduces. +_RANK_CHUNK = 8 + + +class K3LatentArRmsnorm(OpWrapper): + """``out = bf16((s * rsqrt(mean(s^2) + eps)) * gamma)`` with ``s`` the BF16 all-reduce of every + rank's ``x [M, 3584]`` (``norm=False``: ``s`` alone), over the MNNVL Lamport workspace at + ``uc_addr`` (this rank's unicast buffer) / ``mc_addr`` (multicast base) / ``flags`` (rotation + state), as this ``rank`` of ``world_size``. Returns ``out``, or ``(out, xq, xs)`` with + ``mxfp8_out = (xq, xs)``: also the MXFP8 form of ``out``, scales row-major ``[M, 112]`` or with + ``sf_blocks`` in the mid-M GEMM's atoms. ``row``: the one-CTA-per-row kernel (up to 128 rows, + 2..16 ranks), else the cluster kernel (up to 38 rows); same bits.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The group is the test harness's (2 and 4 GPUs of one node); K3 runs 8 (helix8) and 16 (helix16 ULL). + # mxfp8: "rows" (row-major scales, the low-M GEMM's operand), "atoms" (the mid-M GEMM's, one + # 128-row block) or None. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=why, + spec=dict( + rows=rows, + hidden=3584, + row=row, + mxfp8=mxfp8, + early_state=early_state, + eps=1e-6, + ), + ) + for why, rows, row, mxfp8, early_state in ( + ("one token, row kernel", 1, True, None, False), + ("helix8 own-workspace latent all-reduce: 16 rows, row kernel", 16, True, None, False), + ( + "helix8 b16 mtp5 verify: 96 rows, MXFP8 in the mid-M GEMM's atoms", + 96, + True, + "atoms", + False, + ), + ("cluster kernel at its 38-row maximum", 38, False, None, False), + ("helix16 ULL tail sandwich: 8 rows, MXFP8 for the low-M GEMM", 8, True, "rows", False), + ( + "helix16 ULL column-parallel tail: 32 rows, MXFP8, early state", + 32, + True, + "rows", + True, + ), + ) + ) + + note = """ + Bit-identical to the C++ one-shot MNNVL all-reduce with the RMSNorm + epilogue (oneshotAllreduceFusionKernel), so C++ and these calls alternate on + one workspace (same uint32[9] rotation state); every rank makes the same + sequence of calls. Each Lamport buffer must hold M x W x 3584 x 2 bytes. + Pushes x (sanitized) through the multicast address, re-arms the previous + call's buffer, polls the W slots and sums them in FP32 (chunks of eight + ranks), rounds to BF16; the norm squares in BF16, sums in the C++ reduction + tree, rsqrt.approx. Triggers dependents right after griddepcontrol.wait. + Writes out (and mxfp8_out); x and gamma are read-only. + prepushed: the producer already stored x into this rank's slots (the kernel + skips its push). early_state: the state read, arrival count and re-arm run + before griddepcontrol.wait (the previous call on the workspace completed + before x's producer started). sf_blocks > 0: scales in the mid-M GEMM's + tcgen05 atoms of sf_blocks 128-row blocks, padding rows neutral 127. + Process-wide compile caches (`_compiled`, `_compiled_row`) keyed by + (W, PDL, norm, mxfp8[, sf_blocks, prepushed, early_state]); compile each + key before graph capture. + """ + + def __call__( + self, + x: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + *, + uc_addr: int, + mc_addr: int, + flags: torch.Tensor, + rank: int, + world_size: int, + use_pdl: bool = True, + norm: bool = True, + mxfp8_out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + sf_blocks: int = 0, + prepushed: bool = False, + early_state: bool = False, + row: bool = True, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + if row: + run_k3_latent_ar_norm_row( + x, + gamma, + eps, + out, + uc_addr=uc_addr, + mc_addr=mc_addr, + flags=flags, + rank=rank, + world_size=world_size, + use_pdl=use_pdl, + norm=norm, + mxfp8_out=mxfp8_out, + sf_blocks=sf_blocks, + prepushed=prepushed, + early_state=early_state, + ) + else: + run_k3_latent_ar_norm( + x, + gamma, + eps, + out, + uc_addr=uc_addr, + mc_addr=mc_addr, + flags=flags, + rank=rank, + world_size=world_size, + use_pdl=use_pdl, + norm=norm, + mxfp8_out=mxfp8_out, + ) + if mxfp8_out is None: + return out + return (out, *mxfp8_out) + + def reference( + self, + inputs_in_rank_order: Sequence[torch.Tensor], + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + *, + uc_addr: int = 0, + mc_addr: int = 0, + flags: Optional[torch.Tensor] = None, + rank: int = 0, + world_size: int = 0, + use_pdl: bool = True, + norm: bool = True, + mxfp8_out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + sf_blocks: int = 0, + prepushed: bool = False, + early_state: bool = False, + row: bool = True, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + """The collective exception: every rank's ``x`` in rank order replaces ``x``. The sum in + the C++ order (chunks of eight ranks), the norm in FP32 rounded once; with ``mxfp8_out`` + the MXFP8 form of the expected ``out`` with row-major scales (``sf_blocks`` only moves + them). ``out`` and the workspace arguments are not read.""" + total = lamport_sum_reference(inputs_in_rank_order, _RANK_CHUNK) + if not norm: + return total + s = total.float() + expected = (s * torch.rsqrt(s.square().mean(dim=-1, keepdim=True) + eps)) * gamma.float() + expected = expected.to(torch.bfloat16) + if mxfp8_out is None: + return expected + return (expected, *mxfp8_rows_reference(expected)) + + def is_valid( + self, + x: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + *, + uc_addr: int, + mc_addr: int, + flags: torch.Tensor, + rank: int, + world_size: int, + use_pdl: bool = True, + norm: bool = True, + mxfp8_out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + sf_blocks: int = 0, + prepushed: bool = False, + early_state: bool = False, + row: bool = True, + ) -> None: + # The launchers do not re-validate: the kernels read dense 3584-wide rows from x and gamma + # and store dense rows into out. + from ...kernels.k3_latent_ar_norm import can_run, can_run_row + + served = can_run_row(x, gamma, world_size) if row else can_run(x, gamma, world_size) + assert served, "outside the kernel's shape (3584-wide contiguous BF16 rows, M, W)" + assert out.shape == x.shape and out.dtype == x.dtype and out.is_contiguous() + assert 0 <= rank < world_size, f"rank {rank} outside a group of {world_size}" + # The cluster kernel takes none of these and would drop them without an error. + assert row or not (sf_blocks or prepushed or early_state), ( + "sf_blocks / prepushed / early_state are row-kernel options" + ) + + +k3_latent_ar_rmsnorm = K3LatentArRmsnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/embedding/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/embedding/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/embedding/k3_embedding_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/embedding/k3_embedding_gather.py new file mode 100644 index 000000000000..14b3b944f0f9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/embedding/k3_embedding_gather.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Token-embedding lookup of a whole table as a vectorized row copy (Triton, +``kernels/embedding_gather.py``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.embedding_gather import embedding_gather + + +class K3EmbeddingGather(OpWrapper): + """Return the new ``[len(ids), hidden]`` tensor of rows ``ids`` of ``weight [vocab, hidden]`` + (1-D ``ids``, contiguous ``weight``); an id outside ``[0, vocab)`` gives a zero row.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The replicated table of Kimi K3 (vocab 163840, hidden 7168), rows: tokens of the decode step. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{why}: {rows} rows of the replicated [163840, 7168] table", + spec=dict(rows=rows, vocab=163840, hidden=7168, padded=False), + ) + for why, rows in ( + ("helix16 ULL 1-request step", 8), + ("helix8 / helix16 ULL 2-request step", 16), + ("helix16 ULL 4-request step", 32), + ("helix8 b16 mtp5 verify step", 96), + ) + ), + Cell( + why="ids padded with -1 (rejection sampler): zero rows, 8 rows", + spec=dict(rows=8, vocab=163840, hidden=7168, padded=True), + ), + ) + + note = """ + A copy: bitwise F.embedding on in-range ids. One program per (token, 1024 + columns) with 16-byte accesses instead of the scalar index-select kernel; + the launch shape depends only on the token count, so it is capturable after + a warm-up per Triton specialization (dtype, divisibility). No metadata or + state; ids and weight are read-only. The decode step's embedding where the + target runs the speculative-decode head (helix8, helix16 ULL); the context + worker uses F.embedding. + """ + + def __call__(self, ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return embedding_gather(ids, weight) + + def reference(self, ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``F.embedding`` with out-of-range ids mapped to a zero row.""" + valid = (ids >= 0) & (ids < weight.shape[0]) + rows = F.embedding(torch.where(valid, ids, torch.zeros_like(ids)), weight) + return rows.masked_fill(~valid.unsqueeze(-1), 0) + + +k3_embedding_gather = K3EmbeddingGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/_fp8_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/_fp8_reference.py new file mode 100644 index 000000000000..260d612e7cb0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/_fp8_reference.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-torch dequantize-then-matmul reference shared by the FP8-weight GEMM entries. + +Every GEMM of this family multiplies an E4M3 operand pair whose scales are powers of two (UE8M0 +bytes, ``2^(e - 127)``): the weight has one exponent per 128 x 128 block, the activation one per +1 x 32 (MXFP8) or 1 x 128 group. The products of the decoded operands are exact in FP64, so the +reference is the FP64 matmul of the decoded operands rounded once to BF16; the kernels accumulate +the same products in FP32. The scale layouts decoded here: + +* DeepGEMM-packed weight scale: ``int32 [N, ceil(K / 512)]`` (strides ``(1, N)``), four K-block + exponents per word, little-endian. +* K32 R128c4 (the CuTe MXFP8 GEMM's): byte ``((r // 128) * (K / 128) + c // 4) * 512 + + (r % 32) * 16 + ((r % 128) // 32) * 4 + c % 4`` for row ``r``, K32 column ``c``. +* The mid-M GEMM's activation layout: the same 512-byte atoms ordered ``[K / 128][row blocks]``. +* Row-major ``uint8 [M, K / 32]`` (the MXFP8 low-M GEMM's activation scales). +""" + +from __future__ import annotations + +import torch + +#: Smallest magnitude (BF16 bits) the kernels' activation quantization takes as a group maximum. +AMAX_FLOOR_BF16 = 0x38D2 + + +def exp2_ue8m0(e: torch.Tensor) -> torch.Tensor: + """``2^(e - 127)`` in FP64 for UE8M0 bytes ``e``.""" + return torch.pow(2.0, e.to(torch.float64) - 127.0) + + +def packed_block_exponents(w_scale: torch.Tensor, n: int, k: int) -> torch.Tensor: + """``uint8 [N, K / 128]``: each weight row's K-block exponents from the DeepGEMM-packed scale.""" + kb = k // 128 + words = w_scale.to(torch.int64) & 0xFFFFFFFF + parts = [(words >> (8 * j)) & 0xFF for j in range(4)] + return torch.stack(parts, dim=-1).reshape(n, -1)[:, :kb].to(torch.uint8) + + +def _atom_offsets(rows: int, k: int, blocks: int, k_major: bool, device) -> torch.Tensor: + r = torch.arange(rows, device=device).unsqueeze(1) + c = torch.arange(k // 32, device=device).unsqueeze(0) + if k_major: + atom = (c // 4) * blocks + r // 128 + else: + atom = (r // 128) * (k // 128) + c // 4 + return atom * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + c % 4 + + +def r128c4_exponents(sf: torch.Tensor, rows: int, k: int) -> torch.Tensor: + """``uint8 [rows, K / 32]`` from a K32 R128c4 scale (the CuTe MXFP8 GEMM's layout).""" + blocks = -(-rows // 128) + flat = sf.reshape(-1).view(torch.uint8) + return flat[_atom_offsets(rows, k, blocks, False, sf.device)] + + +def mid_exponents(xs: torch.Tensor, rows: int, k: int, sf_blocks: int) -> torch.Tensor: + """``uint8 [rows, K / 32]`` from the mid-M GEMM's ``[K / 128][sf_blocks][512 B]`` layout.""" + flat = xs.reshape(-1).view(torch.uint8) + return flat[_atom_offsets(rows, k, sf_blocks, True, xs.device)] + + +def dequant(codes: torch.Tensor, exponents: torch.Tensor, group: int) -> torch.Tensor: + """FP64 ``codes * 2^(e - 127)`` with one exponent per ``group`` consecutive columns.""" + scale = exp2_ue8m0(exponents).repeat_interleave(group, dim=1)[:, : codes.shape[1]] + return codes.view(torch.float8_e4m3fn).to(torch.float64) * scale + + +def fp8_block_weight(w: torch.Tensor, w_scale: torch.Tensor) -> torch.Tensor: + """FP64 ``[N, K]`` weight from its E4M3 codes and either 128 x 128 block-scale layout the + projections hold (DeepGEMM-packed ``int32`` or K32 R128c4 ``uint8``).""" + n, k = w.shape + if w_scale.dtype == torch.int32: + return dequant(w, packed_block_exponents(w_scale, n, k), 128) + return dequant(w, r128c4_exponents(w_scale, n, k), 32) + + +def mxfp8_quantize(x: torch.Tensor, group: int = 32) -> tuple[torch.Tensor, torch.Tensor]: + """``(codes E4M3 [M, K], exponents uint8 [M, K / group])`` of BF16 ``x`` as the kernels + quantize their activation: per group ``e`` = the biased exponent of ``2^ceil(log2(amax / + 448))`` (``amax`` floored at BF16 ``0x38D2``), codes ``satfinite_rn_e4m3(x * 2^(127 - e))`` + (an exact BF16 multiply).""" + m, k = x.shape + g = x.reshape(m, k // group, group) + bits = g.contiguous().view(torch.int16).to(torch.int32) & 0x7FFF + amax = bits.amax(dim=-1).clamp_min(AMAX_FLOOR_BF16) + e = (amax >> 7) - 8 + ((amax & 0x7F) > 0x60).to(torch.int32) + inv = torch.pow(2.0, (127 - e).float()).to(torch.bfloat16) + scaled = (g * inv.unsqueeze(-1)).float() + codes = scaled.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).reshape(m, k) + return codes, e.to(torch.uint8) + + +def matmul_bf16(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """BF16 ``a @ b.T`` of FP64 operands, rounded once.""" + return (a.to(torch.float64) @ b.to(torch.float64).T).to(torch.bfloat16) + + +def lowm_gemm(x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, group: int = 32): + """BF16 ``x @ W^T`` as the FP8-weight GEMMs compute it on BF16 rows: ``x`` quantized to E4M3 + with one exponent per ``group`` columns, times the dequantized weight.""" + codes, e = mxfp8_quantize(x, group) + return matmul_bf16(dequant(codes, e, group), fp8_block_weight(w, w_scale)) + + +def all_reduce_bf16(parts) -> torch.Tensor: + """The MNNVL all-reduce of BF16 per-rank tensors given in rank order: FP32 sum in rank order, + rounded to BF16 (the one-shot and two-shot kernels' order up to eight ranks).""" + total = parts[0].float() + for part in parts[1:]: + total = total + part.float() + return total.to(torch.bfloat16) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/_mxfp8_gemm.py new file mode 100644 index 000000000000..a18196eb031a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/_mxfp8_gemm.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""":func:`k3_mxfp8_gemm_rubin`: the SM107 dense MXFP8 GEMM over K3's tactic set. + +The kernel is main's CuTe DSL persistent dense block-scaled GEMM (``CuteDSLMXFP8RubinLinear``, +compiled-kernel cache shared with it) with K3's fine-grained low-M tuning buckets +(``CuteDSLKimiK3MXFP8RubinLinear``). Its autotuner candidates are those of K3's projections in +the agentic recipes: + +* no 1-CTA b-reuse tile (``mma_tiler_m, mma_inst_m == 256, 128``); +* prefetch only on the unswapped 2-CTA b-reuse tile (``512, 256``); +* split-K (atomic BF16 accumulation, so not run-to-run deterministic) only up to M = 256. + +``get_valid_tactics`` keeps exactly those tactics of main's list, in main's order; the tactic tuples +and the launch are main's. Inputs: FP8 E4M3 ``[M, K]`` activation and ``[N, K]`` weight, their +K32 R128c4 UE8M0 scales (``fp8_quantize_1x128_cutedsl_ue8m0`` / +``transform_k128_scales_to_cutedsl_mxfp8_layout``); output BF16 ``[M, N]``. :data:`OP_NAME` is the +AutoTuner cache key of the GEMM. +""" + +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._utils import get_sm_version + +#: (mma_tiler_m, mma_inst_shape_m) of the tiles K3 tunes over. +K3_MMA_M = ((128, 128), (256, 256), (512, 256)) +#: The tile whose unswapped tactics also tune with prefetch. +K3_PREFETCH_MMA_M = (512, 256) +#: The largest (bucket) M that tunes split-K > 1. +K3_SPLIT_K_MAX_M = 256 + +OP_NAME = "trtllm::k3_mxfp8_gemm_rubin" + +# One read-only FP32 scalar one per CUDA device: the GEMM's alpha operand, so steady-state calls +# (CUDA graph replays included) launch no fill kernel. +_MXFP8_GEMM_ALPHA_CACHE: "dict[torch.device, torch.Tensor]" = {} + + +def _get_mxfp8_gemm_alpha(device: torch.device) -> torch.Tensor: + """The cached FP32 scalar one of ``device``. + + Allocating it inside CUDA graph capture raises: one eager call per device (which autotuning + needs anyway) has to precede the capture. + """ + device = torch.device(device) + if device.type != "cuda": + raise ValueError(f"MXFP8 GEMM alpha requires a CUDA device, got {device}.") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + alpha = _MXFP8_GEMM_ALPHA_CACHE.get(device) + if alpha is None: + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MXFP8 GEMM alpha cache must be initialized before CUDA graph " + f"capture on {device}; run one eager GEMM warmup first." + ) + alpha = torch.ones((), dtype=torch.float32, device=device) + _MXFP8_GEMM_ALPHA_CACHE[device] = alpha + return alpha + + +def keeps_tactic(tactic, m: int) -> bool: + """Whether K3 tunes over main's tactic ``tactic`` for a (bucket) M of ``m`` rows. + + ``tactic``: ``("base", mma_tiler_mnk, mma_inst_shape, cluster, swap_ab, use_prefetch, + scheduler, raster, split_k)`` or ``("mixed_clusters", mma_tiler_mnk, mma_inst_shape, preferred, + fallback, swap_ab, use_prefetch, [scheduler,] raster)``. + """ + variant, mma_tiler_mnk, mma_inst_shape = tactic[0], tactic[1], tactic[2] + # base: (.., cluster, swap_ab, use_prefetch, ..); mixed: (.., preferred, fallback, swap_ab, ..) + flags = 4 if variant == "base" else 5 + swap_ab, use_prefetch = tactic[flags], tactic[flags + 1] + mma_m = (mma_tiler_mnk[0], mma_inst_shape[0]) + if mma_m not in K3_MMA_M: + return False + if use_prefetch and (swap_ab or mma_m != K3_PREFETCH_MMA_M): + return False + if variant == "base" and tactic[8] > 1 and m > K3_SPLIT_K_MAX_M: + return False + return True + + +if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + from tensorrt_llm._torch.autotuner import AutoTuner, OptimizationProfile + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLKimiK3MXFP8RubinLinear + + class K3MXFP8RubinRunner(CuteDSLKimiK3MXFP8RubinLinear): + """Main's K3 MXFP8 runner with K3's tactic set (see the module docstring).""" + + def get_valid_tactics( + self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs + ) -> List[Tuple]: + m = inputs[0].shape[0] + return [ + tactic + for tactic in super().get_valid_tactics(inputs, profile, **kwargs) + if keeps_tactic(tactic, m) + ] + + def k3_mxfp8_gemm_rubin( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + ) -> torch.Tensor: + """BF16 ``input @ weight.T`` of MXFP8 operands on SM107, autotuned over K3's tactics.""" + if (sm_version := get_sm_version()) != 107: + raise ValueError(f"{OP_NAME} requires SM107, got SM{sm_version}") + alpha = _get_mxfp8_gemm_alpha(input.device) + runner = K3MXFP8RubinRunner(output_dtype=torch.bfloat16, use_tvm_ffi=True) + inputs = [input, weight, input_scale, weight_scale, alpha] + _, tactic = AutoTuner.get().choose_one( + OP_NAME, [runner], K3MXFP8RubinRunner.tuning_config, inputs + ) + return runner(inputs, tactic=tactic) + + +def run_with_tactic( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + tactic: Optional[tuple], +) -> torch.Tensor: + """The GEMM with one given tactic (None: the runner's default), bypassing the autotuner.""" + alpha = _get_mxfp8_gemm_alpha(input.device) + runner = K3MXFP8RubinRunner(output_dtype=torch.bfloat16, use_tvm_ffi=True) + inputs = [input, weight, input_scale, weight_scale, alpha] + return runner(inputs, tactic=-1 if tactic is None else tactic) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_gemm.py new file mode 100644 index 000000000000..1389434570a3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_gemm.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M FP8-weight GEMM of BF16 rows (activation quantized to MXFP8 in the kernel), with the +output-gated and sigmoid variants.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_fp8w_lowm_gemm as _lowm +from . import _fp8_reference as ref + + +class K3Fp8wLowmGemm(OpWrapper): + """BF16 ``[M, N] = q(x) @ dequant(w, w_scale)^T`` (``q``: MXFP8, one UE8M0 exponent per 1 x 32), + ``w`` the UE8M0-resmoothed E4M3 ``[N, K]`` weight and ``w_scale`` its DeepGEMM-packed scale. + + ``gate``: ``x`` is first multiplied by ``gate`` (PyTorch's BF16 multiply). ``sigmoid``: + ``torch.sigmoid`` of the BF16 output. ``out``: a contiguous BF16 ``[M, N]`` tensor the result + goes into (returned); else a new one. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: K3 decode projections at the row counts their CUDA graphs run: MTP with seven draft tokens + #: puts 8 rows per request, the helix deployments batch 1..8 requests (8..64 rows), dep16 up + #: to 32 per rank. ``variant``: ``new`` (no ``out``), ``out``, ``gated`` or ``sigmoid``. + CELLS: tuple[Cell, ...] = ( + Cell( + why="helix16 KDA fused q/k/v/g head shard, one request: the b1 decode step", + spec={"n": 3072, "k": 7168, "m": 8, "variant": "out"}, + ), + Cell( + why="MoE latent_down (every deployment) on a single row: the smallest MMA bucket", + spec={"n": 3584, "k": 7168, "m": 1, "variant": "new"}, + ), + Cell( + why="MoE latent_up, four requests: the 32-row bucket both tables run on the low-M kernel", + spec={"n": 7168, "k": 3584, "m": 32, "variant": "out"}, + ), + Cell( + why="MLA q_a|kv_a|k_pe (2112 rows, a partial last N tile), eight requests: low-M on " + "GB300, mid-M above Rubin's 32-row cut-over", + spec={"n": 2112, "k": 7168, "m": 64, "variant": "out"}, + ), + Cell( + why="helix16 KDA / MLA o_proj with the output gate multiplied in (gated kernel)", + spec={"n": 7168, "k": 768, "m": 8, "variant": "gated"}, + ), + Cell( + why="helix16 MLA output gate projection with the sigmoid epilogue, two requests", + spec={"n": 768, "k": 7168, "m": 16, "variant": "sigmoid"}, + ), + Cell( + why="dep16 shared-expert gate_up, 16 requests: quantization kernel + mid-M GEMM", + spec={"n": 12288, "k": 7168, "m": 128, "variant": "out"}, + ), + Cell( + why="dep16 KDA / MLA o_proj (K = 12288), 12 requests: the 96-row mid-M bucket", + spec={"n": 7168, "k": 12288, "m": 96, "variant": "new"}, + ), + ) + + note = """ + ``M <= low_m_max_for(N, K)`` runs the low-M kernel (activation quantized inside the GEMM); + larger ``M`` (up to 256, plain variant only) a quantization kernel followed by the mid-M GEMM, + chained with PDL. The gated and sigmoid variants are low-M only. Per-shape tactics and + cut-overs come from per-architecture tables (GB300 for SM100 / SM103, Rubin for SM107); split-K + partials are reduced in a fixed order, so repeated calls give the same bits. A PDL dependent + launch: the weight streams before ``griddepcontrol.wait``, the activation (and gate) after. + """ + + def __call__( + self, + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: Optional[torch.Tensor] = None, + gate: Optional[torch.Tensor] = None, + sigmoid: bool = False, + ) -> torch.Tensor: + if gate is not None: + return _lowm.k3_fp8w_lowm_gemm_gated(x, gate, w, w_scale, out) + if sigmoid: + return _lowm.k3_fp8w_lowm_gemm_sigmoid(x, w, w_scale, out) + if out is not None: + return _lowm.k3_fp8w_lowm_gemm_out(x, w, w_scale, out) + return _lowm.k3_fp8w_lowm_gemm(x, w, w_scale) + + def reference( + self, + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: Optional[torch.Tensor] = None, + gate: Optional[torch.Tensor] = None, + sigmoid: bool = False, + ) -> torch.Tensor: + """The FP64 product of the MXFP8-quantized rows and the dequantized weight, rounded to + BF16 (then ``torch.sigmoid``); ``out`` is not written.""" + if gate is not None: + x = x * gate + y = ref.lowm_gemm(x, w, w_scale) + return torch.sigmoid(y) if sigmoid else y + + def is_valid( + self, + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: Optional[torch.Tensor] = None, + gate: Optional[torch.Tensor] = None, + sigmoid: bool = False, + ) -> None: + # With a gate the gated kernel runs and `sigmoid` would be dropped without an error. + assert gate is None or not sigmoid, "the gated GEMM has no sigmoid epilogue" + + +k3_fp8w_lowm_gemm = K3Fp8wLowmGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_gemm_push.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_gemm_push.py new file mode 100644 index 000000000000..5fe445d64d4e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_gemm_push.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M FP8-weight GEMM of BF16 rows whose epilogue pushes the output into the next MNNVL +all-reduce on a Lamport workspace (instead of storing it).""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_fp8w_lowm_gemm as _lowm +from ...kernels import k3_lamport_ar +from . import _fp8_reference as ref + + +class K3Fp8wLowmGemmPush(OpWrapper): + """Push this rank's BF16 ``x @ dequant(w, w_scale)^T`` (``[M, N]``) into the Lamport slots of + the next all-reduce on ``ws`` (``one_shot`` or two-shot transport); returns None. + ``k3_lamport_ar.run_lamport_finish`` with the same ``one_shot`` (or a consumer that finishes + the call) completes it. + + ``gate``: the output-gated GEMM (``x * gate``, PyTorch's BF16 multiply). ``strided``: ``x`` and + ``w`` are column slices of wider tensors (unit column stride, 16-byte aligned rows) read in + place, ``w_scale`` the packed scale of the ``[N, K]`` slice, and with ``addend`` (contiguous + BF16 ``[M, N]``) the pushed values are ``bf16(result + addend)``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: Row-parallel decode projections of the helix deployments, at 8 rows per request. + #: ``k_total`` / ``strided``: each rank's ``[rank * k, (rank + 1) * k)`` column slice of a + #: ``k_total``-wide operand. ``sm_107``: the push of a split-K tactic, which only SM107's + #: ``ar_push_tactic`` admits (GB300 keeps its ``split_k == 1`` rule there). + CELLS: tuple[Cell, ...] = ( + Cell( + why="helix16 KDA / MLA o_proj all-reduce of one request, one-shot transport", + spec={"n": 7168, "k": 768, "m": 8, "one_shot": True}, + ), + Cell( + why="helix16 o_proj at eight requests, two-shot transport", + spec={"n": 7168, "k": 768, "m": 64, "one_shot": False}, + ), + Cell( + why="helix16 o_proj with the output gate multiplied in (gated kernel)", + spec={"n": 7168, "k": 768, "m": 8, "one_shot": True, "gate": True}, + ), + Cell( + why="MoE row-parallel tail: latent_up on this rank's K slice of the latent, read in " + "place, plus the shared-expert addend (split-K owners push)", + spec={ + "n": 7168, + "k": 896, + "k_total": 3584, + "m": 8, + "one_shot": True, + "strided": True, + "addend": True, + "sm_107": True, + }, + ), + ) + + note = """ + The low-M kernel's tactic, so with the finish the call is bitwise the GEMM followed by the + MNNVL all-reduce of its output; above the low-M range (ungated, unsliced) the mid-M GEMM's push + epilogue after its quantization kernel. ``-0`` is sent as ``+0``. PDL: the dependent (the + finish) is released right after the GEMM's own ``griddepcontrol.wait``; the weight streams + before it. The workspace's rotation state is read, not advanced (the finish advances it). + """ + + def __call__( + self, + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + gate: Optional[torch.Tensor] = None, + addend: Optional[torch.Tensor] = None, + strided: bool = False, + ) -> None: + if strided: + _lowm.run_k3_fp8w_lowm_gemm_split_ar_push(x, w, w_scale, ws, one_shot, addend=addend) + else: + _lowm.run_k3_fp8w_lowm_gemm_ar_push(x, w, w_scale, ws, one_shot, gate=gate) + + def reference( + self, + x: Sequence[torch.Tensor], + w: Sequence[torch.Tensor], + w_scale: Sequence[torch.Tensor], + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + gate: Optional[Sequence[torch.Tensor]] = None, + addend: Optional[Sequence[torch.Tensor]] = None, + strided: bool = False, + ) -> torch.Tensor: + """A collective: ``x``, ``w``, ``w_scale`` (and ``gate`` / ``addend``) are the group's + inputs in rank order. Returns the BF16 ``[M, N]`` the finish of the call yields on every + rank: the all-reduce of each rank's BF16 GEMM output. ``ws``, ``one_shot`` and + ``strided`` do not change the result.""" + parts = [] + for r in range(len(x)): + xr = x[r] if gate is None else x[r] * gate[r] + y = ref.lowm_gemm(xr.contiguous(), w[r].contiguous(), w_scale[r]) + if addend is not None: + y = (y.float() + addend[r].float()).to(torch.bfloat16) + parts.append(y) + return ref.all_reduce_bf16(parts) + + def is_valid( + self, + x: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + gate: Optional[torch.Tensor] = None, + addend: Optional[torch.Tensor] = None, + strided: bool = False, + ) -> None: + # Each launcher ignores the other variant's argument without an error. + assert not (strided and gate is not None), "the strided push has no gated variant" + assert strided or addend is None, "only the strided push adds an addend" + + +k3_fp8w_lowm_gemm_push = K3Fp8wLowmGemmPush() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_mxfp8_gemm.py new file mode 100644 index 000000000000..63da1b4f88b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_mxfp8_gemm.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M FP8-weight GEMM on an MXFP8 activation a producer already quantized (row-major 1 x 32 +exponents).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_fp8w_lowm_mxfp8 as _mx +from . import _fp8_reference as ref + + +class K3Fp8wLowmMxfp8Gemm(OpWrapper): + """``out [M, N]`` = BF16 ``dequant(xq, xs) @ dequant(w, w_scale)^T``; returns ``out``. + + ``xq``: E4M3 ``[M, K]``; ``xs``: UE8M0 ``[M, K / 32]`` row-major; ``w_scale``: the weight's + DeepGEMM-packed scale. ``stage_weights``: the tactic that keeps the CTA's weight slice on chip + before ``griddepcontrol.wait`` (``staged_mx_tactic``, the same bits). + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: Helix decode projections whose producer emits the MXFP8 operand, at 8 rows per request. + #: ``stage``: ``stage_weights``. + CELLS: tuple[Cell, ...] = ( + Cell( + why="helix16 KDA q/k/v/g shard with its weight slice staged before the wait, b1", + spec={"n": 3072, "k": 7168, "m": 8, "stage": True}, + ), + Cell( + why="helix8 KDA q/k/v/g shard, eight requests: above Rubin's 32-row cut-over", + spec={"n": 6144, "k": 7168, "m": 64, "stage": False}, + ), + Cell( + why="MoE latent_down, two requests", + spec={"n": 3584, "k": 7168, "m": 16, "stage": False}, + ), + Cell( + why="MLA q_a|kv_a|k_pe (a partial last N tile), four requests", + spec={"n": 2112, "k": 7168, "m": 32, "stage": False}, + ), + Cell( + why="helix16 KDA / MLA o_proj on one row (decode without MTP)", + spec={"n": 7168, "k": 768, "m": 1, "stage": False}, + ), + ) + + note = """ + Same tiling, split-K, K-tile order and scale folding as the low-M kernel of + :class:`K3Fp8wLowmGemm`, so on the operand that kernel's quantization produces the output has + its bits. Inputs are not re-validated (``k3_fp8w_lowm_mxfp8.can_run_mxfp8`` is the caller's + check). A PDL dependent: the weight warp releases the dependent grid at entry and streams before + ``griddepcontrol.wait``. + """ + + def __call__( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + stage_weights: bool = False, + ) -> torch.Tensor: + return _mx.run_k3_fp8w_lowm_gemm_mxfp8(xq, xs, w, w_scale, out, stage_weights=stage_weights) + + def reference( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + stage_weights: bool = False, + ) -> torch.Tensor: + """The expected ``out``: FP64 product of the decoded operands, rounded to BF16.""" + return ref.matmul_bf16(ref.dequant(xq, xs, 32), ref.fp8_block_weight(w, w_scale)) + + def is_valid( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + stage_weights: bool = False, + ) -> None: + # The launcher passes raw pointers; can_run_mxfp8 is the contract it does not re-check. + assert _mx.can_run_mxfp8(xq, xs, w, w_scale, allow_partial_n_tile=True, above_cutover=True) + assert out.dtype == torch.bfloat16 and out.is_contiguous() + assert out.shape == (xq.shape[0], w.shape[0]) + + +k3_fp8w_lowm_mxfp8_gemm = K3Fp8wLowmMxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_mxfp8_push.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_mxfp8_push.py new file mode 100644 index 000000000000..40509437211a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_lowm_mxfp8_push.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M FP8-weight GEMM on a prequantized MXFP8 activation whose epilogue pushes the output into +the next MNNVL Lamport call: all-reduce slots, a K slice of wider operands, or all-gather columns.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_fp8w_lowm_mxfp8 as _mx +from ...kernels import k3_lamport_ar +from . import _fp8_reference as ref + + +class K3Fp8wLowmMxfp8Push(OpWrapper): + """Push this rank's BF16 ``dequant(xq, xs) @ dequant(w, w_scale)^T`` into the next Lamport call + on ``ws``; returns None. ``xq`` E4M3 ``[M, K]``, ``xs`` UE8M0 ``[M, K / 32]`` row-major, + ``w_scale`` DeepGEMM-packed. + + * Default: into the all-reduce slots (``one_shot`` or two-shot transport). + * ``k1``: on the K columns ``[k0, k0 + k1)`` of wider ``xq`` / ``xs`` / ``w``, read in place + (``w_scale`` the packed scale of the ``[N, k1]`` slice); with ``addend`` (contiguous BF16 + ``[M, N]``) the pushed values are ``bf16(result + addend)``. + * ``hidden``: into columns ``[col0, col0 + N)`` of the call's all-gather rows of width + ``hidden`` (one-shot transport; ``one_shot`` is ignored). + + ``tactic``: the low-M MXFP8 tactic (default ``default_mx_tactic`` of the GEMM's own shape; + the all-gather caller passes the full projection's so every 128-row tile is computed as there). + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: The helix decode pushes, at 8 rows per request: ``variant`` ``ar`` (all-reduce), ``slice`` + #: (each rank's K slice ``[rank * k, (rank + 1) * k)`` of a ``k_total``-wide operand, plus an + #: addend) or ``ag`` (each rank's ``n`` output columns of a ``hidden``-wide projection, with + #: the full projection's tactic). + CELLS: tuple[Cell, ...] = ( + Cell( + why="helix16 o_proj on the KDA verify's MXFP8 rows, one-shot all-reduce, b1", + spec={"variant": "ar", "n": 7168, "k": 768, "m": 8, "one_shot": True}, + ), + Cell( + why="helix8 o_proj (K = 1536), four requests, two-shot all-reduce", + spec={"variant": "ar", "n": 7168, "k": 1536, "m": 32, "one_shot": False}, + ), + Cell( + why="MoE row-parallel tail: latent_up on this rank's K slice of the MXFP8 latent, " + "plus the shared-expert addend", + spec={"variant": "slice", "n": 7168, "k": 896, "k_total": 3584, "m": 8}, + ), + Cell( + why="MoE column-parallel tail: this rank's latent_up output columns into the " + "all-gather rows, with the full projection's tactic", + spec={"variant": "ag", "hidden": 7168, "k": 3584, "m": 8}, + ), + ) + + note = """ + The GEMM of :class:`K3Fp8wLowmMxfp8Gemm` (same tactic, K-tile order and rounding) with the + Lamport push epilogue; ``k3_lamport_ar.run_lamport_finish`` (``ag=True`` for all-gather rows) + or a consumer's fused finish completes the call. ``-0`` is sent as ``+0``. A PDL dependent: + the weight streams before ``griddepcontrol.wait``. + """ + + def __call__( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool = True, + k0: int = 0, + k1: Optional[int] = None, + addend: Optional[torch.Tensor] = None, + col0: int = 0, + hidden: Optional[int] = None, + tactic: Optional["_mx.K3LowMMxTactic"] = None, + ) -> None: + if hidden is not None: + _mx.run_k3_fp8w_lowm_gemm_mxfp8_ag_push( + xq, xs, w, w_scale, ws, col0, hidden, tactic=tactic + ) + elif k1 is not None: + _mx.run_k3_fp8w_lowm_gemm_mxfp8_slice_ar_push( + xq, xs, w, w_scale, k0, k1, ws, one_shot, addend=addend, tactic=tactic + ) + else: + _mx.run_k3_fp8w_lowm_gemm_mxfp8_ar_push(xq, xs, w, w_scale, ws, one_shot, tactic=tactic) + + def reference( + self, + xq: Sequence[torch.Tensor], + xs: Sequence[torch.Tensor], + w: Sequence[torch.Tensor], + w_scale: Sequence[torch.Tensor], + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool = True, + k0: Optional[Sequence[int]] = None, + k1: Optional[int] = None, + addend: Optional[Sequence[torch.Tensor]] = None, + col0: Optional[Sequence[int]] = None, + hidden: Optional[int] = None, + tactic: Optional["_mx.K3LowMMxTactic"] = None, + ) -> torch.Tensor: + """A collective: ``xq``, ``xs``, ``w``, ``w_scale`` (``k0``, ``addend``, ``col0``) are the + group's inputs in rank order. All-reduce: the BF16 ``[M, N]`` the finish yields, the all-reduce of + the ranks' BF16 GEMM outputs. All-gather (``hidden``): the BF16 ``[M, hidden]`` all-gather + rows, rank ``r``'s output in columns ``[col0[r], col0[r] + N_r)`` (zero elsewhere; the + finish adds them to the call's reduced slots). ``ws``, ``one_shot`` and ``tactic`` do not + change the result.""" + parts = [] + for r in range(len(xq)): + a, b = xq[r], w[r] + e = xs[r] + if k1 is not None: + lo = 0 if k0 is None else k0[r] + a, b = a[:, lo : lo + k1], b[:, lo : lo + k1] + e = e[:, lo // 32 : (lo + k1) // 32] + y = ref.matmul_bf16(ref.dequant(a, e, 32), ref.fp8_block_weight(b, w_scale[r])) + if addend is not None: + y = (y.float() + addend[r].float()).to(torch.bfloat16) + parts.append(y) + if hidden is None: + return ref.all_reduce_bf16(parts) + rows = torch.zeros( + (parts[0].shape[0], hidden), dtype=torch.bfloat16, device=parts[0].device + ) + for r, y in enumerate(parts): + rows[:, col0[r] : col0[r] + y.shape[1]] = y + return rows + + def is_valid( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool = True, + k0: int = 0, + k1: Optional[int] = None, + addend: Optional[torch.Tensor] = None, + col0: int = 0, + hidden: Optional[int] = None, + tactic: Optional["_mx.K3LowMMxTactic"] = None, + ) -> None: + # Each launcher ignores the other variants' arguments without an error. + assert hidden is None or (k1 is None and addend is None), ( + "the all-gather push takes no K slice and no addend" + ) + assert k1 is not None or (k0 == 0 and addend is None), ( + "k0 and addend need the K-slice push (k1)" + ) + assert xq.dtype == torch.float8_e4m3fn and xs.dtype == torch.uint8 + assert w_scale.dtype == torch.int32, "w_scale must be the DeepGEMM-packed UE8M0 scale" + + +k3_fp8w_lowm_mxfp8_push = K3Fp8wLowmMxfp8Push() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_mid_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_mid_gemm.py new file mode 100644 index 000000000000..e69214cb4754 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_mid_gemm.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mid-M FP8-weight GEMM on an MXFP8 operand a producer already quantized (no quantization +kernel).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_fp8w_lowm_gemm as _lowm +from . import _fp8_reference as ref + + +class K3Fp8wMidGemm(OpWrapper): + """``out [M, N]`` = BF16 ``dequant(xq, xs) @ dequant(w, w_scale)^T`` with the mid-M kernel of + ``tactic`` (``k3_fp8w_lowm_gemm.mid_prequant_plan``); returns ``out``. + + ``xq``: E4M3 ``[M, K]``; ``xs``: UE8M0 1 x 32 exponents in the mid-M layout + (``[K / 128][sf_blocks][512 B]``, ``sf_blocks = tactic.mma_n`` rows / 128 rounded up); + ``w_scale``: the weight's DeepGEMM-packed scale. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: Rows above the low-M cut-over of both tables (GB300: 64 for these weights), on the operand + #: ``quantize_mid_operand`` produces for ``mid_prequant_plan``'s tactic. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MoE latent_up on the latent all-reduce's MXFP8 rows, 12 requests", + spec={"n": 7168, "k": 3584, "m": 96}, + ), + Cell( + why="MoE latent_up at dep16's largest graph (32 requests x 8 rows): mma_n 256", + spec={"n": 7168, "k": 3584, "m": 256}, + ), + Cell( + why="dep16 KDA / MLA o_proj (K = 12288) on the producer's MXFP8 rows, 20 requests", + spec={"n": 7168, "k": 12288, "m": 160}, + ), + Cell( + why="dep16 MLA q_b (N = 18432, K = 1536), 28 requests: the 224-row bucket", + spec={"n": 18432, "k": 1536, "m": 224}, + ), + ) + + note = """ + The GEMM :class:`K3Fp8wLowmGemm` runs above the low-M cut-over, without its quantization + kernel: on the operand ``K3MxFp8QuantKernel`` (or a producer reproducing it) writes, the output + has the BF16 path's bits. A PDL dependent: the weight streams before ``griddepcontrol.wait``. + Inputs are not validated. + """ + + def __call__( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + tactic: "_lowm.K3MidMTactic", + ) -> torch.Tensor: + return _lowm.run_mid_prequantized(xq, xs, w, w_scale, out, tactic) + + def reference( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + tactic: "_lowm.K3MidMTactic", + ) -> torch.Tensor: + """The expected ``out``: FP64 product of the decoded operands, rounded to BF16.""" + m, k = xq.shape + sf_blocks = -(-tactic.mma_n // 128) + a = ref.dequant(xq, ref.mid_exponents(xs, m, k, sf_blocks), 32) + return ref.matmul_bf16(a, ref.fp8_block_weight(w, w_scale)) + + def is_valid( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + tactic: "_lowm.K3MidMTactic", + ) -> None: + # The launcher passes raw pointers: nothing checks shapes, dtypes or strides. + m, k = xq.shape + n = w.shape[0] + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous() + assert w.dtype == torch.float8_e4m3fn and w.is_contiguous() and w.shape[1] == k + assert w_scale.dtype == torch.int32, "w_scale must be the DeepGEMM-packed UE8M0 scale" + sf_blocks = -(-tactic.mma_n // 128) + assert xs.dtype == torch.uint8 and xs.numel() == (k // 128) * sf_blocks * 512, ( + "xs must be the mid-M layout of the tactic" + ) + assert out.dtype == torch.bfloat16 and out.is_contiguous() and out.shape == (m, n) + assert m <= tactic.mma_n, f"M={m} exceeds the tactic's {tactic.mma_n} rows" + + +k3_fp8w_mid_gemm = K3Fp8wMidGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_mid_gemm_push.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_mid_gemm_push.py new file mode 100644 index 000000000000..008bfc75f068 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_fp8w_mid_gemm_push.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mid-M FP8-weight GEMM on a prequantized MXFP8 operand whose epilogue pushes the output into the +next MNNVL all-reduce on a Lamport workspace.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_fp8w_lowm_gemm as _lowm +from ...kernels import k3_lamport_ar +from . import _fp8_reference as ref + + +class K3Fp8wMidGemmPush(OpWrapper): + """Push this rank's BF16 ``dequant(xq, xs) @ dequant(w, w_scale)^T`` (``[M, N]``, the mid-M + kernel of ``tactic``) into the Lamport slots of the next all-reduce on ``ws``; returns None. + Operands as :class:`K3Fp8wMidGemm`'s (``xs`` in the mid-M layout).""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: Rows above the low-M cut-over of both tables, on ``quantize_mid_operand``'s operand for + #: ``mid_ar_push_tactic``'s tactic. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MoE latent_up all-reduce on 12 requests, two-shot transport", + spec={"n": 7168, "k": 3584, "m": 96, "one_shot": False}, + ), + Cell( + why="o_proj (K = 768) all-reduce on 16 requests, two-shot transport", + spec={"n": 7168, "k": 768, "m": 128, "one_shot": False}, + ), + Cell( + why="MoE latent_up all-reduce on 9 requests, one-shot transport", + spec={"n": 7168, "k": 3584, "m": 72, "one_shot": True}, + ), + ) + + note = """ + :class:`K3Fp8wMidGemm` with the all-reduce push epilogue: with the finish + (``k3_lamport_ar.run_lamport_finish``, same ``one_shot``) the call is bitwise that GEMM followed + by the MNNVL all-reduce of its output. A PDL dependent: the weight streams before + ``griddepcontrol.wait``. + """ + + def __call__( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + tactic: "_lowm.K3MidMTactic", + ) -> None: + _lowm.run_mid_prequantized_ar_push(xq, xs, w, w_scale, ws, one_shot, tactic) + + def reference( + self, + xq: Sequence[torch.Tensor], + xs: Sequence[torch.Tensor], + w: Sequence[torch.Tensor], + w_scale: Sequence[torch.Tensor], + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + tactic: "_lowm.K3MidMTactic", + ) -> torch.Tensor: + """A collective: ``xq``, ``xs``, ``w``, ``w_scale`` are the group's inputs in rank order. + Returns the BF16 ``[M, N]`` the finish yields: the all-reduce of the ranks' BF16 GEMM + outputs. ``ws`` and ``one_shot`` do not change the result.""" + sf_blocks = -(-tactic.mma_n // 128) + parts = [] + for r in range(len(xq)): + m, k = xq[r].shape + a = ref.dequant(xq[r], ref.mid_exponents(xs[r], m, k, sf_blocks), 32) + parts.append(ref.matmul_bf16(a, ref.fp8_block_weight(w[r], w_scale[r]))) + return ref.all_reduce_bf16(parts) + + def is_valid( + self, + xq: torch.Tensor, + xs: torch.Tensor, + w: torch.Tensor, + w_scale: torch.Tensor, + ws: "k3_lamport_ar.LamportWorkspace", + one_shot: bool, + tactic: "_lowm.K3MidMTactic", + ) -> None: + # The launcher passes raw pointers: nothing checks dtypes, layouts or strides. + k = xq.shape[1] + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous() + assert w.dtype == torch.float8_e4m3fn and w.is_contiguous() and w.shape[1] == k + assert w_scale.dtype == torch.int32, "w_scale must be the DeepGEMM-packed UE8M0 scale" + sf_blocks = -(-tactic.mma_n // 128) + assert xs.dtype == torch.uint8 and xs.numel() == (k // 128) * sf_blocks * 512, ( + "xs must be the mid-M layout of the tactic" + ) + + +k3_fp8w_mid_gemm_push = K3Fp8wMidGemmPush() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_latent_down_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_latent_down_fp4.py new file mode 100644 index 000000000000..65518b9c5bfc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_latent_down_fp4.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MoE latent_down projection with the routed experts' NVFP4 input quantization fused into +its epilogue.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) + +from ...kernels import k3_latent_fp8_gemm as _latent +from . import _fp8_reference as ref + + +class K3LatentDownFp4(OpWrapper): + """NVFP4 ``(codes uint8 [M, N / 2], scales uint8 [M, N / 16])`` of BF16 ``x @ dequant(w)^T``. + + ``w``: the UE8M0-resmoothed E4M3 ``[N, K]`` weight, ``w_scale`` its DeepGEMM-packed scale, + ``global_sf``: the consumer's FP32 ``[1]`` NVFP4 global scale (``fp4_quantize``'s). Scales in + the linear (row-major) layout. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: The MoE latent_down (N = 3584, K = 7168) of every deployment: the decode buckets (on SM107 + #: the low-M route, 1 x 32) and a tp1 context chunk (the persistent GEMM, 1 x 128). + CELLS: tuple[Cell, ...] = ( + Cell(why="one row: decode without MTP", spec={"m": 1}), + Cell(why="one request with seven draft tokens", spec={"m": 8}), + Cell(why="eight requests: the largest helix graph", spec={"m": 64}), + Cell(why="a tp1 context chunk: the persistent MXFP8 GEMM on every arch", spec={"m": 1024}), + ) + + note = """ + Two kernels chained with PDL: the activation quantized to E4M3 with one UE8M0 exponent per + 1 x 128, then a persistent MXFP8 GEMM (no split-K) whose epilogue rounds to BF16 and quantizes + with ``fp4_quantize``'s arithmetic. Up to ``k3_latent_fp8_gemm.lowm_route_max_m(N, K)`` rows + (a per-architecture table; on SM107 the decode token counts) the split-K GEMM of + ``k3_fp8w_lowm_gemm`` runs instead, with 1 x 32 activation exponents and the same epilogue. The + outputs are new contiguous tensors; the current stream is used. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, global_sf: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return _latent.k3_latent_down_fp4(x, w, w_scale, global_sf) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, global_sf: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """``fp4_quantize`` (catalog reference, linear scales) of the BF16 projection of the + quantized activation (1 x 32 exponents where the op runs the low-M route, else 1 x 128).""" + m = x.shape[0] + n, k = w.shape + group = 32 if m <= _latent.lowm_route_max_m(n, k) else 128 + y = ref.lowm_gemm(x, w, w_scale, group) + codes, scales = fp4_quantize.reference(y, global_sf, 16, False, False) + return codes, scales.view(m, n // 16) + + +k3_latent_down_fp4 = K3LatentDownFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_mxfp8_gemm.py new file mode 100644 index 000000000000..48ed96c30a8e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/k3_mxfp8_gemm.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SM107 dense MXFP8 GEMM of an FP8 block-scale projection, autotuned over K3's tactics.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from . import _fp8_reference as ref +from . import _mxfp8_gemm + + +class K3Mxfp8Gemm(OpWrapper): + """BF16 ``input @ weight.T`` of E4M3 ``input [M, K]`` / ``weight [N, K]`` with their K32 R128c4 + UE8M0 scales (``fp8_quantize_1x128_cutedsl_ue8m0`` / the projection's prepared scale).""" + + ARCHS = frozenset({Arch.SM_107}) + #: The FP8 block-scale projections that reach the MXFP8 GEMM on SM107 (tp1 context, and the + #: shapes the low-M kernels leave to it), with rows from decode to a context chunk. + CELLS: tuple[Cell, ...] = ( + Cell( + why="MLA q_a|kv_a|k_pe (a partial last N tile) on a tp1 context chunk", + spec={"n": 2112, "k": 7168, "m": 1000}, + ), + Cell( + why="o_proj-like K = 1024 weight on 48 rows: a split-K bucket (BF16 atomics)", + spec={"n": 7168, "k": 1024, "m": 48}, + ), + Cell( + why="a decode row of the K = 7168, N = 4096 projection", + spec={"n": 4096, "k": 7168, "m": 1}, + ), + Cell( + why="K = 512 projection on a 4096-row context chunk", + spec={"n": 2304, "k": 512, "m": 4096}, + ), + ) + + note = """ + Main's CuTe DSL persistent block-scaled GEMM (``CuteDSLKimiK3MXFP8RubinLinear``) with K3's + tactic subset (``ops.gemm._mxfp8_gemm.keeps_tactic``); the AutoTuner picks the tactic per bucket of + M under the cache key ``ops.gemm._mxfp8_gemm.OP_NAME``. Split-K tactics accumulate BF16 atomically + (up to M = 256), so those buckets are not run-to-run deterministic. The output is a new + contiguous BF16 ``[M, N]`` tensor; the current stream is used. Raises on any GPU but SM107. + """ + + def __call__( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + ) -> torch.Tensor: + return _mxfp8_gemm.k3_mxfp8_gemm_rubin(input, weight, input_scale, weight_scale) + + def reference( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + ) -> torch.Tensor: + """FP64 product of the decoded operands (both scales K32 R128c4), rounded to BF16.""" + m, k = input.shape + n = weight.shape[0] + a = ref.dequant(input, ref.r128c4_exponents(input_scale, m, k), 32) + b = ref.dequant(weight, ref.r128c4_exponents(weight_scale, n, k), 32) + return ref.matmul_bf16(a, b) + + +k3_mxfp8_gemm = K3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_moe_residual_add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_moe_residual_add.py new file mode 100644 index 000000000000..d5cd19a5201e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_moe_residual_add.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MoE block's residual update with both BF16 rounding boundaries kept, in one launch.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.moe import kimi_k3_moe_residual as _residual + + +class K3MoeResidualAdd(OpWrapper): + """``residual + bf16(routed + shared)`` as PyTorch's two BF16 adds compute it; new tensor.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: Decode rows of every deployment that fuses the residual (helix8, helix16, dep16), up to the + #: kernel's 96-row bound, and one row above it (the PyTorch adds). + CELLS: tuple[Cell, ...] = ( + Cell(why="one row", spec={"m": 1}), + Cell(why="eight requests of 8 rows", spec={"m": 64}), + Cell(why="the kernel's 96-row bound", spec={"m": 96}), + Cell(why="above the bound: the PyTorch path", spec={"m": 128}), + ) + + note = """ + One Triton kernel for contiguous BF16 ``[M, 7168]`` operands of ``1 <= M <= 96`` rows on one + CUDA device (no autograd); every other input takes ``residual + (routed + shared)`` in PyTorch, + with the same result. The output is a new tensor shaped like ``routed``. + """ + + def __call__( + self, routed: torch.Tensor, shared: torch.Tensor, residual: torch.Tensor + ) -> torch.Tensor: + return _residual._fused_moe_residual_add(routed, shared, residual) + + def reference( + self, routed: torch.Tensor, shared: torch.Tensor, residual: torch.Tensor + ) -> torch.Tensor: + """FP32 adds rounded to BF16 after each, as the two PyTorch BF16 adds.""" + moe = (routed.float() + shared.float()).to(torch.bfloat16) + return (residual.float() + moe.float()).to(torch.bfloat16) + + +k3_moe_residual_add = K3MoeResidualAdd() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_moe_tail1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_moe_tail1.py new file mode 100644 index 000000000000..89967d291c62 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_moe_tail1.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MoE back half of at most 8 rows in one kernel: the routed latent all-reduce, its +RMSNorm and this rank's output columns of latent_up, multicast into the tail call's all-gather +rows.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_moe_tail1 as _tail1 + + +class K3MoeTail1(OpWrapper): + """For ``m`` rows on ``world`` ranks: polls the BF16 latent partials ``[m, 3584]`` every rank's + routed experts pushed into slot ``(token, rank)`` of the latent Lamport workspace (``lat_uc`` + local base, ``lat_flags`` its rotation state), sums them, applies the RMSNorm of ``gamma`` / + ``eps`` and writes rows ``[c0, c0 + 7168 / world)`` of ``latent_up`` (``weight`` E4M3 + ``[7168, 3584]``, ``wexp`` its UE8M0 block exponents ``[56, 28]``) into the all-gather rows of + the tail call (``ag_base``: multicast base, or a local one with ``mc_out=False``; ``ag_flags`` + its rotation state). Returns None; the next pre-attention mix finishes the call.""" + + # No cells yet: the op reads the partials from a Lamport workspace of every rank. + ARCHS: frozenset[Arch] = frozenset() + CELLS: tuple[Cell, ...] = () + + note = """ + A collective of the latent workspace's ranks. Clusters of 16 CTAs; each lane loads its + latent_up fragments before polling. The latent sum is the C++ one-shot all-reduce's (FP32 in + rank order, chunks of eight ranks, rounded to BF16); the RMS uses FP32 squares of it; the + projection's activation is ``bf16(x * gamma)`` and ``rsqrt(mean + eps)`` scales the FP32 + product sum before its rounding, so the output is not bitwise the column-parallel tail. + ``-0`` is stored as ``+0``. The kernel re-arms the previous call's latent buffer and, from the + CTA that arrives last, rotates the latent workspace after its ``griddepcontrol.wait``; the + tail call's state is only read. PDL: the dependent is released at entry. Compiled per + (world, m, mc_out, device). + """ + + def __call__( + self, + *, + m: int, + world: int, + lat_uc: int, + lat_flags: torch.Tensor, + weight: torch.Tensor, + wexp: torch.Tensor, + gamma: torch.Tensor, + eps: float, + ag_base: int, + ag_flags: torch.Tensor, + c0: int, + mc_out: bool = True, + ) -> None: + _tail1.run_k3_moe_tail1( + m=m, + world=world, + lat_uc=lat_uc, + lat_flags=lat_flags, + weight=weight, + wexp=wexp, + gamma=gamma, + eps=eps, + ag_base=ag_base, + ag_flags=ag_flags, + c0=c0, + mc_out=mc_out, + ) + + def reference( + self, + *, + m: int, + world: int, + lat_uc: Sequence[torch.Tensor], + lat_flags: torch.Tensor, + weight: torch.Tensor, + wexp: torch.Tensor, + gamma: torch.Tensor, + eps: float, + ag_base: int, + ag_flags: torch.Tensor, + c0: int, + mc_out: bool = True, + ) -> torch.Tensor: + """A collective: ``lat_uc`` is the group's input, the ranks' BF16 latent partials + ``[m, 3584]`` in rank order (what their experts pushed). Returns the BF16 ``[m, 7168]`` + all-gather rows once every rank ran: ``latent_up(rmsnorm(bf16(sum)))`` with the + normalization applied in FP64. The workspace and output addresses are not read.""" + total = lat_uc[0].float() + for part in lat_uc[1:]: + total = total + part.float() + latent = total.to(torch.bfloat16).double() + rms = torch.rsqrt(latent.pow(2).mean(-1, keepdim=True) + eps) + scale = torch.pow(2.0, wexp.double() - 127.0) + scale = scale.repeat_interleave(128, 0).repeat_interleave(128, 1) + w = weight.to(torch.float64) * scale + y = ((latent * gamma.double()) @ w.T) * rms + return y.to(torch.bfloat16) + + +k3_moe_tail1 = K3MoeTail1() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_route_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_route_tail.py new file mode 100644 index 000000000000..da02fd96b75d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_route_tail.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 routing of FP32 router logits: ``noaux_tc`` top-16 and this rank's ``moe_sort`` tables +in one launch.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.noaux_tc_op import noaux_tc_op + +from ...kernels import k3_route_tail as _route_tail + + +def sort_tables_reference( + ids: torch.Tensor, local_expert_offset: int, local_num_experts: int, tile: int +) -> tuple[torch.Tensor, ...]: + """``moe_sort``'s six tables of the top-k ``ids [M, top_k]`` for the local experts + ``[local_expert_offset, local_expert_offset + local_num_experts)``, int32: a counting sort by + local expert, rows of an expert in increasing expanded index (``token * top_k + slot``), each + expert's rows padded to whole ``tile``-row tiles in local-expert order. Entries the op leaves + unwritten are ``-1``.""" + dev = ids.device + num_tokens, width = ids.shape + num_expanded = num_tokens * width + local = ids.reshape(-1).long() - local_expert_offset + is_local = (local >= 0) & (local < local_num_experts) + counts = torch.bincount(local[is_local], minlength=local_num_experts) + tiles_per_expert = (counts + tile - 1) // tile + num_tiles = int(tiles_per_expert.sum()) + filled = min(local_num_experts, num_expanded) + max_tiles = filled + (num_expanded - filled) // tile + run_start = (torch.cumsum(tiles_per_expert, 0) - tiles_per_expert) * tile + tile_expert = torch.repeat_interleave( + torch.arange(local_num_experts, device=dev), tiles_per_expert + ) + tile_idx = torch.arange(num_tiles, device=dev) + mn_limit = torch.minimum((tile_idx + 1) * tile, run_start[tile_expert] + counts[tile_expert]) + expanded = torch.arange(num_expanded, device=dev)[is_local] + expert = local[is_local] + order = torch.sort(expert, stable=True).indices + expanded, expert = expanded[order], expert[order] + rank = torch.arange(expert.numel(), device=dev) - (torch.cumsum(counts, 0) - counts)[expert] + permuted = run_start[expert] + rank + + def unwritten(n: int) -> torch.Tensor: + return torch.full((n,), -1, dtype=torch.long, device=dev) + + tile_idx_to_expert_idx = unwritten(max_tiles) + tile_idx_to_expert_idx[:num_tiles] = tile_expert + tile_idx_to_mn_limit = unwritten(max_tiles) + tile_idx_to_mn_limit[:num_tiles] = mn_limit + permuted_idx_to_expanded_idx = unwritten(max_tiles * tile) + permuted_idx_to_expanded_idx[permuted] = expanded + expanded_idx_to_permuted_idx = unwritten(num_expanded) + expanded_idx_to_permuted_idx[expanded] = permuted + outs = ( + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx.view(num_tokens, width), + permuted_idx_to_expanded_idx, + torch.tensor([num_tiles * tile], device=dev), + torch.tensor([num_tiles], device=dev), + ) + return tuple(t.to(torch.int32) for t in outs) + + +class K3RouteTail(OpWrapper): + """``noaux_tc`` top-16 of FP32 router logits ``[M, 896]`` (or ``S`` split-K partials + ``[S, R, 896]``, summed in split order) and the ``moe_sort`` tables (tile 128) of the experts + ``[local_expert_offset, local_expert_offset + local_num_experts)``. + + Returns ``(topk_ids int32 [M, 16], topk_weights fp32 [M, 16], tile_idx_to_expert_idx, + tile_idx_to_mn_limit, expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, + total_num_padded_tokens, num_non_exiting_tiles)``, new tensors with ``moe_sort``'s shapes; + only the valid table entries are written. ``num_tokens``: ``M`` (default every row). + ``workspace``: above 256 tokens, the device's persistent multi-CTA sort workspace + (``k3_route_tail.wide_workspace``). + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: The local expert windows of the deployments (896 experts): tp1 (all), dep16 EP16 (56), + #: helix8 EP8 (112) and helix16 MoE TP4 x EP4 (224), at decode rows and above the cluster + #: kernel's 256 tokens; ``splits``: split-K partials ``[S, M, 896]`` summed by the op. + CELLS: tuple[Cell, ...] = ( + Cell( + why="tp1: every expert local, one row", + spec={"m": 1, "offset": 0, "local": 896, "splits": 0}, + ), + Cell( + why="dep16 EP16: a 56-expert window in the middle, 32 rows", + spec={"m": 32, "offset": 392, "local": 56, "splits": 0}, + ), + Cell( + why="helix8 EP8: the last 112-expert window, one request with seven draft tokens", + spec={"m": 8, "offset": 784, "local": 112, "splits": 0}, + ), + Cell( + why="helix16 MoE TP4 x EP4: a 224-expert window on two split-K logit partials", + spec={"m": 16, "offset": 224, "local": 224, "splits": 2}, + ), + Cell( + why="dep16 at 32 requests x 8 rows: the cluster kernel's 256-token limit", + spec={"m": 256, "offset": 840, "local": 56, "splits": 0}, + ), + Cell( + why="dep16 EP16 context, 896 tokens: the multi-CTA sort with its workspace (it serves " + "at most 224 local experts, so not tp1 above 256 tokens)", + spec={"m": 896, "offset": 392, "local": 56, "splits": 0}, + ), + ) + + note = """ + Up to 256 tokens one cluster of at most eight CTAs; 257..2048 tokens independent CTAs whose + last arrival (a counter in ``workspace``, reset by the same launch) builds the tables. The + top-k is bit-identical to ``noaux_tc_op`` on the same logits (selection, order, weights); an + expert's rows are in ascending token order. ``workspace`` must be created outside CUDA graph + capture; the kernels keep no other state between launches. Unsupported inputs raise ValueError. + """ + + def __call__( + self, + logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + routed_scaling_factor: float, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + num_tokens: int = -1, + workspace: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, ...]: + return _route_tail.k3_route_tail( + logits, + e_score_correction_bias, + routed_scaling_factor, + local_expert_offset, + local_num_experts, + tile_tokens_dim, + num_tokens, + workspace, + ) + + def reference( + self, + logits: torch.Tensor, + e_score_correction_bias: torch.Tensor, + routed_scaling_factor: float, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + num_tokens: int = -1, + workspace: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, ...]: + """The catalog ``noaux_tc_op`` reference (one group, top-16) on the logits, then + :func:`sort_tables_reference` of its top-k; ``workspace`` is not read.""" + if logits.dim() == 3: + rows = logits.shape[1] if num_tokens < 0 else num_tokens + total = logits[0, :rows].float() + for split in range(1, logits.shape[0]): + total = total + logits[split, :rows].float() + logits = total + elif num_tokens >= 0: + logits = logits[:num_tokens] + weights, ids = noaux_tc_op.reference( + logits, e_score_correction_bias, 1, 1, _route_tail.TOP_K, routed_scaling_factor + ) + tables = sort_tables_reference(ids, local_expert_offset, local_num_experts, tile_tokens_dim) + return (ids, weights.float(), *tables) + + +k3_route_tail = K3RouteTail() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_router_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_router_logits.py new file mode 100644 index 000000000000..36c467e46d29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_router_logits.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 router logits of up to 16 tokens with the router weight read before +``griddepcontrol.wait``.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ...kernels import k3_router_exact as _router + + +class K3RouterLogits(OpWrapper): + """FP32 ``[M, 896]`` logits ``hidden_states @ weight.T`` of BF16 ``hidden_states [M, 7168]`` + (``1 <= M <= 16``) and the BF16 ``[896, 7168]`` router weight.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: Decode rows of the helix deployments (router_exact_kda): one row, one and two requests. + CELLS: tuple[Cell, ...] = ( + Cell(why="one row: decode without MTP", spec={"m": 1}), + Cell(why="one request with seven draft tokens", spec={"m": 8}), + Cell(why="two requests: the op's 16-row limit", spec={"m": 16}), + ) + + note = """ + Bit-identical to ``dsv3_router_gemm_op(hidden_states, weight.t())``: the same 14 K slices of + 512, ``mma.sync`` fragments and slice-order sum, cut over clusters of two CTAs that hand the + running sum over distributed shared memory. Each CTA copies its 112 KB weight slice into shared + memory by TMA before ``griddepcontrol.wait`` and releases its dependents right after it. The + output is a new FP32 tensor; unsupported inputs raise ValueError (``k3_router_exact.can_run``). + """ + + def __call__(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return _router.k3_router_logits(hidden_states, weight) + + def reference(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """FP32 product of the BF16 operands (IEEE FP32 accumulation).""" + with true_fp32_matmul(): + return hidden_states.float() @ weight.float().T + + +k3_router_logits = K3RouterLogits() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_shared_persist.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_shared_persist.py new file mode 100644 index 000000000000..b6f51c633bf7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_shared_persist.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 shared expert of at most 8 rows on a few persistent CTAs beside the routed experts, +pushing its row-parallel partial into the MoE tail's Lamport call.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_shared_persist as _persist + + +class K3SharedPersist(OpWrapper): + """This rank's shared-expert partial ``situ(x Wg^T, x Wu^T) Wd^T`` (BF16 ``[m, 7168]``) of BF16 + ``x [m, 7168]``, stored into slot ``(token, rank)`` of the tail call's current buffer + (``out_base``: multicast base, or a local one with ``mc_out=False``; ``out_flags`` its rotation + state); returns None. ``frags``: ``k3_shared_persist.persist_fragments`` of the rank's gate/up + and down weights for ``sms`` CTAs; ``ws``: the exchange buffer and counters + (``k3_shared_persist.PersistWorkspace``).""" + + # No cells yet: the op writes into a Lamport workspace the next pre-attention mix finishes. + ARCHS: frozenset[Arch] = frozenset() + CELLS: tuple[Cell, ...] = () + + note = """ + ``sms`` persistent CTAs stream their weights once through a shared-memory ring, compute the + gate/up products (BF16 ``mma.sync``, FP32 accumulation, outputs rounded to BF16), the SiTU + activation with ``ex2.approx`` tanh / sigmoid (rounded to BF16), exchange the activation + through ``ws`` and multiply their down m-tiles; ``-0`` is stored as ``+0``. Not bitwise the + GEMM + activation-op path. Launched without PDL; consecutive launches on a stream reuse ``ws`` + (its counters are zero between launches). Compiled per (world, m, sms, mc_out, betas, device). + """ + + def __call__( + self, + *, + x: torch.Tensor, + frags: torch.Tensor, + world: int, + rank: int, + sms: int, + ws: "_persist.PersistWorkspace", + out_base: int, + out_flags: torch.Tensor, + beta: float, + linear_beta: Optional[float], + mc_out: bool = True, + ) -> None: + _persist.run_k3_shared_persist( + x=x, + frags=frags, + world=world, + rank=rank, + sms=sms, + ws=ws, + out_base=out_base, + out_flags=out_flags, + beta=beta, + linear_beta=linear_beta, + mc_out=mc_out, + ) + + def reference( + self, + *, + x: torch.Tensor, + frags: Tuple[torch.Tensor, torch.Tensor], + world: int, + rank: int, + sms: int, + ws: "_persist.PersistWorkspace", + out_base: int, + out_flags: torch.Tensor, + beta: float, + linear_beta: Optional[float], + mc_out: bool = True, + ) -> torch.Tensor: + """The BF16 ``[m, 7168]`` slot rows. ``frags`` is the ``(gate_up [2 I, 7168], down + [7168, I])`` BF16 pair ``persist_fragments`` packs (the packing is a permutation). FP32 + GEMMs rounded to BF16 like the op's, SiTU in FP32 rounded to BF16. The workspace and + output addresses are not read.""" + gate_up, down = frags + inter = gate_up.shape[0] // 2 + xf = x.float() + g = (xf @ gate_up[:inter].float().T).to(torch.bfloat16).float() + u = (xf @ gate_up[inter:].float().T).to(torch.bfloat16).float() + act = beta * torch.tanh(g / beta) * torch.sigmoid(g) + if linear_beta is not None: + u = linear_beta * torch.tanh(u / linear_beta) + act = (act * u).to(torch.bfloat16) + return (act.float() @ down.float().T).to(torch.bfloat16) + + +k3_shared_persist = K3SharedPersist() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_small_batch_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_small_batch_moe.py new file mode 100644 index 000000000000..10bdcdb6d037 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_small_batch_moe.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 routed NVFP4 SiTU experts of at most 32 tokens in one persistent kernel (FC1 + SiTU + +NVFP4 requantization + FC2 + weighted combine).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_small_batch_moe_op as _op +from .k3_route_tail import sort_tables_reference + +_SF_VEC = 16 +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def _fp4_to_f64(packed: torch.Tensor) -> torch.Tensor: + lut = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + p = packed.contiguous().view(torch.uint8).long() + return torch.stack([lut[p & 0xF], lut[p >> 4]], dim=-1).flatten(-2) + + +def _e4m3_to_f64(sf: torch.Tensor) -> torch.Tensor: + return sf.contiguous().view(torch.uint8).view(torch.float8_e4m3fn).to(torch.float64) + + +def _unswizzle_128x4(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Row-major ``[rows, cols]`` scale bytes from the 128x4-blocked layout.""" + padded_cols = -(-cols // 4) * 4 + r = torch.arange(rows, device=sf.device).view(-1, 1) + c = torch.arange(cols, device=sf.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * (4 * 128) + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + return sf.reshape(-1)[index] + + +def _dequant_expert(w: torch.Tensor, sf: torch.Tensor, e: int) -> torch.Tensor: + rows, cols = w.shape[1], w.shape[2] * 2 + lin = _unswizzle_128x4(sf[e].contiguous().view(torch.uint8), rows, cols // _SF_VEC) + return _fp4_to_f64(w[e]) * _e4m3_to_f64(lin).repeat_interleave(_SF_VEC, dim=1) + + +def _nvfp4_roundtrip(h: torch.Tensor, global_sf: float) -> torch.Tensor: + """``h`` quantized as the FC1 epilogue quantizes it (16-element blocks, E4M3 scales of + ``amax / 6 * global_sf``), decoded in the scaled domain the FC2 alphas expect.""" + grid = torch.tensor(_E2M1[:8], dtype=torch.float64, device=h.device) + blocks = h.view(h.size(0), -1, _SF_VEC) * global_sf + amax = blocks.abs().amax(-1, keepdim=True) + sf = (amax / 6.0).float().to(torch.float8_e4m3fn).double() + v = torch.where(sf > 0, blocks / sf, torch.zeros_like(blocks)).clamp(-6.0, 6.0) + q = grid[(v.abs()[..., None] - grid).abs().argmin(-1)] * v.sign() + return (q * sf).view_as(h) + + +class K3SmallBatchMoe(OpWrapper): + """Writes into every element of ``output [tokens, hidden]`` (BF16) the sum of this rank's + experts' routing-weighted down-projections; returns None. ``counters`` is zero again after + the launch. + + ``input`` / ``input_scale``: NVFP4 tokens and their linear E4M3 scales; ``fc1_*`` / ``fc2_*``: + the local experts' NVFP4 weights ``[L, rows, cols / 2]``, 128x4-blocked E4M3 scales and FP32 + per-expert alphas (FC1 rows in ``[gate16, up16]`` groups); the ``moe_sort`` tables of tile 128; + ``global_sf``: the FC2 input scale; ``counters``: ``new_small_batch_moe_counters``. + ``token_indexed``: the token-indexed FC1 output layout (MoE-TP slices). ``published_routing`` + (token-indexed): ``expanded_idx_to_permuted_idx`` holds the top-k ids a routing kernel + published into the routing word of ``counters``, the other four tables are None. + ``push_world > 0``: the rows are also pushed into slot ``(token, push_rank)`` of the one-shot + MNNVL all-reduce whose multicast base is ``push_mc`` and state ``push_flags``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + #: K3 routed experts (896 experts, top 16, latent 3584) per deployment, at decode rows: + #: ``local`` experts from ``offset``, ``inter`` per rank, ``layout`` ``permuted`` or + #: ``token`` (token-indexed), ``published``: routing published by the routing kernel. + CELLS: tuple[Cell, ...] = ( + Cell( + why="dep16 EP16: 56 experts with the full intermediate, one row", + spec={"m": 1, "local": 56, "offset": 392, "inter": 3072, "layout": "permuted"}, + ), + Cell( + why="helix8 EP8: 112 experts, one request with seven draft tokens", + spec={"m": 8, "local": 112, "offset": 336, "inter": 3072, "layout": "permuted"}, + ), + Cell( + why="helix16 MoE TP4 x EP4: 224 experts, a quarter of the intermediate, token-indexed, " + "four requests", + spec={"m": 32, "local": 224, "offset": 224, "inter": 768, "layout": "token"}, + ), + Cell( + why="helix16 ULL: the same window on the routing the front kernel publishes", + spec={ + "m": 8, + "local": 224, + "offset": 672, + "inter": 768, + "layout": "token", + "published": True, + }, + ), + ) + + note = """ + The default combine adds the weighted rows with BF16 vector reductions, whose order (so the + output's bits) can change from run to run; ``deterministic_combine`` sums the slices in a fixed + order (bitwise independent of the CTA count). ``num_ctas`` > 0 overrides the persistent grid, + which is otherwise one CTA per SM less ``sm_reserve`` and ``cute_dsl_utils.sm_reserve()``. + Compiled kernels are cached per process by (token tile, shapes, top-k, SiTU betas, layout, + CTA count, combine, routing, push world, PDL); the first call of each key compiles (outside + CUDA graph capture). ``die_major`` raises (not supported). + """ + + def __call__( + self, + input: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: Optional[torch.Tensor], + tile_idx_to_mn_limit: Optional[torch.Tensor], + permuted_idx_to_expanded_idx: Optional[torch.Tensor], + expanded_idx_to_permuted_idx: torch.Tensor, + num_non_exiting_tiles: Optional[torch.Tensor], + global_sf: torch.Tensor, + token_final_scales: torch.Tensor, + counters: torch.Tensor, + top_k: int, + situ_beta: float, + situ_linear_beta: float, + token_indexed: bool, + num_ctas: int = 0, + deterministic_combine: bool = False, + die_major: bool = False, + published_routing: bool = False, + expert_offset: int = 0, + push_mc: int = 0, + push_flags: Optional[torch.Tensor] = None, + push_rank: int = 0, + push_world: int = 0, + sm_reserve: int = 0, + ) -> None: + _op.k3r_nvfp4_small_batch_moe( + input, + input_scale, + fc1_weight, + fc1_weight_scale, + fc1_alpha, + fc2_weight, + fc2_weight_scale, + fc2_alpha, + output, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + expanded_idx_to_permuted_idx, + num_non_exiting_tiles, + global_sf, + token_final_scales, + counters, + top_k, + situ_beta, + situ_linear_beta, + token_indexed, + num_ctas, + deterministic_combine, + die_major, + published_routing, + expert_offset, + push_mc, + push_flags, + push_rank, + push_world, + sm_reserve, + ) + + def reference( + self, + input: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: Optional[torch.Tensor], + tile_idx_to_mn_limit: Optional[torch.Tensor], + permuted_idx_to_expanded_idx: Optional[torch.Tensor], + expanded_idx_to_permuted_idx: torch.Tensor, + num_non_exiting_tiles: Optional[torch.Tensor], + global_sf: torch.Tensor, + token_final_scales: torch.Tensor, + counters: torch.Tensor, + top_k: int, + situ_beta: float, + situ_linear_beta: float, + token_indexed: bool, + num_ctas: int = 0, + deterministic_combine: bool = False, + die_major: bool = False, + published_routing: bool = False, + expert_offset: int = 0, + push_mc: int = 0, + push_flags: Optional[torch.Tensor] = None, + push_rank: int = 0, + push_world: int = 0, + sm_reserve: int = 0, + ) -> torch.Tensor: + """The expected ``output`` in FP64 (round to BF16 to compare bits): per served slot, FC1 + of the decoded operands, SiTU, the FC1 epilogue's NVFP4 requantization, FC2, the routing + weight, summed per token. The served slots are the tables' (with ``published_routing``, + the top-k ids of the experts ``[expert_offset, expert_offset + L)``). The layout, combine, + CTA count and push arguments do not change the value.""" + num_tokens = token_final_scales.shape[0] + num_local = fc1_weight.shape[0] + hidden = fc2_weight.shape[1] + inter = fc2_weight.shape[2] * 2 + if published_routing: + ids = expanded_idx_to_permuted_idx.reshape(num_tokens, top_k) + t2g, t2lim, _, p2e, _, nnet = sort_tables_reference(ids, expert_offset, num_local, 128) + else: + t2g, t2lim, p2e, nnet = ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + ) + num_tiles = int(nnet[0]) + limits = t2lim[:num_tiles].tolist() + rows = [(t, r) for t in range(num_tiles) for r in range(t * 128, limits[t])] + out = torch.zeros(num_tokens, hidden, dtype=torch.float64, device=input.device) + if not rows: + return out + dev = input.device + row_idx = torch.tensor([r for _, r in rows], device=dev) + tiles = torch.tensor([t for t, _ in rows], device=dev) + expanded = p2e[row_idx].long() + tokens, slots = expanded // top_k, expanded % top_k + experts = t2g[tiles].long() + x_sf = _e4m3_to_f64(input_scale).reshape(input.shape[0], -1) + x = _fp4_to_f64(input) * x_sf.repeat_interleave(_SF_VEC, dim=1) + h = torch.empty(len(rows), inter, dtype=torch.float64, device=dev) + for e in experts.unique().tolist(): + sel = (experts == e).nonzero().flatten() + w1 = _dequant_expert(fc1_weight, fc1_weight_scale, e) + acc = (x[tokens[sel]] @ w1.t()) * float(fc1_alpha[e]) + acc = acc.view(len(sel), -1, 2, 16) # [gate16, up16] row groups + gate, up = acc[:, :, 0], acc[:, :, 1] + act = situ_beta * torch.tanh(gate / situ_beta) * torch.sigmoid(gate) + act = act * situ_linear_beta * torch.tanh(up / situ_linear_beta) + h[sel] = act.reshape(len(sel), -1) + h = _nvfp4_roundtrip(h, float(global_sf.reshape(-1)[0])) + partial = torch.empty(len(rows), hidden, dtype=torch.float64, device=dev) + for e in experts.unique().tolist(): + sel = (experts == e).nonzero().flatten() + w2 = _dequant_expert(fc2_weight, fc2_weight_scale, e) + partial[sel] = h[sel] @ w2.t() * float(fc2_alpha[e]) + weights = token_final_scales[tokens, slots].double()[:, None] + out.index_add_(0, tokens, partial * weights) + return out + + +k3_small_batch_moe = K3SmallBatchMoe() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_ull_front.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_ull_front.py new file mode 100644 index 000000000000..67e4b7cac365 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/k3_ull_front.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MoE front of small decode steps over the TP group: router logits, routing tail and the +latent down-projection with its NVFP4 quantization, sharded over the ranks in two kernels.""" + +from __future__ import annotations + +from typing import Callable, Optional, Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) + +from ...kernels import k3_ull_front as _front +from ..gemm import _fp8_reference as gemm_ref +from .k3_route_tail import k3_route_tail +from .k3_router_logits import k3_router_logits + + +class K3UllFront(OpWrapper): + """Returns ``(routing, codes, scales)`` for this rank's BF16 ``x [M, 7168]`` (``M <= 32``, the + same rows on every rank of the group): + + * ``routing``: the 8 outputs of :class:`K3RouteTail` on the router logits of every expert, with + the ``moe_sort`` tables of ``plan.local_offset`` / ``plan.local_num`` (with ``publish``, the + address of a zero int32 word, the route kernel adds 1 to it per routed token instead and + leaves the tables unwritten); + * ``codes`` uint8 ``[M, 1792]`` / ``scales`` uint8 ``[M, 224]``: :class:`K3LatentDownFp4` of + ``x`` (linear scales), each rank computing its ``plan`` tiles of latent_down and multicasting + them through ``ws``. + + ``plan``: ``k3_ull_front.make_plan`` of the layer; ``ws``: the TP group's + ``k3_ull_front.ull_front_workspace``. ``after_head`` runs between the two launches. + """ + + # No cells yet: the op needs an MNNVL multicast workspace over a TP group of ranks. + ARCHS: frozenset[Arch] = frozenset() + CELLS: tuple[Cell, ...] = () + + note = """ + A collective over ``plan.world`` ranks; every rank calls it in the same order. Head kernel: + the latent GEMM CTAs (the low-M split-K MXFP8 GEMM, one-shot push of this rank's tiles into + slot ``(token, rank)`` of every rank's workspace half) and the router CTAs (``K3RouterLogits``' + arithmetic on all 896 experts into this rank's own workspace). Route kernel, a PDL dependent of + the head: the ``K3RouteTail`` cluster polling the logits and gather CTAs polling every rank's + latent slots, then ``fp4_quantize``'s arithmetic. Every reader re-arms what it read; + consecutive calls alternate between the workspace's two halves (``ws.calls`` is advanced on + the host). The first call of each shape compiles, which must not happen inside CUDA graph + capture (``k3_ull_front.is_compiled``). + """ + + def __call__( + self, + x: torch.Tensor, + plan: "_front.UllFrontPlan", + ws: "_front.UllFrontWorkspace", + after_head: Optional[Callable[[], None]] = None, + publish: int = 0, + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: + return _front.run_k3_ull_front(x, plan, ws, after_head=after_head, publish=publish) + + def reference( + self, + x: Sequence[torch.Tensor], + plan: Sequence["_front.UllFrontPlan"], + ws: "_front.UllFrontWorkspace", + after_head: Optional[Callable[[], None]] = None, + publish: int = 0, + ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]]: + """A collective: ``x`` and ``plan`` are the group's in rank order; returns each rank's + ``(routing, codes, scales)`` in rank order, composed from :class:`K3RouterLogits`, + :class:`K3RouteTail` and the latent projection of the ranks' tiles together (activation + MXFP8 1 x 32, then ``fp4_quantize`` with ``plan.global_sf``). ``ws``, ``after_head`` and + ``publish`` are not read (with ``publish`` only the top-k is defined).""" + latent_w = torch.cat([p.latent_w for p in plan]) + exps = torch.cat( + [ + gemm_ref.packed_block_exponents(p.latent_ws, p.latent_w.shape[0], _front.HIDDEN) + for p in plan + ] + ) + w = gemm_ref.dequant(latent_w, exps, 128) + out = [] + for xr, p in zip(x, plan): + logits = k3_router_logits.reference(xr, p.router_w) + routing = k3_route_tail.reference( + logits, + p.bias, + p.routed_scaling_factor, + p.local_offset, + p.local_num, + _front.TILE, + ) + codes, e = gemm_ref.mxfp8_quantize(xr, 32) + y = gemm_ref.matmul_bf16(gemm_ref.dequant(codes, e, 32), w) + q, sf = fp4_quantize.reference(y, p.global_sf, 16, False, False) + out.append((routing, q, sf.view(xr.shape[0], -1))) + return out + + +k3_ull_front = K3UllFront() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/k3_attn_res_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/k3_attn_res_norm.py new file mode 100644 index 000000000000..12b7f63248b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/k3_attn_res_norm.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 decode attention-residual mix + RMSNorm, optionally after the residual add, in one PDL +kernel (CuTe DSL, ``kernels/k3_attn_res.py``).""" + +from __future__ import annotations + +from typing import Tuple, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.norm.attn_res_fwd import ( + attn_res_fwd, +) + +from ...kernels import k3_attn_res as _kernel + +#: Rows of one mid-M scale-atom block (``mx_atom``). +_ATOM_ROWS = 128 + + +def mxfp8_rows_reference(rows: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 form ``(xq, xs)`` of BF16 ``rows [M, K]`` that the low-M FP8 GEMM quantizes its + BF16 operand to: per 1x32 group the UE8M0 exponent of ``2^ceil(log2(amax / 448))`` (amax + floored at the BF16 value 0x38D2, ~1e-4), and each value times ``2^(127 - e)`` rounded to E4M3 + (round to nearest, saturating). ``xs`` is ``uint8 [M, K / 32]``.""" + m, k = rows.shape + groups = rows.contiguous().view(torch.int16).to(torch.int32).reshape(m, k // 32, 32) & 0x7FFF + amax = groups.amax(dim=-1).clamp_min(0x38D2) + exponent = (amax >> 7) - 8 + ((amax & 0x7F) > 0x60).to(torch.int32) + scale = torch.exp2((127 - exponent).to(torch.float32)).repeat_interleave(32, dim=1) + scaled = (rows.float() * scale).clamp(-448.0, 448.0) + return scaled.to(torch.float8_e4m3fn), exponent.to(torch.uint8) + + +def mx_atom_offsets(num_rows: int, hidden: int, device: torch.device) -> torch.Tensor: + """Byte offsets, row-major over ``[num_rows, hidden / 32]``, of the scales ``mx_atom`` stores: + group ``g`` of row ``t`` at ``(g / 4) * 512 + (t % 32) * 16 + (t / 32) * 4 + g % 4`` (one + 128-row block of the mid-M FP8 GEMM's tcgen05 scale atoms).""" + t = torch.arange(num_rows, device=device).unsqueeze(1) + g = torch.arange(hidden // 32, device=device).unsqueeze(0) + return ((g >> 2) * 512 + (t & 31) * 16 + ((t >> 5) & 3) * 4 + (g & 3)).reshape(-1) + + +def attn_res_rmsnorm_reference( + layer: torch.Tensor, + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, +) -> torch.Tensor: + """``bf16(bf16(mixed * rsqrt(mean(mixed^2) + out_eps)) * out_weight)`` of the family + ``attn_res_fwd`` mixture ``mixed`` of ``block [N - 1, M, H]`` and ``layer [M, H]``.""" + mixed, _, _, _ = attn_res_fwd.reference(layer, block, res_weight, rms_weight, rms_eps) + mixed = mixed.float() + normed = mixed * torch.rsqrt(mixed.square().mean(dim=-1, keepdim=True) + out_eps) + return out_weight.reshape(-1) * normed.to(torch.bfloat16) + + +class K3AttnResNorm(OpWrapper): + """RMSNorm of the attention-residual mix of ``block`` and the layer row, with the layer row + + * ``addends == ()``: ``layer`` itself; returns ``output``; + * ``(addend,)``: ``updated = bf16(layer + addend)``; returns ``(updated, output)``; + * ``(routed, shared)``: ``updated = bf16(layer + bf16(routed + shared))`` (the previous MoE + layer's residual update); returns ``(updated, output)``, with ``mxfp8`` also the MXFP8 form + ``(xq, xs)`` of ``output`` (``xs`` in the mid-M GEMM's scale layout with ``mx_atom``). + + ``layer`` and the addends are BF16 ``[M, 7168]``, ``block`` BF16 ``[N - 1, M, 7168]`` (the + snapshots), the three weights BF16 ``[7168]``; outputs are new tensors. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # rows: tokens of the decode step; candidates: N = snapshots + 1; addends: 0 (the mix alone), + # 1 (the residual add) or 2 (the deferred MoE residual update). tp1 (context) never runs it. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=why, + spec=dict( + rows=rows, + candidates=candidates, + hidden=7168, + addends=addends, + mxfp8=mxfp8, + entry_trigger=entry_trigger, + mx_atom=mx_atom, + rms_eps=1e-6, + out_eps=1e-6, + ), + ) + for why, rows, candidates, addends, mxfp8, entry_trigger, mx_atom in ( + ("helix16 ULL 1-request step, mix alone: 8 rows, N=2", 8, 2, 0, False, False, False), + ( + "dep16 per-rank step, mix alone: 64 rows (its cap), N=12", + 64, + 12, + 0, + False, + False, + False, + ), + ("helix8 b16 mtp5 verify, mix alone: 96 rows, N=9", 96, 9, 0, False, False, False), + ("one token, residual add, N=3: 4-CTA cluster", 1, 3, 1, False, False, False), + ("helix8 A_mid add: 16 rows, N=5, 8-CTA cluster", 16, 5, 1, False, False, False), + ("helix8 add at its 128-row cap, N=4", 128, 4, 1, False, False, False), + ("dep16 deferred MoE residual: 32 rows, N=7", 32, 7, 2, False, False, False), + ("helix8 deferred MoE residual: 96 rows, N=12", 96, 12, 2, False, False, False), + ( + "helix8 MoE add + MXFP8 for the KDA QKVG low-M GEMM: 16 rows, N=8", + 16, + 8, + 2, + True, + False, + False, + ), + ("helix16 ULL MoE add + MXFP8, entry trigger: 8 rows, N=6", 8, 6, 2, True, True, False), + ( + "helix8 MoE add + MXFP8 in the mid-M GEMM's scale atoms: 96 rows, N=9", + 96, + 9, + 2, + True, + False, + True, + ), + ( + "helix16 ULL 4-request step, MXFP8 atoms + entry trigger: 32 rows, N=4", + 32, + 4, + 2, + True, + True, + True, + ), + ) + ) + + note = """ + Decode-size fusion: up to TRTLLM_K3_ATTN_RES_PDL_MAX_M rows (helix8: 128, + others: 64) where `kernels/k3_attn_res.can_use_k3_attn_res` holds (H 7168, + 2 <= N <= 12, M <= 1024, contiguous BF16). Inputs are read-only. + A cluster of 4 or 8 CTAs per token (`fast_cluster_size`) loads the weights + and snapshot rows and reduces their statistics before griddepcontrol.wait; + only the layer row and the addends are read after it. The updated prefix + has the bits of the unfused adds; the normed output can differ from the + C++ per-token kernels by one BF16 step in a few elements (other sum order). + A process-wide cache (`kernels/k3_attn_res._compiled`) holds one compiled + kernel per (N, add mode, PDL, cluster, mxfp8, entry_trigger, mx_atom); + compile each key eagerly before CUDA-graph capture. TRTLLM_ENABLE_PDL is + read at import. `entry_trigger` releases the dependent grid at kernel entry + and prefers the largest shared-memory carve-out; `mx_atom` (M <= 128) stores + the scales of one 128-row block of the mid-M GEMM's atoms, rows past M are + not written. + """ + + def __call__( + self, + layer: torch.Tensor, + addends: Tuple[torch.Tensor, ...], + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + mxfp8: bool = False, + entry_trigger: bool = False, + mx_atom: bool = False, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + if len(addends) == 2: + routed, shared = addends + if mxfp8: + return _kernel.k3_attn_res_moe_add_rmsnorm_mx( + layer, + routed, + shared, + block, + res_weight, + rms_weight, + out_weight, + rms_eps, + out_eps, + entry_trigger, + mx_atom, + ) + return _kernel.k3_attn_res_moe_add_rmsnorm( + layer, routed, shared, block, res_weight, rms_weight, out_weight, rms_eps, out_eps + ) + if mxfp8: + raise ValueError("the MXFP8 output is served with the (routed, shared) addends only") + if addends: + return _kernel.k3_attn_res_add_rmsnorm( + layer, addends[0], block, res_weight, rms_weight, out_weight, rms_eps, out_eps + ) + return _kernel.k3_attn_res_rmsnorm( + layer, block, res_weight, rms_weight, out_weight, rms_eps, out_eps + ) + + def reference( + self, + layer: torch.Tensor, + addends: Tuple[torch.Tensor, ...], + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + mxfp8: bool = False, + entry_trigger: bool = False, + mx_atom: bool = False, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + """The adds in FP32 rounded to BF16 where the kernel rounds, then the family + ``attn_res_fwd`` mixture and the twice-rounded trailing RMSNorm. With ``mxfp8`` the + scales come row-major ``[M, 224]``, or with ``mx_atom`` at :func:`mx_atom_offsets` of a + zeroed buffer of the op's size. ``entry_trigger`` only changes scheduling.""" + if len(addends) == 2: + routed, shared = addends + moe = (routed.float() + shared.float()).to(torch.bfloat16) + row = (layer.float() + moe.float()).to(torch.bfloat16) + elif addends: + row = (layer.float() + addends[0].float()).to(torch.bfloat16) + else: + row = layer + output = attn_res_rmsnorm_reference( + row, block, res_weight, rms_weight, out_weight, rms_eps, out_eps + ) + if not addends: + return output + if not mxfp8: + return row, output + xq, xs = mxfp8_rows_reference(output) + if mx_atom: + hidden = layer.shape[1] + atom = torch.zeros((hidden // _ATOM_ROWS) * 512, dtype=torch.uint8, device=layer.device) + atom[mx_atom_offsets(layer.shape[0], hidden, layer.device)] = xs.reshape(-1) + xs = atom + return row, output, xq, xs + + def is_valid( + self, + layer: torch.Tensor, + addends: Tuple[torch.Tensor, ...], + block: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + out_weight: torch.Tensor, + rms_eps: float, + out_eps: float, + mxfp8: bool = False, + entry_trigger: bool = False, + mx_atom: bool = False, + ) -> None: + # The launcher does not re-validate: a strided, mis-sized or foreign-device operand is + # read as dense rows of 7168 and gives a plausible wrong mix. + assert _kernel.can_use_k3_attn_res( + layer, tuple(addends), block, res_weight, rms_weight, out_weight + ), "outside can_use_k3_attn_res (H 7168, 2 <= N <= 12, M <= 1024, contiguous BF16)" + # Without the MXFP8 output both flags are dropped without an error. + assert mxfp8 or not (entry_trigger or mx_atom), "entry_trigger / mx_atom need mxfp8" + + +k3_attn_res_norm = K3AttnResNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/_dspark_mla_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/_dspark_mla_rope.py new file mode 100644 index 000000000000..0b7f508a6d26 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/_dspark_mla_rope.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""RMSNorm + RoPE launchers of the Kimi K3 MLA DSpark drafter. + +* :func:`k3r_dspark_rmsnorm_rope`: main's ``trtllm::cute_dsl_dspark_rmsnorm_rope`` (same + kernel, same support check), compiled for + regular row-strided input, so the drafter normalizes column slices of its fused GEMM outputs + in place instead of staging them in a contiguous copy. Main's op is compiled for compact + rows. +* :func:`dspark_mla_ctx_norm_rope` / :func:`dspark_mla_block_rope` (``TLLM_DSPARK_TAIL_V2``): + the fused context-latent and block kernels of ``kernels/dspark_mla_rope.py``; bitwise the + per-layer launches they replace. + +Every operand is detached before the DLPack export, which refuses tensors that require grad +(module parameters do, also under inference mode). +""" + +import functools + +import cutlass +import cutlass.cute as cute +import torch + +from tensorrt_llm._torch.attention.kernels.blackwell.dspark_rmsnorm_rope import ( + DSparkRMSNormRoPEKernel, +) +from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( + is_fused_dspark_rmsnorm_rope_supported, +) +from tensorrt_llm._utils import get_sm_version + +from ...kernels.dspark_mla_rope import DSparkMlaBlockRoPEKernel, DSparkMlaCtxNormRoPEKernel + + +@functools.cache +def _compile_rmsnorm_rope_strided( + hidden_dim: int, + rope_dim: int, + num_heads: int, + eps: float, + apply_weight: bool, + apply_rmsnorm: bool, + inverse_rope: bool, + norm_dim: int, +): + """``DSparkRMSNormRoPEKernel`` for ``[rows, hidden_dim]`` input of any row stride (unit inner + stride), writing compact ``[rows, hidden_dim]`` output.""" + rows = cute.sym_int() + freq_rows = cute.sym_int() + # Regular row-strided input, as is_fused_dspark_rmsnorm_rope_supported admits: the MLA + # drafter normalizes column slices of fused GEMM outputs. + x_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (rows, hidden_dim), + stride=(cute.sym_int64(), 1), + ) + # None, not a fake tensor, when the kernel does not scale by a weight: the operand then does + # not exist in the compiled signature. + weight_fake = ( + cute.runtime.make_fake_compact_tensor(cutlass.BFloat16, (norm_dim,), stride_order=(0,)) + if apply_weight + else None + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (freq_rows, cute.sym_int(), 2), + stride_order=(2, 1, 0), + ) + output_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (rows, hidden_dim), stride_order=(1, 0) + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkRMSNormRoPEKernel( + hidden_dim, + rope_dim, + num_heads, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + norm_dim=norm_dim, + ) + return cute.compile( + kernel, + x_fake, + weight_fake, + freqs_fake, + output_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + +def k3r_dspark_rmsnorm_rope( + x: torch.Tensor, + weight: torch.Tensor | None, + freqs: torch.Tensor, + num_heads: int, + rope_dim: int, + eps: float, + apply_weight: bool, + apply_rmsnorm: bool, + inverse_rope: bool, + norm_dim: int | None = None, +) -> torch.Tensor: + """``cute_dsl_dspark_rmsnorm_rope`` on regular row-strided BF16 rows, read in place. + + ``weight`` may be None when ``apply_weight`` is False. The result is compact. + """ + if apply_weight and weight is None: + raise ValueError("k3r_dspark_rmsnorm_rope needs a weight when apply_weight is set") + weight = weight.detach() if apply_weight else None + if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim, norm_dim): + raise ValueError( + "k3r_dspark_rmsnorm_rope requires regular row-strided BF16 tensors on an SM100, " + f"SM103, or SM107 GPU with a valid FP32 frequency view; got SM {get_sm_version()}" + ) + original_shape = x.shape + x_flat = x.detach().view(-1, x.shape[-1]) + output = torch.empty(x_flat.shape, dtype=x.dtype, device=x.device) + compiled = _compile_rmsnorm_rope_strided( + x.shape[-1], + rope_dim, + num_heads, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + x.shape[-1] if norm_dim is None else norm_dim, + ) + compiled(x_flat, weight, freqs.detach(), output) + return output.view(original_shape) + + +@functools.cache +def _compile_dspark_mla_ctx_norm_rope( + num_layers: int, kv_lora_rank: int, rope_dim: int, eps: float +): + head_dim = kv_lora_rank + rope_dim + tokens = cute.sym_int() + x_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, (tokens, num_layers * head_dim), stride=(cute.sym_int64(), 1) + ) + weight_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (num_layers, kv_lora_rank), stride_order=(1, 0) + ) + positions_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, (tokens,), stride_order=(0,) + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), rope_dim // 2, 2), stride_order=(2, 1, 0) + ) + output_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (cute.sym_int(), head_dim), stride_order=(1, 0) + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkMlaCtxNormRoPEKernel(num_layers, kv_lora_rank, rope_dim, eps) + return cute.compile( + kernel, + x_fake, + weight_fake, + positions_fake, + freqs_fake, + output_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + +def dspark_mla_ctx_norm_rope( + latent: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + eps: float, +) -> torch.Tensor: + """RMSNorm + RoPE of the context latents of every drafter layer, one launch. + + Args: + latent: ``[N, L, kv_lora_rank + rope_dim]`` BF16 with contiguous + ``(L, head_dim)`` rows (any token stride). + norm_weight: ``[L, kv_lora_rank]`` BF16 latent RMSNorm weights. + positions: ``[N]`` int64 RoPE positions. + freqs: ``[P, rope_dim // 2, 2]`` FP32 (cos, sin) table indexed by position. + eps: RMSNorm epsilon. + Returns: + ``[L, N, head_dim]``: for each layer the values of + ``cute_dsl_dspark_rmsnorm_rope(latent[:, l], norm_weight[l], freqs[positions], + 1, rope_dim, eps, True, True, False, kv_lora_rank)``, bitwise. + """ + num_tokens, num_layers, head_dim = latent.shape + kv_lora_rank = norm_weight.shape[1] + rope_dim = head_dim - kv_lora_rank + assert latent.dtype == torch.bfloat16 and norm_weight.dtype == torch.bfloat16 + assert latent.stride(2) == 1 and latent.stride(1) == head_dim + assert norm_weight.shape[0] == num_layers and norm_weight.is_contiguous() + assert positions.dtype == torch.int64 and positions.shape == (num_tokens,) + assert freqs.dtype == torch.float32 and freqs.is_contiguous() + assert freqs.shape[1:] == (rope_dim // 2, 2) + out = torch.empty((num_layers, num_tokens, head_dim), dtype=latent.dtype, device=latent.device) + if num_tokens == 0: + return out + compiled = _compile_dspark_mla_ctx_norm_rope(num_layers, kv_lora_rank, rope_dim, float(eps)) + compiled( + latent.detach().as_strided((num_tokens, num_layers * head_dim), (latent.stride(0), 1)), + norm_weight.detach(), + positions.contiguous(), + freqs.detach(), + out.view(num_layers * num_tokens, head_dim), + ) + return out + + +@functools.cache +def _compile_dspark_mla_block_rope( + num_heads: int, qk_nope_dim: int, kv_lora_rank: int, rope_dim: int, eps: float +): + head_dim = kv_lora_rank + rope_dim + q_rows = cute.sym_int() + tokens = cute.sym_int() + q_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, (q_rows, qk_nope_dim + rope_dim), stride=(cute.sym_int64(), 1) + ) + query_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, (q_rows, head_dim), stride=(cute.sym_int64(), 1) + ) + kv_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, (tokens, head_dim), stride=(cute.sym_int64(), 1) + ) + weight_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (kv_lora_rank,), stride_order=(0,) + ) + positions_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, (tokens,), stride_order=(0,) + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), rope_dim // 2, 2), stride_order=(2, 1, 0) + ) + pool_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (cute.sym_int(), cute.sym_int(), head_dim), + stride=(cute.sym_int64(), cute.sym_int64(), 1), + ) + pages_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int64, (tokens,), stride_order=(0,)) + offsets_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, (tokens,), stride_order=(0,) + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkMlaBlockRoPEKernel(num_heads, qk_nope_dim, kv_lora_rank, rope_dim, eps) + return cute.compile( + kernel, + q_fake, + query_fake, + kv_fake, + weight_fake, + positions_fake, + freqs_fake, + pool_fake, + pages_fake, + offsets_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + +def dspark_mla_block_rope( + q: torch.Tensor, + query: torch.Tensor, + kv: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + eps: float, +) -> None: + """The block RoPE work of one MLA drafter layer in one launch. + + Args: + q: ``[T, H, qk_nope_dim + rope_dim]`` BF16 query rows (unit inner stride, + ``H`` heads per token). + query: ``[T, H, kv_lora_rank + rope_dim]`` BF16 absorbed query; only its + rope columns are written. + kv: ``[T, kv_lora_rank + rope_dim]`` BF16 block latents (row-strided). + norm_weight: ``[kv_lora_rank]`` BF16 latent RMSNorm weight. + positions: ``[T]`` int64 RoPE positions of the block tokens. + freqs: ``[P, rope_dim // 2, 2]`` FP32 (cos, sin) table indexed by position. + pool: ``[pages, page_size, kv_lora_rank + rope_dim]`` paged latent pool. + pages, offsets: ``[T]`` int64 destination page and in-page offset. + eps: RMSNorm epsilon. + + Writes what ``cute_dsl_dspark_rope_into(q, freqs[positions], query, H, rope_dim, + kv_lora_rank)`` writes into ``query`` and what + ``scatter_rows_to_pool(cute_dsl_dspark_rmsnorm_rope(kv, norm_weight, + freqs[positions], 1, rope_dim, eps, True, True, False, kv_lora_rank), pool, + pages, offsets)`` writes into ``pool``, bitwise. + """ + num_tokens, num_heads, q_dim = q.shape + head_dim = query.shape[-1] + kv_lora_rank = norm_weight.shape[0] + rope_dim = head_dim - kv_lora_rank + assert q.dtype == query.dtype == kv.dtype == pool.dtype == torch.bfloat16 + assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous() + assert query.shape[:2] == (num_tokens, num_heads) and kv.shape == (num_tokens, head_dim) + assert q.stride(2) == 1 and q.stride(0) == num_heads * q.stride(1) + assert query.stride(2) == 1 and query.stride(0) == num_heads * query.stride(1) + assert kv.stride(1) == 1 and pool.stride(2) == 1 and pool.shape[2] == head_dim + assert positions.dtype == pages.dtype == offsets.dtype == torch.int64 + assert positions.numel() == pages.numel() == offsets.numel() == num_tokens + assert freqs.dtype == torch.float32 and freqs.is_contiguous() + assert freqs.shape[1:] == (rope_dim // 2, 2) + if num_tokens == 0: + return + compiled = _compile_dspark_mla_block_rope( + num_heads, q_dim - rope_dim, kv_lora_rank, rope_dim, float(eps) + ) + compiled( + q.detach().as_strided((num_tokens * num_heads, q_dim), (q.stride(1), 1)), + query.detach().as_strided((num_tokens * num_heads, head_dim), (query.stride(1), 1)), + kv.detach(), + norm_weight.detach(), + positions.reshape(-1).contiguous(), + freqs.detach(), + pool.detach(), + pages.reshape(-1).contiguous(), + offsets.reshape(-1).contiguous(), + ) + + +__all__ = [ + "dspark_mla_block_rope", + "dspark_mla_ctx_norm_rope", + "k3r_dspark_rmsnorm_rope", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/_reference.py new file mode 100644 index 000000000000..90710c1ac354 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/_reference.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plain-torch building blocks shared by the references of the speculative-decoding entries.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +#: Index a vocab-parallel split without valid columns reports (above every vocab id). +EMPTY_INDEX = 2**24 +INT32_MAX = 2**31 - 1 + + +def argmax_pairs(vals: torch.Tensor, idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(value, index)`` the ``torch.argmax`` order selects along the last dim of ``(vals, + idx)``: a NaN beats every number and the lowest index wins among NaNs; otherwise the largest + value wins and the lowest index among equal values (``-0.0 == +0.0``).""" + if vals.dtype != torch.float64: + vals = vals.float() + idx = idx.to(torch.int64) + is_nan = torch.isnan(vals) + any_nan = is_nan.any(dim=-1, keepdim=True) + vmax = torch.where(is_nan, float("-inf"), vals).amax(dim=-1, keepdim=True) + cand = torch.where(any_nan, is_nan, (vals == vmax) & ~is_nan) + best = torch.where(cand, idx, INT32_MAX).amin(dim=-1) + value = torch.where(any_nan.squeeze(-1), float("nan"), vmax.squeeze(-1)) + return value, best + + +def interleaved_pairs( + buf: torch.Tensor, stride_row: int, stride_slot: int, n_slots: int, rows: int +) -> tuple[torch.Tensor, torch.Tensor]: + """``(values, indices)`` ``[rows, n_slots]`` of exchanged ``(value, global index)`` float32 + pairs at ``buf.view(-1)[r * stride_row + s * stride_slot + {0, 1}]``.""" + flat = buf.reshape(-1) + r = torch.arange(rows, device=buf.device)[:, None] * stride_row + s = torch.arange(n_slots, device=buf.device)[None, :] * stride_slot + at = (r + s).reshape(-1) + vals = flat[at].view(rows, n_slots) + # The kernels convert the float32 index with a truncating cast; the indices are integers. + idx = flat[at + 1].to(torch.int32).view(rows, n_slots) + return vals, idx + + +def strict_accept( + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """``SpecWorkerBase._sample_and_accept_draft_tokens_base`` on already-selected target + tokens ``[num_contexts + num_gens * (K + 1)]``: ``accepted [batch, K + 1]`` int32 (context + rows define column 0 only; the rest is 0 here) and ``num_accepted [batch]`` int32. + + ``force``: None, ``(base_total,)``, or ``(base_total, frac, max_total, pool, counter, + counter_stride, slot_stride)``, in which case the draw reads ``counter + 1`` (the kernels + advance ``counter`` by one) and ``counter`` is not written here. + """ + num_gens = batch_size - num_contexts + k = draft_tokens.shape[-1] + device = target_tokens.device + accepted = torch.zeros((batch_size, k + 1), dtype=torch.int32, device=device) + num_accepted = torch.ones((batch_size,), dtype=torch.int32, device=device) + accepted[:num_contexts, 0] = target_tokens[:num_contexts].to(torch.int32) + if num_gens == 0: + return accepted, num_accepted + t_gen = target_tokens[num_contexts : num_contexts + num_gens * (k + 1)].to(torch.int32) + t_gen = t_gen.view(num_gens, k + 1) + accepted[num_contexts:] = t_gen + match = (draft_tokens[:num_gens].to(torch.int32) == t_gen[:, :k]).to(torch.int32) + n = 1 + match.cumprod(dim=1).sum(dim=1) + if force is not None and len(force) == 1: + n = torch.full_like(n, int(force[0])) + elif force is not None: + base_total, frac, max_total, pool, counter, counter_stride, slot_stride = force + nxt = counter.reshape(-1)[0].to(torch.int64) + 1 + rows = torch.arange(num_gens, device=device, dtype=torch.int64) + at = (nxt * counter_stride + rows * slot_stride) & (pool.numel() - 1) + extra = (pool.reshape(-1)[at] < frac).to(torch.int64) + n = (base_total + extra).clamp(max=max_total) + num_accepted[num_contexts:] = n.to(torch.int32) + return accepted, num_accepted + + +def u32_mul(a: torch.Tensor, b: int) -> tuple[torch.Tensor, torch.Tensor]: + """``(high, low)`` 32-bit words of ``a * b`` for int64 tensors holding uint32 values.""" + b_hi, b_lo = b >> 16, b & 0xFFFF + p_lo = a * b_lo + p_hi = a * b_hi + high = (p_hi + (p_lo >> 16)) >> 16 + low = (p_lo + ((p_hi & 0xFFFF) << 16)) & 0xFFFFFFFF + return high, low + + +def philox4x32(seed: torch.Tensor, c0, c1, c2, c3, rounds: int = 10): + """Philox4x32 (``tl.philox``) on int64 tensors holding uint32 counters; ``seed`` int64.""" + mask = 0xFFFFFFFF + k0 = seed & mask + k1 = (seed >> 32) & mask + for _ in range(rounds): + hi_b, lo_b = u32_mul(c2, 0xCD9E8D57) + hi_a, lo_a = u32_mul(c0, 0xD2511F53) + c0, c1, c2, c3 = hi_b ^ c1 ^ k0, lo_b, hi_a ^ c3 ^ k1, lo_a + k0 = (k0 + 0x9E3779B9) & mask + k1 = (k1 + 0xBB67AE85) & mask + return c0, c1, c2, c3 + + +def gumbel(bits: torch.Tensor) -> torch.Tensor: + """``-log(-log(u))`` with ``u = ((bits >> 9) + 0.5) / 2**23``, in float64.""" + u = ((bits >> 9).double() + 0.5) / 8388608.0 + return -torch.log(-torch.log(u)) + + +def publish_pairs(layout, vals: torch.Tensor, col_base: int) -> torch.Tensor: + """The pair buffer of ``layout`` (a ``vocab_parallel_kernels._PairLayout``) after each + split of ``vals`` ``[rows, shard]`` published its argmax ``(value, col_base + column)`` at + slot ``slot0 + split``, and the columns outside ``[zero_lo, zero_hi)`` were zeroed when + ``zero_width > 0``. A clone; ``layout.buffer`` is not written.""" + buf = layout.buffer.clone() + rows, n_cols = vals.shape + if rows == 0: + return buf + if layout.zero_width > 0: + cols = torch.arange(layout.zero_width, device=buf.device) + outside = (cols < layout.zero_lo) | (cols >= layout.zero_hi) + buf[:, : layout.zero_width][:, outside] = 0.0 + for split in range(layout.splits): + lo = split * layout.chunk + hi = min(lo + layout.chunk, n_cols) + slot = layout.slot0 + split + if lo >= hi: + value = torch.full((rows,), float("-inf"), device=buf.device) + index = torch.full((rows,), EMPTY_INDEX, device=buf.device) + else: + cols = torch.arange(lo, hi, device=buf.device).expand(rows, -1) + value, index = argmax_pairs(vals[:, lo:hi], cols) + index = index + col_base + buf[:, 2 * slot] = value.to(buf.dtype) + buf[:, 2 * slot + 1] = index.to(buf.dtype) + return buf + + +def markov_chain_tokens( + block_logits: torch.Tensor, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + vocab: int, +) -> torch.Tensor: + """Greedy DSpark vanilla-Markov chain: ``[B, K]`` int32 tokens of the block logits + ``[B, K, V]`` (``V`` the columns of ``markov_w2`` rows given). Step ``i`` adds the bias + ``markov_w2 @ markov_w1[prev_i]`` (FP32 product rounded to the weights' dtype, as + ``F.linear`` returns it) to position ``i`` in FP32 and takes the ``torch.argmax``; ``prev_0`` + is the anchor, and a ``prev`` outside ``[0, vocab)`` gives a zero bias.""" + num_reqs, block_size, width = block_logits.shape + device = block_logits.device + tokens = torch.empty((num_reqs, block_size), dtype=torch.int32, device=device) + prev = anchors.to(torch.int64) + cols = torch.arange(width, device=device).expand(num_reqs, -1) + w2 = markov_w2.double() + for step in range(block_size): + valid = (prev >= 0) & (prev < vocab) + row = markov_w1[prev.clamp(0, markov_w1.shape[0] - 1)].double() * valid[:, None] + bias = (row @ w2.t()).to(markov_w2.dtype).float() + _, prev = argmax_pairs(block_logits[:, step].float() + bias, cols) + tokens[:, step] = prev.to(torch.int32) + return tokens diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_ctx_positions.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_ctx_positions.py new file mode 100644 index 000000000000..2d290c93778f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_ctx_positions.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RoPE positions of a DFlash generation step's context latents (Triton, one program).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import dflash_ctx_positions + + +class K3DflashCtxPositions(OpWrapper): + """``out[g * width + j] = ctx_len[slots[g]] + j`` for ``g < len(slots)``, ``j < width``: + int64 ``[len(slots) * width]``, new. ``ctx_len`` is the drafter's per-slot context length + before the step advances it.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="helix8 c24/c42: 2 requests of 8 positions on the side stream", + spec=dict(num_gens=2, width=8, max_batch=2), + ), + Cell( + why="helix8 c75/c174: 8 requests of 8 positions", + spec=dict(num_gens=8, width=8, max_batch=8), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4): 4 requests of 8 positions, padding rows on the " + "dummy slot", + spec=dict(num_gens=4, width=8, max_batch=4), + ), + ) + + note = """ + Runs on the side stream that computes a TP-sharded drafter's context latents, before the step + preparation advances ``ctx_len`` on the main stream (the caller orders the two with an event). + """ + + def __call__(self, ctx_len: torch.Tensor, slots: torch.Tensor, width: int) -> torch.Tensor: + return dflash_ctx_positions(ctx_len, slots, width) + + def reference(self, ctx_len: torch.Tensor, slots: torch.Tensor, width: int) -> torch.Tensor: + old = ctx_len[slots.to(torch.int64)].to(torch.int64) + j = torch.arange(width, device=ctx_len.device, dtype=torch.int64) + return (old[:, None] + j[None, :]).reshape(-1) + + +k3_dflash_ctx_positions = K3DflashCtxPositions() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_gen_prepare.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_gen_prepare.py new file mode 100644 index 000000000000..41c077676c92 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_gen_prepare.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DFlash paged-context bookkeeping of one generation step, Helix-sharded or not, in one +Triton launch (one program).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import dflash_gen_prepare + + +def gen_prepare_reference( + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + first_row: int, + dummy_slot: int, + max_ctx: int, + page_size: int, + block_size: int, + width: int, + cp_size: int = 1, + cp_rank: int = 0, +) -> tuple[dict, torch.Tensor, torch.Tensor]: + """``(outputs, ctx_len after, seq_lens_cuda after)`` of the step preparation for the + ``len(slots)`` generation requests at batch rows ``first_row + g``; nothing is written.""" + device = ctx_len.device + num_gens = slots.shape[0] + g_rows = first_row + torch.arange(num_gens, device=device) + nacc = num_accepted_tokens[g_rows].to(torch.int64) + slot = slots.to(torch.int64) + old = ctx_len[slot].to(torch.int64) + counts = block_counts[g_rows].to(torch.int64) + table = block_tables[g_rows].to(torch.int64) + max_blocks = block_tables.shape[1] + + def page_of(local_pos): + return table.gather(1, (local_pos // page_size).clamp(max=max_blocks - 1)) + + bonus = accepted_tokens[g_rows].gather(1, (nacc - 1).clamp(min=0)[:, None]).squeeze(1) + bonus = bonus.to(torch.int64) + now = (old + nacc).clamp(max=max_ctx) + j_block = torch.arange(block_size, device=device) + j_width = torch.arange(width, device=device) + pos = old[:, None] + j_width[None, :] + capacity = counts * page_size + gpage = pos // page_size + col = (gpage // cp_size) * page_size + pos % page_size + col = torch.minimum(col, capacity[:, None] - 1).clamp(min=0) + store = (j_width[None, :] < nacc[:, None]).to(torch.int8) + store = torch.where(gpage % cp_size == cp_rank, store, torch.full_like(store, -1)) + new_len = torch.where(slot == dummy_slot, torch.zeros_like(now), now) + num_ctx = torch.minimum(new_len, (capacity * cp_size - block_size).clamp(min=0)) + bpos = num_ctx[:, None] + j_block[None, :] + bgpage = bpos // page_size + bslot = (bgpage // cp_size) * page_size + bpos % page_size + total = num_ctx + block_size + out = dict( + bonus=bonus, + query_positions=now[:, None] + j_block[None, :], + ctx_positions=pos.reshape(-1), + store_pages=page_of(col).reshape(-1), + store_offsets=(col % page_size).reshape(-1), + store_mask=store.reshape(-1), + num_ctx=num_ctx, + blk_pages=page_of(bslot), + blk_offsets=bslot % page_size, + ) + if cp_size > 1: + out["blk_owned"] = (bgpage % cp_size == cp_rank).to(torch.int8) + span = page_size * cp_size + rem = total % span - cp_rank * page_size + total = (total // span) * page_size + rem.clamp(min=0, max=page_size) + out["blk_seq_lens"] = total.to(torch.int32) + out["blk_kv_bounds"] = total.to(torch.int32).repeat_interleave(block_size) + ctx_out = ctx_len.clone() + ctx_out[slot] = new_len.to(ctx_len.dtype) + ctx_out[dummy_slot] = 0 + seq_out = seq_lens_cuda.clone() + seq_out[g_rows] = width + return out, ctx_out, seq_out + + +class K3DflashGenPrepare(OpWrapper): + """The per-request bookkeeping of a DFlash generation step whose drafter context is the + paged single-latent cache, as a dict of new tensors, and its two side effects: ``ctx_len`` + advanced in place (the dummy slot zeroed) and ``seq_lens_cuda`` of the generation rows set + to ``width``. + + Generation request ``g`` is batch row ``num_contexts + g`` and context slot ``slots[g]``. The + dict holds the bonus token, ``embed_ids`` (the bonus tokens then ``mask_token_id``), the block + query positions, the RoPE positions / destination page / offset / write mask of the ``width`` + target latents, the advertised context length and the block-decode page / offset / bound + operands. With ``cp_size > 1`` the pages are Helix-sharded: the store mask is -1 on other + ranks' positions, the block operands are rank-local and ``blk_owned`` marks this rank's. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="dep16 (c383): 16 requests, DSpark 5 (width and block 6), one rank's pages", + spec=dict(num_gens=16, draft_len=5, cp_size=1, cp_rank=0, dummies=0), + ), + Cell( + why="helix8 c240: 20 requests of DSpark 5, pages over 8 ranks (rank 3)", + spec=dict(num_gens=20, draft_len=5, cp_size=8, cp_rank=3, dummies=0), + ), + Cell( + why="helix8 c75/c174: 8 requests of DSpark 7 on the last of 8 ranks", + spec=dict(num_gens=8, draft_len=7, cp_size=8, cp_rank=7, dummies=0), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4): 4 requests of DSpark 7 over 16 ranks (rank 5), two " + "padding rows on the dummy slot", + spec=dict(num_gens=4, draft_len=7, cp_size=16, cp_rank=5, dummies=2), + ), + ) + + note = """ + ``K3DFlashWorker.prepare_1st_drafter_inputs`` and ``_build_mla_block_fixup`` for that cache, in + one program of 4 warps; every lane reads its old context length before any lane writes one. + The advertised context is capped so the block fits the allocation. ``accepted_tokens`` and + ``block_tables`` need unit inner strides. + """ + + def __call__( + self, + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + num_contexts: int, + num_gens: int, + dummy_slot: int, + max_ctx: int, + mask_token_id: int, + page_size: int, + block_size: int, + width: int, + cp_size: int = 1, + cp_rank: int = 0, + ) -> dict: + return dflash_gen_prepare( + num_accepted_tokens, + accepted_tokens, + slots, + ctx_len, + block_counts, + block_tables, + seq_lens_cuda, + num_contexts, + num_gens, + dummy_slot, + max_ctx, + mask_token_id, + page_size, + block_size, + width, + cp_size=cp_size, + cp_rank=cp_rank, + ) + + def reference( + self, + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + num_contexts: int, + num_gens: int, + dummy_slot: int, + max_ctx: int, + mask_token_id: int, + page_size: int, + block_size: int, + width: int, + cp_size: int = 1, + cp_rank: int = 0, + ) -> tuple[dict, torch.Tensor, torch.Tensor]: + """``(the op's dict, ctx_len after, seq_lens_cuda after)``; nothing is written.""" + out, ctx_out, seq_out = gen_prepare_reference( + num_accepted_tokens, + accepted_tokens, + slots[:num_gens], + ctx_len, + block_counts, + block_tables, + seq_lens_cuda, + num_contexts, + dummy_slot, + max_ctx, + page_size, + block_size, + width, + cp_size, + cp_rank, + ) + mask = torch.full((1,), int(mask_token_id), dtype=torch.int64, device=ctx_len.device) + out["embed_ids"] = torch.cat((out["bonus"], mask)) + return out, ctx_out, seq_out + + def is_valid( + self, + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + num_contexts: int, + num_gens: int, + dummy_slot: int, + max_ctx: int, + mask_token_id: int, + page_size: int, + block_size: int, + width: int, + cp_size: int = 1, + cp_rank: int = 0, + ) -> None: + # ctx_len is stored through an int64 pointer; one program covers every request. + assert ctx_len.dtype == torch.int64, "ctx_len must be int64" + assert slots.shape[0] >= num_gens, "one slot per generation request" + assert 0 <= cp_rank < cp_size, "cp_rank out of range" + live = slots[:num_gens][slots[:num_gens] != dummy_slot] + # Two live requests on one slot race on the advanced length. + assert live.unique().numel() == live.numel(), "generation requests share a slot" + + +k3_dflash_gen_prepare = K3DflashGenPrepare() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_noise_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_noise_embedding.py new file mode 100644 index 000000000000..75a48f64331d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_noise_embedding.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DFlash drafter's input block embedding from one combined lookup (Triton copy).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import dflash_noise_embedding + + +class K3DflashNoiseEmbedding(OpWrapper): + """``[G, block_size, H]`` (new, the dtype of ``combined_embed``): row ``combined_embed[g]`` + (the bonus token's embedding) at block slot 0 and ``combined_embed[G]`` (the mask token's) + at every other slot, from ``combined_embed = embed([bonus_0 .. bonus_{G-1}, mask])``.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="dep16 (c383): 16 requests, blocks of 6, hidden 7168", + spec=dict(num_gens=16, block_size=6), + ), + Cell( + why="helix8 c75/c174: 8 requests, blocks of 8", + spec=dict(num_gens=8, block_size=8), + ), + Cell( + why="helix16_moetp4ep4 (ULL c1): one request, a block of 8", + spec=dict(num_gens=1, block_size=8), + ), + ) + + note = """ + A pure copy: grid (G, block_size, cdiv(H, 1024)). A non-contiguous ``combined_embed`` is made + contiguous first. + """ + + def __call__(self, combined_embed: torch.Tensor, block_size: int) -> torch.Tensor: + return dflash_noise_embedding(combined_embed, block_size) + + def reference(self, combined_embed: torch.Tensor, block_size: int) -> torch.Tensor: + num_gens = combined_embed.shape[0] - 1 + out = combined_embed[num_gens:].expand(num_gens, -1).unsqueeze(1) + out = out.repeat(1, block_size, 1) + out[:, 0] = combined_embed[:num_gens] + return out + + +k3_dflash_noise_embedding = K3DflashNoiseEmbedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_paged_latent_prepare.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_paged_latent_prepare.py new file mode 100644 index 000000000000..f286baf9e13f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dflash_paged_latent_prepare.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused DFlash generation-step bookkeeping of a paged single-latent context (Triton, +one program): the preparation a generation profile without the tail fusion takes.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import PagedLatentPrepare, dflash_paged_latent_prepare +from .k3_dflash_gen_prepare import gen_prepare_reference + + +class K3DflashPagedLatentPrepare(OpWrapper): + """``PagedLatentPrepare`` of new tensors for the ``len(slots)`` generation requests at batch + rows ``first_row + g``, plus ``ctx_len`` advanced in place (the dummy slot zeroed) and + ``seq_lens_cuda`` of those rows set to ``width``: ``k3_dflash_gen_prepare`` with one rank + and without the embedding ids.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="dep16 (c383) shape: 16 requests of DSpark 5 on one rank (the unfused preparation " + "no shipped profile takes)", + spec=dict(num_gens=16, draft_len=5, first_row=0, dummies=1), + ), + Cell( + why="the same unfused preparation with 2 context rows ahead of 2 DSpark 7 requests: " + "the generation rows start at batch row 2", + spec=dict(num_gens=2, draft_len=7, first_row=2, dummies=0), + ), + ) + + note = """ + ``K3DFlashWorker.prepare_1st_drafter_inputs`` and ``_build_mla_block_fixup`` in one program; + every lane reads its old context length before any lane writes one. No shipped K3 generation + profile reaches it (each runs with the tail fusion, which takes ``k3_dflash_gen_prepare``). + """ + + def __call__( + self, + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + first_row: int, + dummy_slot: int, + max_ctx: int, + page_size: int, + block_size: int, + width: int, + ) -> PagedLatentPrepare: + return dflash_paged_latent_prepare( + num_accepted_tokens, + accepted_tokens, + slots, + ctx_len, + block_counts, + block_tables, + seq_lens_cuda, + first_row, + dummy_slot, + max_ctx, + page_size, + block_size, + width, + ) + + def reference( + self, + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + first_row: int, + dummy_slot: int, + max_ctx: int, + page_size: int, + block_size: int, + width: int, + ) -> tuple[PagedLatentPrepare, torch.Tensor, torch.Tensor]: + """``(the op's PagedLatentPrepare, ctx_len after, seq_lens_cuda after)``; nothing is + written.""" + out, ctx_out, seq_out = gen_prepare_reference( + num_accepted_tokens, + accepted_tokens, + slots, + ctx_len, + block_counts, + block_tables, + seq_lens_cuda, + first_row, + dummy_slot, + max_ctx, + page_size, + block_size, + width, + ) + prep = PagedLatentPrepare(**{name: out[name] for name in PagedLatentPrepare._fields}) + if slots.shape[0] == 0: + # Nothing is launched without a generation request: the dummy slot keeps its length. + ctx_out = ctx_len.clone() + return prep, ctx_out, seq_out + + def is_valid( + self, + num_accepted_tokens: torch.Tensor, + accepted_tokens: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + block_counts: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_cuda: torch.Tensor, + first_row: int, + dummy_slot: int, + max_ctx: int, + page_size: int, + block_size: int, + width: int, + ) -> None: + assert ctx_len.dtype == torch.int64, "ctx_len must be int64" + live = slots[slots != dummy_slot] + assert live.unique().numel() == live.numel(), "generation requests share a slot" + + +k3_dflash_paged_latent_prepare = K3DflashPagedLatentPrepare() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_ar_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_ar_norm.py new file mode 100644 index 000000000000..3d99390b331c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_ar_norm.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark drafter's all-reduce site: finish a pushed MNNVL all-reduce, add the residual and +apply the following RMSNorm in one CuTe DSL launch.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper + +from ...kernels.dspark_draft_ar import ar_norm_config, run_draft_ar_norm + + +class K3DraftArNorm(OpWrapper): + """``residual_out = bf16(bf16(sum over ranks of the pushed partials) + residual)`` and + ``out = RMSNorm(residual_out) * gamma`` (BF16), written into the two contiguous ``[M, + hidden]`` outputs. Returns None. + + Every rank pushed its BF16 partial into ``workspace`` (a ``k3_lamport_ar.LamportWorkspace``) + with the ``one_shot`` or two-shot transport; ``cpp_one_shot`` selects which C++ fused + all-reduce (``RESIDUAL_RMS_NORM``) the outputs reproduce bit for bit: its one-shot kernel + (``bf16((x * r) * gamma)``) or two-shot + ``rmsNormLamport`` (``bf16((gamma * x) * r)``). + """ + + ARCHS = frozenset({Arch.SM_107}) + + note = """ + The ranks are summed in FP32 in rank order; the RMS reduction sums BF16 squares in a fixed + order and uses ``rsqrt.approx``. Launched behind the producer with PDL: it releases its + dependents first and polls the Lamport slots while the producer runs; only CTA (0, 0) waits + for the producer before advancing the workspace's rotation state. Up to 64 rows with a + single-load launch shape (``ar_norm_config``); raises ValueError otherwise. A process-wide + compile cache holds one kernel per (world, hidden, threads, cluster, transport, C++ path, + PDL). Every rank must make the same call. + """ + + def __call__( + self, + workspace, + one_shot: bool, + cpp_one_shot: bool, + residual: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + residual_out: torch.Tensor, + use_pdl: bool = True, + ) -> None: + run_draft_ar_norm( + workspace, + one_shot, + cpp_one_shot, + residual, + gamma, + eps, + out, + residual_out, + use_pdl=use_pdl, + ) + + def reference( + self, + partials: Sequence[torch.Tensor], + one_shot: bool, + cpp_one_shot: bool, + residual: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + residual_out: torch.Tensor, + use_pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``(out, residual_out)``, BF16. A collective: ``partials`` (in place of ``workspace``) + is the group's pushed BF16 partials in rank order. The sum is FP32 in rank order, rounded + to BF16 and added to the residual with one rounding; the norm is FP64 from that BF16 + residual, rounded once. The transport flags and the output buffers are not read.""" + acc = torch.zeros(residual.shape, dtype=torch.float32, device=residual.device) + for partial in partials: + acc += partial.float() + new_residual = (acc.to(torch.bfloat16).float() + residual.float()).to(torch.bfloat16) + x = new_residual.double() + inv_rms = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) + norm = (x * inv_rms * gamma.double()).to(torch.bfloat16) + return norm, new_residual + + def is_valid( + self, + workspace, + one_shot: bool, + cpp_one_shot: bool, + residual: torch.Tensor, + gamma: torch.Tensor, + eps: float, + out: torch.Tensor, + residual_out: torch.Tensor, + use_pdl: bool = True, + ) -> None: + # Raw addresses: shapes, dtypes and layouts are not checked by the launcher. + for t in (residual, out, residual_out): + assert t.dtype == torch.bfloat16 and t.is_contiguous() and t.shape == out.shape + assert gamma.dtype == torch.bfloat16 and gamma.numel() == out.shape[1] + assert ar_norm_config(out.shape[0], out.shape[1], out.device)[0], "no launch shape" + + +k3_draft_ar_norm = K3DraftArNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_block_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_block_table.py new file mode 100644 index 000000000000..5ba49284d449 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_block_table.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DFlash drafter's page-table decode: layer-relative block offsets to pool pages, and the +pages each row holds (Triton, one row per program or a row's columns over programs).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import ( + decode_draft_block_table, + decode_draft_block_table_parallel, +) + + +class K3DraftBlockTable(OpWrapper): + """For ``encoded = src - block_offset``: ``table[:rows] = encoded.clamp(min=0) // divisor`` + (int32) and ``counts[:rows] = (encoded >= 0).sum(-1)`` (int64). Returns None. + + ``src`` is ``[rows, blocks]`` of any integer dtype and strides and is not written; ``table`` is + ``[>= rows, blocks]`` int32 with a unit inner stride, ``counts`` ``[>= rows]`` int64. With + ``partial`` (an int32 scratch of at least ``rows * cdiv(blocks, 2048)`` elements) a row's + columns are spread over programs and a second launch sums the counts; the values are the same. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="tp1 and dep16 (c383): 16 rows of the 7795-page table (max_seq_len 997696, 128 " + "tokens per page), one program per row", + spec=dict(rows=16, blocks=7795, partial=False, divisor=5, block_offset=0), + ), + Cell( + why="helix8 c24/c42: 2 rows, each row's columns over 4 programs, then the count sum", + spec=dict(rows=2, blocks=7795, partial=True, divisor=1, block_offset=3), + ), + Cell( + why="helix16_moetp4ep4 (ULL c1/c2/c4): 4 rows, the parallel decode", + spec=dict(rows=4, blocks=7795, partial=True, divisor=1, block_offset=3), + ), + Cell( + why="helix8 c240: 20 rows, the parallel decode at the largest helix batch", + spec=dict(rows=20, blocks=7795, partial=True, divisor=1, block_offset=3), + ), + ) + + note = """ + Integer arithmetic only; rows past ``rows`` of ``table`` and ``counts`` are not touched. The + caller passes the metadata's ``kv_cache_block_offsets`` plane of the drafter pool and the + worker's persistent table, counts and scratch. Launch shapes depend on host sizes only. + """ + + def __call__( + self, + src: torch.Tensor, + divisor: int, + table: torch.Tensor, + counts: torch.Tensor, + partial: Optional[torch.Tensor] = None, + block_offset: int = 0, + ) -> None: + if partial is None: + decode_draft_block_table(src, divisor, table, counts, block_offset=block_offset) + else: + decode_draft_block_table_parallel( + src, divisor, table, counts, partial, block_offset=block_offset + ) + + def reference( + self, + src: torch.Tensor, + divisor: int, + table: torch.Tensor, + counts: torch.Tensor, + partial: Optional[torch.Tensor] = None, + block_offset: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``(table, counts)`` after the call; the arguments are not written, ``partial`` is + scratch and not read.""" + rows = src.shape[0] + encoded = src.to(torch.int64) - int(block_offset) + table_out, counts_out = table.clone(), counts.clone() + table_out[:rows] = torch.div(encoded.clamp(min=0), int(divisor), rounding_mode="floor").to( + torch.int32 + ) + counts_out[:rows] = (encoded >= 0).sum(dim=-1).to(counts.dtype) + return table_out, counts_out + + def is_valid( + self, + src: torch.Tensor, + divisor: int, + table: torch.Tensor, + counts: torch.Tensor, + partial: Optional[torch.Tensor] = None, + block_offset: int = 0, + ) -> None: + # The kernels index `table` with its row stride only and store int32 / int64 through raw + # pointers: another dtype is written with the wrong width. + assert table.dtype == torch.int32, "table must be int32" + assert counts.dtype == torch.int64 and counts.stride(-1) == 1, "counts must be int64" + assert int(divisor) > 0, "divisor must be positive" + + +k3_draft_block_table = K3DraftBlockTable() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_embed.py new file mode 100644 index 000000000000..eea751c36e5c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_embed.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark drafter's vocab-parallel embedding as one owner-multicast CuTe DSL launch: the +owner rank of each token stores its row into every rank's one-shot all-reduce buffer.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper + +from ...kernels.dspark_draft_ar import ELTS, EMBED_THREADS, run_draft_embed + + +class K3DraftEmbed(OpWrapper): + """``out[t]`` (BF16 ``[M, hidden]``, contiguous) = the full-vocab row of ``ids[t]``, zero for + an id outside ``[0, vocab)``: the all-reduce of the vocab-sharded lookup. This rank holds + rows ``[rank * slice_width, ...)`` of the table in ``table`` (BF16, contiguous); ``ids`` is + contiguous int32 / int64 ``[M]``. Returns None. + """ + + ARCHS = frozenset({Arch.SM_107}) + + note = """ + Through the one-shot Lamport buffer of ``workspace`` (the drafter embedding all-reduce's + MNNVL workspace): the owner stores its row, ``-0`` sent as ``+0``, into slot (token, owner) of + every rank and every rank polls that one slot, so the result is the all-reduce's, bit for + bit, in one launch and one hop. ``hidden`` must be a multiple of ELTS * EMBED_THREADS. A + process-wide compile cache holds one kernel per (world, hidden, id dtype, PDL). Every rank + must make the same call. + """ + + def __call__( + self, + workspace, + ids: torch.Tensor, + table: torch.Tensor, + slice_width: int, + vocab: int, + out: torch.Tensor, + use_pdl: bool = True, + ) -> None: + run_draft_embed(workspace, ids, table, slice_width, vocab, out, use_pdl=use_pdl) + + def reference( + self, + workspace, + ids: torch.Tensor, + tables: Sequence[torch.Tensor], + slice_width: int, + vocab: int, + out: torch.Tensor, + use_pdl: bool = True, + ) -> torch.Tensor: + """``out`` after the call (new). A collective: ``tables`` (in place of ``table``) is the + group's table slices in rank order; ``workspace`` and ``out`` are not read. Every row is + a copy, ``-0`` becoming ``+0`` as in the sum.""" + full = torch.cat([t[:slice_width] for t in tables])[:vocab] + ids = ids.to(torch.int64) + valid = (ids >= 0) & (ids < vocab) + rows = full[ids.clamp(0, vocab - 1)] + rows = torch.where(valid[:, None], rows, torch.zeros_like(rows)) + return rows + 0.0 + + def is_valid( + self, + workspace, + ids: torch.Tensor, + table: torch.Tensor, + slice_width: int, + vocab: int, + out: torch.Tensor, + use_pdl: bool = True, + ) -> None: + assert table.dtype == out.dtype == torch.bfloat16 and table.is_contiguous() + assert out.is_contiguous() and ids.is_contiguous() and ids.dim() == 1 + assert ids.dtype in (torch.int32, torch.int64) and ids.shape[0] == out.shape[0] + assert out.shape[1] % (ELTS * EMBED_THREADS) == 0, "hidden not a whole number of chunks" + assert table.shape[0] >= slice_width or vocab <= slice_width * (workspace.rank + 1) + + +k3_draft_embed = K3DraftEmbed() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_gemm_push.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_gemm_push.py new file mode 100644 index 000000000000..f0debf81a671 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_draft_gemm_push.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark drafter's low-M BF16 split-K GEMM whose epilogue pushes its output into the Lamport +slots of the MNNVL all-reduce that follows (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper + +from ...kernels.dspark_draft_gemm_push import MAX_M, run_gemm_push, validate_push_tactic +from ...kernels.low_m_bf16_splitk import SplitKTactic + + +class K3DraftGemmPush(OpWrapper): + """Pushes ``x [M, K] @ weight[N, K]^T`` (BF16, FP32 accumulation over split-K slices) into + ``target`` (a ``layers.LamportPushTarget`` of an all-reduce of ``[M, N]``) with the + target's transport and sets ``target.pushed``; with ``store_out`` also writes it into ``out`` + (``[M, N]``, contiguous, whose shape sets the launch either way). Returns None. + + ``x`` and ``weight`` are contiguous BF16, 32-byte aligned; ``tactic`` a ``SplitKTactic`` that + ``validate_push_tactic`` admits for ``(M, N, K)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + note = """ + The owner CTA of each cluster reduces the split-K partials, stages the tile token-major and + stores it as 16-byte vectors into every rank's slot (one-shot, through the multicast address) + or the token owner's scatter slot (two-shot), ``-0`` as ``+0``; ``k3_draft_ar_norm`` or the + Lamport finish then completes the all-reduce. The weight warp streams before + ``griddepcontrol.wait``; dependents are released right after the activation warp's wait. Up + to MAX_M (64) rows. A process-wide compile cache holds one kernel per (dtype, tactic, PDL, + world, N, transport, store_out). Targets are finished in push order by every rank. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + target, + tactic: SplitKTactic, + pdl: bool = True, + store_out: bool = False, + ) -> None: + run_gemm_push(x, weight, out, target, tactic, pdl=pdl, store_out=store_out) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + target, + tactic: SplitKTactic, + pdl: bool = True, + store_out: bool = False, + ) -> torch.Tensor: + """The BF16 ``[M, N]`` product this rank pushes (and stores into ``out`` with + ``store_out``): an FP32 product rounded once. Nothing is written.""" + return (x.float() @ weight.float().t()).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + target, + tactic: SplitKTactic, + pdl: bool = True, + store_out: bool = False, + ) -> None: + m, k = x.shape + n = weight.shape[0] + assert x.dtype == weight.dtype == out.dtype == torch.bfloat16 + assert x.is_contiguous() and weight.is_contiguous() and out.is_contiguous() + assert tuple(out.shape) == (m, n) and m <= MAX_M + validate_push_tactic(tactic, m, n, k) + # The push lays the tile out for the target's [M, N] all-reduce. + assert tuple(target.output.shape) == (m, n), "the target is for another shape" + + +k3_draft_gemm_push = K3DraftGemmPush() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_markov_greedy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_markov_greedy.py new file mode 100644 index 000000000000..44fc1507c7a1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_markov_greedy.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark greedy Markov chain over a vocab-sharded LM head in one CuTe DSL launch +(SM100 / SM103 / SM107), per-CTA keys exchanged through multicast stores.""" + +from __future__ import annotations + +from typing import Optional, Sequence, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.dspark_markov_greedy import run_dspark_markov_greedy +from ._reference import markov_chain_tokens + +Shards = Union[torch.Tensor, Sequence[torch.Tensor]] + + +class K3DsparkMarkovGreedy(OpWrapper): + """Writes ``tokens`` (int32 ``[B, K]``): the chain of ``k3_markov_chain`` on BF16 block + logits ``[B, K, shard]`` (unit vocab stride), ``anchors`` ``[B]`` int32 / int64, the full + BF16 ``markov_w1`` ``[vocab, R]`` and this rank's ``w2_shard`` ``[shard, R]`` (contiguous). + Returns None. + + ``state`` is the rank's int32 ``[STATE_WORDS]`` launch state (zero-initialized), ``box_uc`` / + ``box_mc`` the local and multicast addresses of the ``NUM_SETS`` key-exchange sets + (zero-initialized; equal on one rank, ``multicast`` False). + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # One rank of the helix group with the others' keys in the exchange (no multicast): + # 163840 columns over 16 ranks (40 CTAs) or 8 ranks (80 CTAs), Markov rank 256, 8 block + # positions (DSpark 7); the request widths 1, 2, 4 and 8 of the recipes. + CELLS = ( + Cell( + why="helix16 ULL c1, rank 0 of 16: one request, int32 anchor", + spec=dict(tp=16, rank=0, reqs=1, block=8, anchor="int32"), + ), + Cell( + why="helix16 ULL c4, last rank: 4 requests, int64 anchors, one outside the vocab", + spec=dict(tp=16, rank=15, reqs=4, block=8, anchor="int64"), + ), + Cell( + why="helix8 c24/c42, rank 3 of 8: 2 requests over a 20480-column shard", + spec=dict(tp=8, rank=3, reqs=2, block=8, anchor="int32"), + ), + Cell( + why="helix8 c75/c174, last rank: 8 requests, the 5120 keys of a position", + spec=dict(tp=8, rank=7, reqs=8, block=8, anchor="int64"), + ), + ) + + note = """ + One CTA of 512 threads per 256 shard rows, all resident (the caller checks SM count and + shared memory: R a multiple of 16 up to 256, at most 16 requests and positions). The bias is + ``mma.sync`` BF16 x BF16 accumulated in FP32 over R in order and rounded to BF16, then added to + the BF16 logit in FP32; keys order by value, then lower global index, NaN first. A + process-wide compile cache holds one kernel per (request width, positions, tp, CTAs, R, + multicast, anchor dtype, PDL, staged, gate); compile before capture. The launch ordinal in + ``state`` rotates three exchange sets; every rank makes the same sequence of calls. + """ + + def __call__( + self, + local_block: torch.Tensor, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: torch.Tensor, + tokens: torch.Tensor, + *, + vocab: int, + tp: int, + rank: int, + state: torch.Tensor, + box_uc: int, + box_mc: int, + multicast: bool, + use_pdl: bool = True, + staged: Optional[bool] = None, + gate: Optional[bool] = None, + ) -> None: + run_dspark_markov_greedy( + local_block, + anchors, + markov_w1, + w2_shard, + tokens, + vocab=vocab, + tp=tp, + rank=rank, + state=state, + box_uc=box_uc, + box_mc=box_mc, + multicast=multicast, + use_pdl=use_pdl, + staged=staged, + gate=gate, + ) + + def reference( + self, + local_block: Shards, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: Shards, + tokens: torch.Tensor, + *, + vocab: int, + tp: int, + rank: int, + state: torch.Tensor, + box_uc: int, + box_mc: int, + multicast: bool, + use_pdl: bool = True, + staged: Optional[bool] = None, + gate: Optional[bool] = None, + ) -> torch.Tensor: + """The tokens (new). A collective: with ``tp > 1`` ``local_block`` and ``w2_shard`` are + the group's, in rank order. The launch state and the exchange are not read.""" + if not isinstance(local_block, torch.Tensor): + local_block = torch.cat(list(local_block), dim=-1) + w2_shard = torch.cat(list(w2_shard), dim=0) + return markov_chain_tokens(local_block, anchors, markov_w1, w2_shard, vocab) + + def is_valid( + self, + local_block: torch.Tensor, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: torch.Tensor, + tokens: torch.Tensor, + **kwargs, + ) -> None: + num_reqs, steps, shard = local_block.shape + assert local_block.dtype == markov_w1.dtype == w2_shard.dtype == torch.bfloat16 + assert 1 <= num_reqs <= 16 and 1 <= steps <= 16, "1..16 requests and positions" + assert w2_shard.shape == (shard, markov_w1.shape[1]) and w2_shard.is_contiguous() + assert markov_w1.shape[1] % 16 == 0 and markov_w1.shape[1] <= 256, "R: 16 x n <= 256" + assert local_block.stride(-1) == 1 and markov_w1.is_contiguous() + assert tokens.dtype == torch.int32 and tokens.stride(-1) == 1 + + +k3_dspark_markov_greedy = K3DsparkMarkovGreedy() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_mla_block_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_mla_block_rope.py new file mode 100644 index 000000000000..e41557f5ca9c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_mla_block_rope.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The block RoPE work of one MLA drafter layer in one CuTe DSL launch: the absorbed query's +rope columns and the block latents normalized, rotated and stored into the paged pool.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._dspark_mla_rope import dspark_mla_block_rope +from .k3_dspark_rmsnorm_rope import rmsnorm_rope_reference + + +class K3DsparkMlaBlockRope(OpWrapper): + """Writes ``query[..., kv_lora_rank:]`` = the RoPE of ``q[..., qk_nope_dim:]`` per head and + ``pool[pages[t], offsets[t]]`` = ``kv[t]`` with its latent RMSNorm (times ``norm_weight``) + and its rope columns rotated, both with ``freqs[positions[t]]``. Returns None; the rest of + ``query`` and of the pool is untouched. + + ``q`` is BF16 ``[T, H, qk_nope_dim + rope_dim]``, ``query`` BF16 ``[T, H, kv_lora_rank + + rope_dim]`` (unit inner strides, heads packed), ``kv`` BF16 ``[T, head_dim]`` row-strided, + ``pool`` ``[pages, page_size, head_dim]`` BF16, ``positions`` / ``pages`` / ``offsets`` + int64 ``[T]``, ``freqs`` FP32 ``[P, rope_dim // 2, 2]`` contiguous. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS = ( + Cell( + why="helix16 ULL c4: 4 requests x 8 block tokens (DSpark 7), 64 heads, latents read " + "in place from the fused [q_a | kv_a] GEMM output (row stride 2112), 128-token pages", + spec=dict(tokens=32, kv_row_stride=2112), + ), + Cell( + why="helix16 ULL c1: one request, 8 block tokens", + spec=dict(tokens=8, kv_row_stride=2112), + ), + Cell( + why="helix16 ULL c2: 2 requests x 8 block tokens, kv_a rows from their own projection " + "(compact, 576)", + spec=dict(tokens=16, kv_row_stride=576), + ), + ) + + note = """ + Bitwise ``cute_dsl_dspark_rope_into(q, freqs[positions], query, H, rope_dim, kv_lora_rank)`` + and ``k3_scatter_rows_to_pool(k3_dspark_rmsnorm_rope(kv, norm_weight, freqs[positions], 1, + rope_dim, eps, True, True, False, kv_lora_rank), pool, pages, offsets)`` in one launch + (``dspark_tail_v2``). Distinct (page, offset) per token. A process-wide compile cache holds + one kernel per (H, qk_nope_dim, kv_lora_rank, rope_dim, eps). + """ + + def __call__( + self, + q: torch.Tensor, + query: torch.Tensor, + kv: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + eps: float, + ) -> None: + dspark_mla_block_rope( + q, query, kv, norm_weight, positions, freqs, pool, pages, offsets, eps + ) + + def reference( + self, + q: torch.Tensor, + query: torch.Tensor, + kv: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``(query, pool)`` after the call (clones).""" + num_tokens, num_heads, q_dim = q.shape + kv_lora_rank = norm_weight.shape[0] + rope_dim = query.shape[-1] - kv_lora_rank + table = freqs[positions.reshape(-1)] + query_out = query.clone() + rotated = rmsnorm_rope_reference( + q[..., q_dim - rope_dim :], None, table, num_heads, rope_dim, 0.0, False, False, False + ) + query_out[..., kv_lora_rank:] = rotated + latents = rmsnorm_rope_reference( + kv, norm_weight, table, 1, rope_dim, eps, True, True, False, kv_lora_rank + ) + pool_out = pool.clone() + pool_out[pages.reshape(-1), offsets.reshape(-1)] = latents + return query_out, pool_out + + def is_valid( + self, + q: torch.Tensor, + query: torch.Tensor, + kv: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + eps: float, + ) -> None: + if positions.numel(): + assert int(positions.max()) < freqs.shape[0] and int(positions.min()) >= 0 + keys = pages.reshape(-1) * pool.shape[1] + offsets.reshape(-1) + assert keys.unique().numel() == keys.numel(), "two block tokens share a pool slot" + assert int(offsets.max()) < pool.shape[1] and int(pages.max()) < pool.shape[0] + + +k3_dspark_mla_block_rope = K3DsparkMlaBlockRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_mla_ctx_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_mla_ctx_norm_rope.py new file mode 100644 index 000000000000..ea91d5c28f8a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_mla_ctx_norm_rope.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RMSNorm + RoPE of the MLA drafter's context latents of every layer in one CuTe DSL launch, +the RoPE table looked up in the kernel.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._dspark_mla_rope import dspark_mla_ctx_norm_rope +from .k3_dspark_rmsnorm_rope import rmsnorm_rope_reference + + +class K3DsparkMlaCtxNormRope(OpWrapper): + """``[L, N, kv_lora_rank + rope_dim]`` BF16, new: for each layer ``l`` the RMSNorm of + ``latent[:, l, :kv_lora_rank]`` times ``norm_weight[l]``, and the latent's RoPE columns + rotated by ``freqs[positions]`` (interleaved pairs, not normalized). + + ``latent`` is BF16 ``[N, L, head_dim]`` with contiguous ``(L, head_dim)`` rows and any token + stride; ``norm_weight`` BF16 ``[L, kv_lora_rank]``; ``positions`` int64 ``[N]``; ``freqs`` + FP32 ``[P, rope_dim // 2, 2]`` (cos, sin), contiguous. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS = ( + Cell( + why="helix16 ULL: the accepted tokens of 4 requests x 8 (DSpark 7), 5 layers", + spec=dict(rows=32, layers=5), + ), + Cell( + why="helix8 c240: 20 requests x 6 (DSpark 5), 5 layers", + spec=dict(rows=120, layers=5), + ), + Cell( + why="helix8 / helix16: one accepted token", + spec=dict(rows=1, layers=5), + ), + ) + + note = """ + Bitwise ``k3_dspark_rmsnorm_rope(latent[:, l], norm_weight[l], freqs[positions], 1, rope_dim, + eps, True, True, False, kv_lora_rank)`` per layer, in one launch (``dspark_tail_v2``). A + process-wide compile cache holds one kernel per (L, kv_lora_rank, rope_dim, eps). + """ + + def __call__( + self, + latent: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + eps: float, + ) -> torch.Tensor: + return dspark_mla_ctx_norm_rope(latent, norm_weight, positions, freqs, eps) + + def reference( + self, + latent: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + eps: float, + ) -> torch.Tensor: + num_layers = latent.shape[1] + kv_lora_rank = norm_weight.shape[1] + rope_dim = latent.shape[-1] - kv_lora_rank + table = freqs[positions] + return torch.stack( + [ + rmsnorm_rope_reference( + latent[:, layer], + norm_weight[layer], + table, + 1, + rope_dim, + eps, + True, + True, + False, + kv_lora_rank, + ) + for layer in range(num_layers) + ] + ) + + def is_valid( + self, + latent: torch.Tensor, + norm_weight: torch.Tensor, + positions: torch.Tensor, + freqs: torch.Tensor, + eps: float, + ) -> None: + # The kernel indexes the table by position without a bound. + if positions.numel(): + assert int(positions.max()) < freqs.shape[0], "a position past the RoPE table" + assert int(positions.min()) >= 0, "negative position" + + +k3_dspark_mla_ctx_norm_rope = K3DsparkMlaCtxNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_rmsnorm_rope.py new file mode 100644 index 000000000000..eb050f68098d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_rmsnorm_rope.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark RMSNorm + RoPE of BF16 rows read in place from a row-strided view (CuTe DSL).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._dspark_mla_rope import k3r_dspark_rmsnorm_rope + + +def rmsnorm_rope_reference( + x: torch.Tensor, + weight: Optional[torch.Tensor], + freqs: torch.Tensor, + num_heads: int, + rope_dim: int, + eps: float, + apply_weight: bool, + apply_rmsnorm: bool, + inverse_rope: bool, + norm_dim: Optional[int] = None, +) -> torch.Tensor: + """FP64 RMSNorm over the first ``norm_dim`` columns (all when None), scaled by ``weight``, + then the interleaved-pair RoPE of the last ``rope_dim`` columns with row ``r``'s + ``(cos, sin)`` from ``freqs[r // num_heads]``; rounded once to ``x``'s dtype, compact.""" + rows = x.reshape(-1, x.shape[-1]).double() + n = rows.shape[-1] if norm_dim is None else norm_dim + head = rows[:, :n] + if apply_rmsnorm: + head = head * torch.rsqrt(head.square().mean(dim=-1, keepdim=True) + eps) + if apply_weight: + head = head * weight.double() + out = torch.cat([head, rows[:, n:]], dim=-1) + if rope_dim > 0: + row_freqs = freqs.double().repeat_interleave(num_heads, dim=0)[: out.shape[0]] + cos, sin = row_freqs[..., 0], row_freqs[..., 1] + if inverse_rope: + sin = -sin + real = out[:, -rope_dim::2].clone() + imag = out[:, -rope_dim + 1 :: 2].clone() + out[:, -rope_dim::2] = real * cos - imag * sin + out[:, -rope_dim + 1 :: 2] = imag * cos + real * sin + return out.to(x.dtype).view(x.shape) + + +class K3DsparkRmsnormRope(OpWrapper): + """``[..., hidden]`` BF16, compact and new: RMSNorm of each row's first ``norm_dim`` + columns (all when None; skipped without ``apply_rmsnorm``), times ``weight`` (with + ``apply_weight``), then RoPE of the last ``rope_dim`` columns, pairs interleaved, + ``(cos, sin)`` of row ``r`` from ``freqs[r // num_heads]`` (``[P, rope_dim // 2, 2]`` FP32), + the inverse rotation with ``inverse_rope``. Columns between the norm and the rope pass + through. + + ``x`` is BF16 with a unit inner stride and any regular row stride (read in place); + ``weight`` BF16 ``[norm_dim]``, or None without ``apply_weight``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS = ( + Cell( + why="dep16 draft block latents, 16 requests x 6 rows (DSpark 5): compact [512 latent " + "| 64 rope] rows, the latent normalized with its weight", + spec=dict(rows=96, row_stride=576), + ), + Cell( + why="helix8 c240 draft block latents, 20 requests x 6 rows: compact rows", + spec=dict(rows=120, row_stride=576), + ), + Cell( + why="helix8 c24 / c42 draft block latents, 2 requests x 8 rows (DSpark 7), read in " + "place from the fused [q_a | kv_a] GEMM output (row stride 2112)", + spec=dict(rows=16, row_stride=2112), + ), + ) + + note = """ + ``cute_dsl_dspark_rmsnorm_rope``'s kernel and support check, compiled for row-strided input: + the MLA drafter normalizes column slices of its fused [q_a | kv_a] GEMM output without a + contiguous copy. FP32 throughout, one rounding per output. A process-wide compile cache holds + one kernel per (hidden, rope_dim, heads, eps, weight, norm, inverse, norm_dim); compile before + capture. Raises ValueError on input the support check refuses. + """ + + def __call__( + self, + x: torch.Tensor, + weight: Optional[torch.Tensor], + freqs: torch.Tensor, + num_heads: int, + rope_dim: int, + eps: float, + apply_weight: bool, + apply_rmsnorm: bool, + inverse_rope: bool, + norm_dim: Optional[int] = None, + ) -> torch.Tensor: + return k3r_dspark_rmsnorm_rope( + x, + weight, + freqs, + num_heads, + rope_dim, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + norm_dim, + ) + + def reference( + self, + x: torch.Tensor, + weight: Optional[torch.Tensor], + freqs: torch.Tensor, + num_heads: int, + rope_dim: int, + eps: float, + apply_weight: bool, + apply_rmsnorm: bool, + inverse_rope: bool, + norm_dim: Optional[int] = None, + ) -> torch.Tensor: + return rmsnorm_rope_reference( + x, + weight, + freqs, + num_heads, + rope_dim, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + norm_dim, + ) + + def is_valid( + self, + x: torch.Tensor, + weight: Optional[torch.Tensor], + freqs: torch.Tensor, + num_heads: int, + rope_dim: int, + eps: float, + apply_weight: bool, + apply_rmsnorm: bool, + inverse_rope: bool, + norm_dim: Optional[int] = None, + ) -> None: + rows = x.numel() // x.shape[-1] + # The kernel indexes freqs by row // num_heads without a bound. + assert freqs.shape[0] * num_heads >= rows, "fewer frequency rows than rows" + if apply_weight: + n = x.shape[-1] if norm_dim is None else norm_dim + assert weight.numel() == n, "weight must span the normalized columns" + + +k3_dspark_rmsnorm_rope = K3DsparkRmsnormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_rmsnorm_rope_into.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_rmsnorm_rope_into.py new file mode 100644 index 000000000000..138ba2d12c89 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_dspark_rmsnorm_rope_into.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark RMSNorm + RoPE of row-strided BF16 rows into a caller-owned output, through a +compiled handle (CuTe DSL): the per-layer context-latent path of the MLA drafter.""" + +from __future__ import annotations + +import functools +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._dspark_mla_rope import _compile_rmsnorm_rope_strided +from .k3_dspark_rmsnorm_rope import rmsnorm_rope_reference + + +@dataclass(frozen=True) +class CompiledRmsnormRope: + """A compiled ``k3_dspark_rmsnorm_rope`` kernel and the configuration it was built for.""" + + hidden_dim: int + rope_dim: int + num_heads: int + eps: float + apply_weight: bool + apply_rmsnorm: bool + inverse_rope: bool + norm_dim: int + kernel: Any + + +@functools.cache +def build_rmsnorm_rope_into( + hidden_dim: int, + rope_dim: int, + num_heads: int, + eps: float, + apply_weight: bool, + apply_rmsnorm: bool, + inverse_rope: bool, + norm_dim: int, +) -> CompiledRmsnormRope: + """The handle for this configuration (compiled once per process).""" + kernel = _compile_rmsnorm_rope_strided( + hidden_dim, rope_dim, num_heads, eps, apply_weight, apply_rmsnorm, inverse_rope, norm_dim + ) + return CompiledRmsnormRope( + hidden_dim, + rope_dim, + num_heads, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + norm_dim, + kernel, + ) + + +class K3DsparkRmsnormRopeInto(OpWrapper): + """``k3_dspark_rmsnorm_rope`` with ``compiled``'s configuration, written into ``out`` + (compact BF16 ``[rows, hidden_dim]``). Returns ``out``. + + ``x`` is BF16 ``[rows, hidden_dim]`` with a unit inner stride and any row stride (read in + place), ``weight`` BF16 ``[norm_dim]`` (or None without ``apply_weight``), ``freqs`` FP32 + ``[rows / num_heads, rope_dim // 2, 2]``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS = ( + Cell( + why="tp1 context server: one context chunk of max_num_tokens 16384 rows, each " + "layer's slice of the [N, 5, 576] latent GEMM output (row stride 2880)", + spec=dict(rows=16384, layers=5), + ), + Cell( + why="tp1 context server: a one-token chunk", + spec=dict(rows=1, layers=5), + ), + Cell( + why="dep16 generation: the accepted tokens of 16 requests, 16 x 6 rows", + spec=dict(rows=96, layers=5), + ), + ) + + note = """ + The drafter normalizes each layer's slice of its stacked kv_a GEMM output with one handle + built before the loop, without the support check per call (the caller checks once). Same + kernel and arithmetic as ``k3_dspark_rmsnorm_rope``. + """ + + def __call__( + self, + x: torch.Tensor, + weight: Optional[torch.Tensor], + freqs: torch.Tensor, + out: torch.Tensor, + compiled: CompiledRmsnormRope, + ) -> torch.Tensor: + compiled.kernel(x, weight, freqs, out) + return out + + def reference( + self, + x: torch.Tensor, + weight: Optional[torch.Tensor], + freqs: torch.Tensor, + out: torch.Tensor, + compiled: CompiledRmsnormRope, + ) -> torch.Tensor: + """``out`` after the call (new); ``out`` is not written.""" + return rmsnorm_rope_reference( + x, + weight, + freqs, + compiled.num_heads, + compiled.rope_dim, + compiled.eps, + compiled.apply_weight, + compiled.apply_rmsnorm, + compiled.inverse_rope, + compiled.norm_dim, + ) + + def is_valid( + self, + x: torch.Tensor, + weight: Optional[torch.Tensor], + freqs: torch.Tensor, + out: torch.Tensor, + compiled: CompiledRmsnormRope, + ) -> None: + # The handle is compiled for these widths; another width is read with its geometry. + assert x.shape[-1] == out.shape[-1] == compiled.hidden_dim, "width of another handle" + assert out.is_contiguous() and out.dtype == x.dtype == torch.bfloat16 + assert freqs.shape[0] * compiled.num_heads >= x.shape[0], "fewer frequency rows" + if compiled.apply_weight: + assert weight.numel() == compiled.norm_dim, "weight must span the normalized columns" + + +k3_dspark_rmsnorm_rope_into = K3DsparkRmsnormRopeInto() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_markov_chain_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_markov_chain_logits.py new file mode 100644 index 000000000000..99c04688a930 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_markov_chain_logits.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark vanilla-Markov greedy chain on full-vocab block logits, three launches per block +position (two Triton kernels around the bias GEMV).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.dspark_markov_fused import fused_markov_chain_logits +from ._reference import argmax_pairs + + +class K3FusedMarkovChainLogits(OpWrapper): + """``[B, K, vocab]`` FP32 biased logits (new): position ``i`` is + ``base_logits[:, i] + markov_w2 @ markov_w1[prev_i]`` with the bias in the weights' dtype + (``F.linear``) and the sum in FP32; ``prev_0 = first_prev_tokens`` (int32 / int64) and + ``prev_{i+1}`` the ``torch.argmax`` of position ``i``'s biased logits. A ``prev`` outside + ``[0, vocab)`` adds a zero bias. + + ``base_logits`` is FP32 ``[B, K, vocab]`` with a unit vocab stride; ``markov_w1`` and + ``markov_w2`` are contiguous ``[vocab, R]`` of one dtype. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # Full-vocab drafter logits (163840 columns, Markov rank 256): the deployments whose LM head + # is not vocab-parallel, with DSpark 5 (6 block positions). + CELLS = ( + Cell( + why="dep16 c383: 16 requests per attention-DP rank, int64 anchors", + spec=dict(reqs=16, block=6, anchor="int64"), + ), + Cell( + why="tp1 context server: 4 requests, int32 anchors, one outside the vocab", + spec=dict(reqs=4, block=6, anchor="int32"), + ), + ) + + note = """ + Bit-identical to ``dspark_markov_chain_logits`` driven by a plain full-vocab argmax: the bias + GEMV is that ``F.linear``, the add is FP32 and the argmax keeps NaN-first, lowest-index + order. Reached where the drafter's logits are full-vocab (the shared LM head gathers them) + and the vocab-parallel chains do not apply. 2048 columns per program; launch shapes depend + on (B, vocab) only. + """ + + def __call__( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + ) -> torch.Tensor: + return fused_markov_chain_logits(base_logits, first_prev_tokens, markov_w1, markov_w2) + + def reference( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + ) -> torch.Tensor: + """The biased logits; the bias is an FP64 product rounded to the weights' dtype.""" + num_reqs, block_size, vocab = base_logits.shape + out = torch.empty_like(base_logits) + prev = first_prev_tokens.to(torch.int64) + cols = torch.arange(vocab, device=base_logits.device).expand(num_reqs, -1) + w2 = markov_w2.double() + for step in range(block_size): + valid = (prev >= 0) & (prev < vocab) + row = markov_w1[prev.clamp(0, vocab - 1)].double() * valid[:, None] + bias = (row @ w2.t()).to(markov_w2.dtype).float() + out[:, step] = base_logits[:, step] + bias + _, prev = argmax_pairs(out[:, step], cols) + return out + + def is_valid( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + ) -> None: + # fused_markov_supported, which the caller checks: anything else is read with the wrong + # strides or dtype rather than rejected. + assert base_logits.dim() == 3 and base_logits.dtype == torch.float32 + assert base_logits.stride(-1) == 1, "unit vocab stride" + assert markov_w1.dtype == markov_w2.dtype and markov_w1.shape == markov_w2.shape + assert markov_w1.is_contiguous() and markov_w2.is_contiguous() + assert markov_w2.shape[0] == base_logits.shape[-1], "full-vocab logits" + + +k3_fused_markov_chain_logits = K3FusedMarkovChainLogits() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_select_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_select_accept.py new file mode 100644 index 000000000000..d9ac8a564e2f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_select_accept.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Vocab-parallel target selection, strict acceptance, the KDA replay record and the drafter's +KV-length preparation in one Triton launch (one program).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import fused_select_accept +from ._reference import argmax_pairs, interleaved_pairs, strict_accept + + +class K3FusedSelectAccept(OpWrapper): + """``(accepted [batch, K + 1], num_accepted [batch], rewind [num_gens] or None)``, int32. + + ``pairs`` is ``(buffer, row stride, slot stride == 2, slots)``: every rank's interleaved + ``(value, global index)`` float32 candidates of the ``num_contexts + num_gens * (K + 1)`` + rows (``K3VocabShardedLogits.select_pairs``). Each row's token is the candidate the + ``torch.argmax`` order selects; acceptance and ``force`` are those of ``k3_fused_strict_accept``. + ``kda_replay = (prev_num_accepted_tokens, state_indices, dummy_mask)`` also writes + ``prev[state_indices[c + g]] = max(num_accepted - 1, 0)`` for non-dummy generation rows + (``KDAReplayState.record_acceptance``); ``kv_lens`` also adds 1 to its generation rows and + returns ``rewind = 1 - num_accepted``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="helix8 c24/c42: 2 rows of DSpark 7, 8 ranks x 10 splits = 80 pair slots, with the " + "KDA replay record and the KV lengths", + spec=dict( + num_contexts=0, num_gens=2, draft_len=7, tp=8, splits=10, replay=True, force=0.0 + ), + ), + Cell( + why="helix8 c75/c174: 8 rows of DSpark 7, 8 ranks x 4 splits = 32 pair slots", + spec=dict( + num_contexts=0, num_gens=8, draft_len=7, tp=8, splits=4, replay=True, force=0.0 + ), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4): 4 rows of DSpark 7, 16 ranks x 5 splits = 80 slots", + spec=dict( + num_contexts=0, num_gens=4, draft_len=7, tp=16, splits=5, replay=True, force=0.0 + ), + ), + Cell( + why="helix16_moetp4ep4 (ULL c1): one generation row", + spec=dict( + num_contexts=0, num_gens=1, draft_len=7, tp=16, splits=5, replay=True, force=0.0 + ), + ), + Cell( + why="helix16 shape with a context row ahead, without the replay record or KV lengths", + spec=dict( + num_contexts=1, num_gens=4, draft_len=7, tp=16, splits=5, replay=False, force=0.0 + ), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4) under a fractional forced acceptance (2.5)", + spec=dict( + num_contexts=0, num_gens=4, draft_len=7, tp=16, splits=5, replay=True, force=2.5 + ), + ), + ) + + note = """ + One program of 4 warps covers every row, so the caller limits it to TAIL_V2_MAX_GENS (16) + generation requests. All replay slots are read before any is written; generation rows must + map to distinct slots. The fractional force advances its device counter by one. No host + synchronization; launch shape fixed by host batch sizes. + """ + + def __call__( + self, + pairs: tuple, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, + kda_replay: Optional[tuple] = None, + kv_lens: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + return fused_select_accept( + pairs, draft_tokens, num_contexts, batch_size, force, kda_replay, kv_lens + ) + + def reference( + self, + pairs: tuple, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, + kda_replay: Optional[tuple] = None, + kv_lens: Optional[torch.Tensor] = None, + ) -> tuple: + """``(accepted, num_accepted, rewind, prev_num_accepted_tokens, kv_lens)``: the op's + outputs (context rows hold 0 past column 0) and the expected post-state of the replay + count and of ``kv_lens`` (None when not passed). Nothing is written.""" + buf, stride_row, stride_slot, n_slots = pairs + num_gens = batch_size - num_contexts + k = draft_tokens.shape[-1] + rows = num_contexts + num_gens * (k + 1) + vals, idx = interleaved_pairs(buf, stride_row, stride_slot, n_slots, rows) + _, target = argmax_pairs(vals, idx) + accepted, num_accepted = strict_accept( + target, draft_tokens, num_contexts, batch_size, force + ) + n = num_accepted[num_contexts:] + prev_out = None + if kda_replay is not None: + prev, state_indices, dummy = kda_replay + prev_out = prev.clone() + slots = state_indices[num_contexts:batch_size].to(torch.int64) + keep = dummy[num_contexts:batch_size].to(torch.bool) + prev_out[slots] = torch.where(keep, prev[slots], (n - 1).clamp(min=0).to(prev.dtype)) + rewind = kv_out = None + if kv_lens is not None: + rewind = (1 - n).to(torch.int32) + kv_out = kv_lens.clone() + kv_out[num_contexts:batch_size] += 1 + return accepted, num_accepted, rewind, prev_out, kv_out + + def is_valid( + self, + pairs: tuple, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, + kda_replay: Optional[tuple] = None, + kv_lens: Optional[torch.Tensor] = None, + ) -> None: + num_gens = batch_size - num_contexts + assert draft_tokens.shape[0] >= num_gens, "one draft row per generation request" + if kda_replay is not None and num_gens > 0: + slots = kda_replay[1][num_contexts:batch_size] + # Two rows on one slot race on the same store. + assert slots.unique().numel() == num_gens, "generation rows share a replay slot" + if force is not None and len(force) > 1: + pool = force[3] + assert pool.numel() & (pool.numel() - 1) == 0, "the force pool must be a power of two" + + +k3_fused_select_accept = K3FusedSelectAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_strict_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_strict_accept.py new file mode 100644 index 000000000000..093121e8ee4a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_fused_strict_accept.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Strict acceptance of a uniform-width draft against the sampled target tokens, with the +forced-acceptance override, in one Triton launch (one program).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import fused_strict_accept +from ._reference import strict_accept + + +class K3FusedStrictAccept(OpWrapper): + """``(accepted [batch, K + 1], num_accepted [batch])``, both int32 and new. + + ``target_tokens`` is ``[num_contexts + num_gens * (K + 1)]`` of any integer dtype, + ``draft_tokens`` ``[num_gens, K]`` with a unit inner stride. A context row accepts its one + target token (column 0 only is defined); a generation row takes its ``K + 1`` target tokens + and accepts ``1 +`` the number of leading draft tokens equal to them. ``force``: None, + ``(base_total,)`` (every generation row accepts ``base_total``) or ``(base_total, frac, + max_total, pool, counter, counter_stride, slot_stride)``, which advances ``counter`` by one and + adds 1 where the pool draw is below ``frac``, capped at ``max_total``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="dep16 (c383): 16 generation rows, DSpark 5", + spec=dict(num_contexts=0, num_gens=16, draft_len=5, force=0.0), + ), + Cell( + why="helix8 c240: 20 generation rows of DSpark 5, past the 16 the selection kernel takes", + spec=dict(num_contexts=0, num_gens=20, draft_len=5, force=0.0), + ), + Cell( + why="dep16 shape with 3 context rows ahead of 13 generation rows: column 0 only", + spec=dict(num_contexts=3, num_gens=13, draft_len=5, force=0.0), + ), + Cell( + why="dep16 under a whole forced acceptance (3.0): every row accepts 4", + spec=dict(num_contexts=0, num_gens=16, draft_len=5, force=3.0), + ), + Cell( + why="helix8 c240 under a fractional forced acceptance (2.5): the pool draw and the " + "device counter", + spec=dict(num_contexts=0, num_gens=20, draft_len=5, force=2.5), + ), + ) + + note = """ + One program of 4 warps covers the whole batch (block size the next power of two of the larger + of the context and generation counts). The values are those of + ``SpecWorkerBase._sample_and_accept_draft_tokens_base`` with the same force draws. No host + synchronization; the launch shape depends on host batch sizes only. + """ + + def __call__( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return fused_strict_accept(target_tokens, draft_tokens, num_contexts, batch_size, force) + + def reference( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The op's outputs; context rows hold 0 past column 0, which the op leaves undefined. + The fractional ``force`` reads ``counter + 1`` without writing ``counter``.""" + return strict_accept(target_tokens, draft_tokens, num_contexts, batch_size, force) + + def is_valid( + self, + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + force: Optional[tuple] = None, + ) -> None: + num_gens = batch_size - num_contexts + k = draft_tokens.shape[-1] + # The kernel reads target rows at (K + 1) strides past the contexts: a shorter vector is + # read out of bounds instead of rejected. + assert target_tokens.numel() >= num_contexts + num_gens * (k + 1), "target too short" + assert draft_tokens.shape[0] >= num_gens, "one draft row per generation request" + if force is not None and len(force) > 1: + pool = force[3] + # The draw masks the pool index with numel - 1. + assert pool.numel() & (pool.numel() - 1) == 0, "the force pool must be a power of two" + + +k3_fused_strict_accept = K3FusedStrictAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_add_local_argmax.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_add_local_argmax.py new file mode 100644 index 000000000000..5780e7086724 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_add_local_argmax.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One step of the multi-launch vocab-parallel DSpark Markov chain: the Markov bias added to a +block position's shard logits and each split's argmax published for the exchange (Triton).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.vocab_parallel_kernels import _PairLayout, markov_add_local_argmax +from ._reference import publish_pairs + + +class K3MarkovAddLocalArgmax(OpWrapper): + """``k3_vocab_local_select`` (argmax) of ``step_logits + bias`` (``[requests, shard]``, each + upcast to FP32 and added in FP32): every split's ``(value, col_base + column)`` pair into + ``layout.buffer``, other ranks' columns zeroed when ``zero_width > 0``. Returns None.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The multi-launch chain of helix8 c240 (20 requests, a 6-position block, 20480-column + # shards): the step's logits are a strided view of the block logits. + CELLS = ( + Cell( + why="helix8 c240, rank 0, step 0: 20 requests, 10 splits of 2048 columns", + spec=dict(tp=8, rank=0, reqs=20, block=6, step=0), + ), + Cell( + why="helix8 c240, last rank, last step: the zero fill around the top slots", + spec=dict(tp=8, rank=7, reqs=20, block=6, step=5), + ), + ) + + note = """ + ``step_logits`` is one block position of the shard block logits (a strided view with a unit + inner stride); ``bias`` is ``F.linear(markov_w1[prev], w2_shard)`` in the weights' dtype. + Grid (requests, splits), 8 warps. + """ + + def __call__( + self, step_logits: torch.Tensor, bias: torch.Tensor, col_base: int, layout: _PairLayout + ) -> None: + markov_add_local_argmax(step_logits, bias, col_base, layout) + + def reference( + self, step_logits: torch.Tensor, bias: torch.Tensor, col_base: int, layout: _PairLayout + ) -> torch.Tensor: + """``layout.buffer`` after the call (a clone).""" + return publish_pairs(layout, step_logits.float() + bias.float(), col_base) + + def is_valid( + self, step_logits: torch.Tensor, bias: torch.Tensor, col_base: int, layout: _PairLayout + ) -> None: + assert step_logits.stride(-1) == 1 and bias.stride(-1) == 1, "unit inner strides" + assert bias.shape == step_logits.shape, "one bias row per request, shard wide" + assert layout.splits * layout.chunk >= step_logits.shape[1], "the splits miss columns" + + +k3_markov_add_local_argmax = K3MarkovAddLocalArgmax() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_chain.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_chain.py new file mode 100644 index 000000000000..d3103dd7128d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_chain.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark greedy Markov chain of one draft block over a vocab-sharded LM head in one Triton +launch, the per-rank pairs exchanged through a multicast buffer every program polls.""" + +from __future__ import annotations + +from typing import Sequence, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.vocab_parallel_kernels import markov_chain +from ._reference import markov_chain_tokens + +Shards = Union[torch.Tensor, Sequence[torch.Tensor]] + + +class K3MarkovChain(OpWrapper): + """``tokens`` ``[B, K]`` int32 (written and returned): for each block position ``i``, the + global argmax over the vocab of ``local_block[:, i] + bf16(markov_w2 @ markov_w1[prev_i])`` + (added in FP32), which is draft token ``i`` and ``prev_{i+1}``; ``prev_0`` is + ``first_prev_tokens`` and a ``prev`` outside ``[0, vocab)`` adds nothing. + + Each rank passes its ``[B, K, shard]`` block logits (unit inner stride), its ``[shard, R]`` + rows of ``markov_w2`` and the full ``[vocab, R]`` ``markov_w1`` (contiguous). ``partial``, + ``tickets`` and ``launch`` are the rank's scratch (``_CHAIN_MAX_STEPS * programs * 16 * 2`` + float32, 16 int32 zeros, 1 int32 launch ordinal); ``mcast`` / ``local`` the multicast and + local views of the three-set exchange of ``set_floats`` floats each (one buffer on one rank), + initialized to -0.0. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # One rank of the helix16 ULL group (10240-column shards, 80 programs; Markov rank 256, + # 8 block positions); the other ranks' pairs are in the exchange. The chain where the CuTe DSL + # chain (``k3_dspark_markov_greedy``) does not apply. A helix8 shard (160 programs) needs + # more SMs than one GB300 has, so its cells belong to the archs that hold them. + CELLS = ( + Cell( + why="helix16 ULL c1, rank 0 of 16: one request, int32 anchor", + spec=dict(tp=16, rank=0, reqs=1, block=8, anchor="int32"), + ), + Cell( + why="helix16 ULL c4, last rank: 4 requests, int64 anchors, one outside the vocab", + spec=dict(tp=16, rank=15, reqs=4, block=8, anchor="int64"), + ), + ) + + note = """ + One program per 128 shard columns; every program must be resident (the caller checks against + the SM count) because each waits for all ranks' pairs every step. Up to 16 requests and 16 + positions. The bias is an FP32 dot rounded to the weights' dtype, summed in a different order + than ``F.linear``: a token can differ where two vocab entries tie within that rounding. + Launched with PDL; the program waits before reading the logits. Every rank makes the same + sequence of calls; the launch ordinal rotates the exchange sets, so no fence is needed. + """ + + def __call__( + self, + local_block: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: torch.Tensor, + tokens: torch.Tensor, + partial: torch.Tensor, + tickets: torch.Tensor, + launch: torch.Tensor, + mcast: torch.Tensor, + local: torch.Tensor, + set_floats: int, + vocab: int, + col_base: int, + tp: int, + tp_rank: int, + ) -> torch.Tensor: + return markov_chain( + local_block, + first_prev_tokens, + markov_w1, + w2_shard, + tokens, + partial, + tickets, + launch, + mcast, + local, + set_floats, + vocab, + col_base, + tp, + tp_rank, + ) + + def reference( + self, + local_block: Shards, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: Shards, + tokens: torch.Tensor, + partial: torch.Tensor, + tickets: torch.Tensor, + launch: torch.Tensor, + mcast: torch.Tensor, + local: torch.Tensor, + set_floats: int, + vocab: int, + col_base: int, + tp: int, + tp_rank: int, + ) -> torch.Tensor: + """The tokens (new). A collective: with ``tp > 1`` ``local_block`` and ``w2_shard`` are + the group's, in rank order (a tensor is taken as the whole vocab); ``col_base`` is then + not read. The scratch and exchange buffers are not read.""" + if not isinstance(local_block, torch.Tensor): + local_block = torch.cat(list(local_block), dim=-1) + w2_shard = torch.cat(list(w2_shard), dim=0) + return markov_chain_tokens(local_block, first_prev_tokens, markov_w1, w2_shard, vocab) + + def is_valid( + self, + local_block: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + w2_shard: torch.Tensor, + tokens: torch.Tensor, + partial: torch.Tensor, + tickets: torch.Tensor, + launch: torch.Tensor, + mcast: torch.Tensor, + local: torch.Tensor, + set_floats: int, + vocab: int, + col_base: int, + tp: int, + tp_rank: int, + ) -> None: + num_reqs, block_size, shard = local_block.shape + assert num_reqs <= 16 and block_size <= 16, "the exchange sets hold 16 x 16" + assert w2_shard.shape[0] == shard and w2_shard.stride(-1) == 1, "w2 rows of the shard" + assert markov_w1.is_contiguous() and local_block.stride(-1) == 1 + assert set_floats == 16 * 16 * max(tp, 1) * 2, "set_floats of another exchange" + assert tokens.dtype == torch.int32 and tuple(tokens.shape) == (num_reqs, block_size) + + +k3_markov_chain = K3MarkovChain() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_reduce_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_reduce_embed.py new file mode 100644 index 000000000000..972af3ead49a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_markov_reduce_embed.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One step of the multi-launch vocab-parallel DSpark Markov chain: resolve the previous token +from the exchanged pairs and gather its ``markov_w1`` row (Triton).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.vocab_parallel_kernels import markov_reduce_embed +from ._reference import argmax_pairs, interleaved_pairs + + +class K3MarkovReduceEmbed(OpWrapper): + """For each request ``b < emb.shape[0]``: ``prev = anchors[b]`` at ``step == 0``, else the + winner of the exchanged pairs of row ``b`` (as ``k3_vocab_reduce_pairs``), also stored as + ``tokens[b, step - 1]`` (int32); then ``emb[b] = markov_w1[prev]``, or zeros for a ``prev`` + outside ``[0, vocab)``. Returns None.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # helix8 c240 runs 20 requests, more than the one-launch chains take: its draft block (6 + # positions, Markov rank 256) goes through the multi-launch chain over 8 ranks x 10 splits. + CELLS = ( + Cell( + why="helix8 c240, step 0: the anchors, two of them outside the vocab (zero rows)", + spec=dict(tp=8, reqs=20, block=6, step=0), + ), + Cell( + why="helix8 c240, step 3: the winner of 80 exchanged slots, stored as token 2", + spec=dict(tp=8, reqs=20, block=6, step=3), + ), + Cell( + why="helix8 c240, step 6 = block: the last token only", + spec=dict(tp=8, reqs=20, block=6, step=6), + ), + ) + + note = """ + The multi-launch chain (``K3VocabParallelHead.markov_tokens``) that runs where the one-launch + chains do not apply; its bias GEMV (``F.linear``) and local argmax run between two calls. + One program of 4 warps per request. + """ + + def __call__( + self, + pairs: torch.Tensor, + stride_row: int, + stride_slot: int, + n_slots: int, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + emb: torch.Tensor, + tokens: torch.Tensor, + vocab: int, + step: int, + ) -> None: + markov_reduce_embed( + pairs, + stride_row, + stride_slot, + n_slots, + anchors, + markov_w1, + emb, + tokens, + vocab, + step, + ) + + def reference( + self, + pairs: torch.Tensor, + stride_row: int, + stride_slot: int, + n_slots: int, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + emb: torch.Tensor, + tokens: torch.Tensor, + vocab: int, + step: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``(emb, tokens)`` after the call (clones).""" + num_reqs = emb.shape[0] + tokens_out = tokens.clone() + if step == 0: + prev = anchors[:num_reqs].to(torch.int64) + else: + vals, idx = interleaved_pairs(pairs, stride_row, stride_slot, n_slots, num_reqs) + _, prev = argmax_pairs(vals, idx) + tokens_out[:num_reqs, step - 1] = prev.to(tokens.dtype) + valid = (prev >= 0) & (prev < vocab) + rows = markov_w1[prev.clamp(0, markov_w1.shape[0] - 1)] + emb_out = torch.where(valid[:, None], rows, torch.zeros_like(rows)).to(emb.dtype) + return emb_out, tokens_out + + def is_valid( + self, + pairs: torch.Tensor, + stride_row: int, + stride_slot: int, + n_slots: int, + anchors: torch.Tensor, + markov_w1: torch.Tensor, + emb: torch.Tensor, + tokens: torch.Tensor, + vocab: int, + step: int, + ) -> None: + # Rows are addressed as w1 + prev * RANK and emb + b * RANK. + assert markov_w1.is_contiguous() and emb.is_contiguous(), "w1 and emb must be contiguous" + assert emb.shape[1] == markov_w1.shape[1] and emb.dtype == markov_w1.dtype + assert tokens.dtype == torch.int32 and tokens.stride(-1) == 1, "tokens: int32 rows" + assert vocab <= markov_w1.shape[0], "vocab exceeds markov_w1's rows" + + +k3_markov_reduce_embed = K3MarkovReduceEmbed() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_next_new_tokens.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_next_new_tokens.py new file mode 100644 index 000000000000..365a70b913ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_next_new_tokens.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The next step's new tokens of a speculative batch (the last accepted token, then the new +draft), optionally with the KV-length rewind after the draft forward, in one Triton launch.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import fused_next_new_tokens, fused_next_new_tokens_rewind + + +class K3NextNewTokens(OpWrapper): + """``out[b] = cat([accepted_tokens[b, num_accepted_tokens[b] - 1], next_draft_tokens[b]])``, + ``[batch, K + 1]`` in the promoted dtype of the two token tensors (new). + + ``accepted_tokens`` ``[>= batch, >= K + 1]`` and ``next_draft_tokens`` ``[batch, K]`` have + unit inner strides. With ``kv_lens`` and ``rewind`` (``[num_gens]``, contiguous) also + ``kv_lens[num_contexts + g] = max(kv_lens[num_contexts + g] - rewind[g], 0)``, in one program. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="dep16 (c383): 16 rows of DSpark 5, 128 rows per program", + spec=dict(batch=16, draft_len=5, rewind=False), + ), + Cell( + why="helix8 c240: 20 rows of DSpark 5", + spec=dict(batch=20, draft_len=5, rewind=False), + ), + Cell( + why="helix8 c75/c174: 8 rows of DSpark 7 with the KV-length rewind in one program", + spec=dict(batch=8, draft_len=7, rewind=True), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4): 4 rows of DSpark 7 with the KV-length rewind", + spec=dict(batch=4, draft_len=7, rewind=True), + ), + ) + + note = """ + ``SpecWorkerBase._prepare_next_new_tokens`` with identity batch indices. Without the rewind + the kernel runs 128 rows per program; with it one program covers the batch and the rewind. + """ + + def __call__( + self, + accepted_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + kv_lens: Optional[torch.Tensor] = None, + rewind: Optional[torch.Tensor] = None, + num_contexts: int = 0, + ) -> torch.Tensor: + if kv_lens is None: + return fused_next_new_tokens(accepted_tokens, next_draft_tokens, num_accepted_tokens) + return fused_next_new_tokens_rewind( + accepted_tokens, next_draft_tokens, num_accepted_tokens, kv_lens, rewind, num_contexts + ) + + def reference( + self, + accepted_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + kv_lens: Optional[torch.Tensor] = None, + rewind: Optional[torch.Tensor] = None, + num_contexts: int = 0, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """``(out, kv_lens after the call or None)``; nothing is written.""" + batch, k = next_draft_tokens.shape + dtype = torch.promote_types(accepted_tokens.dtype, next_draft_tokens.dtype) + out = torch.empty((batch, k + 1), dtype=dtype, device=accepted_tokens.device) + last = (num_accepted_tokens[:batch].to(torch.int64) - 1).unsqueeze(1) + out[:, 0] = accepted_tokens[:batch].gather(1, last).squeeze(1).to(dtype) + out[:, 1:] = next_draft_tokens.to(dtype) + kv_out = None + if kv_lens is not None: + num_gens = rewind.shape[0] + kv_out = kv_lens.clone() + gen = kv_out[num_contexts : num_contexts + num_gens] + kv_out[num_contexts : num_contexts + num_gens] = (gen - rewind.to(gen.dtype)).clamp( + min=0 + ) + return out, kv_out + + def is_valid( + self, + accepted_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + kv_lens: Optional[torch.Tensor] = None, + rewind: Optional[torch.Tensor] = None, + num_contexts: int = 0, + ) -> None: + # The kernel reads accepted[b, num_accepted[b] - 1] unchecked. + batch = next_draft_tokens.shape[0] + n = num_accepted_tokens[:batch] + assert bool(((n >= 1) & (n <= accepted_tokens.shape[1])).all()), "num_accepted range" + assert (kv_lens is None) == (rewind is None), "kv_lens and rewind go together" + + +k3_next_new_tokens = K3NextNewTokens() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_scatter_rows_to_pool.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_scatter_rows_to_pool.py new file mode 100644 index 000000000000..28cdcebdd8f2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_scatter_rows_to_pool.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Scatter of the MLA drafter's block latents into its paged pool (Triton, one program per +row).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import scatter_rows_to_pool +from .k3_store_context_latents import to_e4m3_saturated + + +class K3ScatterRowsToPool(OpWrapper): + """``pool[pages[i], offsets[i]] = src[i]`` for a ``[pages, page_size, D]`` pool view; an FP8 + E4M3 pool stores ``to_draft_fp8(src[i])``. With ``rows_mask`` (bool or integer ``[rows]``, + contiguous) rows where it is zero write nothing. Returns None. + + ``src`` is ``[rows, D]`` with a unit inner stride, the pool's dtype or BF16 into FP8. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="dep16 (c383): 16 blocks of 6 latents (576 wide) into the FP8 pool, every row", + spec=dict(rows=96, rows_mask=False), + ), + Cell( + why="helix8 c75/c174: 8 blocks of 8, only the rows on this rank's pages", + spec=dict(rows=64, rows_mask=True), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4): 4 blocks of 8, only this rank's rows", + spec=dict(rows=32, rows_mask=True), + ), + ) + + note = """ + The eager advanced-index assignment as one launch; two rows on one slot race (the drafter's + block positions are distinct). A Helix drafter passes the block positions this rank owns as + ``rows_mask``. + """ + + def __call__( + self, + src: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + rows_mask: Optional[torch.Tensor] = None, + ) -> None: + scatter_rows_to_pool(src, pool, pages, offsets, rows_mask=rows_mask) + + def reference( + self, + src: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + rows_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The pool after the call (a clone).""" + new = pool.clone() + write = torch.ones(src.shape[0], dtype=torch.bool, device=src.device) + if rows_mask is not None: + write = rows_mask.reshape(-1) != 0 + vals, target = src, new + if new.dtype == torch.float8_e4m3fn: + # Assigned as bytes: the bits are the same. + vals, target = to_e4m3_saturated(src).view(torch.uint8), new.view(torch.uint8) + target[pages.to(torch.int64)[write], offsets.to(torch.int64)[write]] = vals[write] + return new + + def is_valid( + self, + src: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + rows_mask: Optional[torch.Tensor] = None, + ) -> None: + keys = pages.reshape(-1).to(torch.int64) * pool.shape[1] + offsets.reshape(-1) + if rows_mask is not None: + keys = keys[rows_mask.reshape(-1) != 0] + # Two rows on one slot: which one lands is unspecified. + assert keys.unique().numel() == keys.numel(), "rows target the same slot" + + +k3_scatter_rows_to_pool = K3ScatterRowsToPool() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_spec_decode_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_spec_decode_head.py new file mode 100644 index 000000000000..3170446a518f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_spec_decode_head.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The head of a speculative decode step in one Triton launch: the overlap-scheduler correction +of position ids and KV lengths and the Helix verify-group buffers derived from the positions.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# The module, not the function: kernels.spec_decode_head imports this entry for its hook. +from ...kernels import spec_decode_head as _head + + +def _helix_local_len(g: torch.Tensor, tokens_per_block: int, cp_size: int, cp_rank: int): + """Rank-local prefix length of the first ``g`` global tokens of a Helix-sharded sequence.""" + ledger = tokens_per_block * cp_size + full = torch.div(g, ledger, rounding_mode="floor") + rest = g - full * ledger - cp_rank * tokens_per_block + return full * tokens_per_block + rest.clamp(min=0, max=tokens_per_block) + + +class K3SpecDecodeHead(OpWrapper): + """In place, returns None. With ``has_overlap``: ``position_ids[:num_pos] += + pos_offsets[:num_pos]`` and ``kv_lens[kv_start : kv_start + num_kv] += kv_offsets[:num_kv]``. + With ``num_helix > 0``, after ``helix_positions[:num_helix] += pos_offsets[:num_helix]`` + (with ``has_overlap`` only), what ``TrtllmAttentionMetadata.recompute_helix_spec_buffers`` + writes: ``helix_slots`` (the rank-local write slot, -1 on other ranks' pages), + ``helix_bounds`` (the rank-local KV bound of each token), ``helix_zero_kv = bounds == 0`` and + ``kv_lens[kv_bound_start + s]`` = the bound of sequence ``s``'s last token. + + Integer tensors are int32 with unit stride; ``helix_zero_kv`` is bool; ``num_helix`` is a + multiple of ``tokens_per_seq``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS = ( + Cell( + why="helix16 ULL c4: 4 requests x 8 tokens (DSpark 7), 16 ranks, 128-token pages, " + "with the overlap scheduler's offsets, on the first rank", + spec=dict(num_seqs=4, tokens_per_seq=8, cp_size=16, cp_rank=0, overlap=True), + ), + Cell( + why="helix16 ULL c1: one request on the last rank", + spec=dict(num_seqs=1, tokens_per_seq=8, cp_size=16, cp_rank=15, overlap=True), + ), + Cell( + why="helix8 c240: 20 requests x 6 tokens (DSpark 5), 8 ranks", + spec=dict(num_seqs=20, tokens_per_seq=6, cp_size=8, cp_rank=3, overlap=True), + ), + Cell( + why="helix8 c75 / c174: 8 requests x 8 tokens without previous-step offsets", + spec=dict(num_seqs=8, tokens_per_seq=8, cp_size=8, cp_rank=7, overlap=False), + ), + ) + + note = """ + The ``RuntimeHooks.decode_prologue`` of helix8 / helix16 (``SpecDecodeHead``) calls it with + the preparer's previous-step offsets and the metadata's Helix buffers; the zero-bound flags go + to the hook's scratch. One program of 4 warps, 512 lanes per pass; the KV-bound pass reads the + bounds back after a barrier. The scalars are not specialized, so one compile serves every + step. + """ + + def __call__( + self, + position_ids: torch.Tensor, + pos_offsets: torch.Tensor, + num_pos: int, + kv_lens: torch.Tensor, + kv_offsets: torch.Tensor, + kv_start: int, + num_kv: int, + *, + has_overlap: bool, + helix_positions: Optional[torch.Tensor] = None, + helix_slots: Optional[torch.Tensor] = None, + helix_bounds: Optional[torch.Tensor] = None, + helix_zero_kv: Optional[torch.Tensor] = None, + num_helix: int = 0, + tokens_per_block: int = 1, + cp_size: int = 1, + cp_rank: int = 0, + kv_bound_start: int = 0, + tokens_per_seq: int = 1, + ) -> None: + _head.fused_spec_head_prep( + position_ids, + pos_offsets, + num_pos, + kv_lens, + kv_offsets, + kv_start, + num_kv, + has_overlap=has_overlap, + helix_positions=helix_positions, + helix_slots=helix_slots, + helix_bounds=helix_bounds, + helix_zero_kv=helix_zero_kv, + num_helix=num_helix, + tokens_per_block=tokens_per_block, + cp_size=cp_size, + cp_rank=cp_rank, + kv_bound_start=kv_bound_start, + tokens_per_seq=tokens_per_seq, + ) + + def reference( + self, + position_ids: torch.Tensor, + pos_offsets: torch.Tensor, + num_pos: int, + kv_lens: torch.Tensor, + kv_offsets: torch.Tensor, + kv_start: int, + num_kv: int, + *, + has_overlap: bool, + helix_positions: Optional[torch.Tensor] = None, + helix_slots: Optional[torch.Tensor] = None, + helix_bounds: Optional[torch.Tensor] = None, + helix_zero_kv: Optional[torch.Tensor] = None, + num_helix: int = 0, + tokens_per_block: int = 1, + cp_size: int = 1, + cp_rank: int = 0, + kv_bound_start: int = 0, + tokens_per_seq: int = 1, + ) -> dict: + """The written tensors after the call, by argument name (clones; only those the call + writes).""" + pos, kv = position_ids.clone(), kv_lens.clone() + out = dict(position_ids=pos, kv_lens=kv) + if has_overlap: + pos[:num_pos] += pos_offsets[:num_pos] + kv[kv_start : kv_start + num_kv] += kv_offsets[:num_kv] + if num_helix > 0: + n = num_helix + hpos = helix_positions.clone() + if has_overlap: + hpos[:n] += pos_offsets[:n] + g = hpos[:n].to(torch.int64) + before = _helix_local_len(g, tokens_per_block, cp_size, cp_rank) + bound = _helix_local_len(g + 1, tokens_per_block, cp_size, cp_rank) + owned = torch.div(g, tokens_per_block, rounding_mode="floor") % cp_size == cp_rank + slots, bounds, zero = helix_slots.clone(), helix_bounds.clone(), helix_zero_kv.clone() + slots[:n] = torch.where(owned, before, -1).to(slots.dtype) + bounds[:n] = bound.to(bounds.dtype) + zero[:n] = bound == 0 + num_seqs = n // tokens_per_seq + last = bounds[tokens_per_seq - 1 : num_seqs * tokens_per_seq : tokens_per_seq] + kv[kv_bound_start : kv_bound_start + num_seqs] = last + out.update( + helix_positions=hpos, helix_slots=slots, helix_bounds=bounds, helix_zero_kv=zero + ) + return out + + def is_valid( + self, + position_ids: torch.Tensor, + pos_offsets: torch.Tensor, + num_pos: int, + kv_lens: torch.Tensor, + kv_offsets: torch.Tensor, + kv_start: int, + num_kv: int, + **kwargs, + ) -> None: + # The launcher checks dtypes and strides; the ranges it reads are not checked. + assert position_ids.numel() >= num_pos and pos_offsets.numel() >= num_pos + assert kv_lens.numel() >= kv_start + num_kv and kv_offsets.numel() >= num_kv + num_helix = kwargs.get("num_helix", 0) + if num_helix > 0: + assert 0 <= kwargs.get("cp_rank", 0) < kwargs.get("cp_size", 1), "cp_rank" + num_seqs = num_helix // kwargs.get("tokens_per_seq", 1) + assert kv_lens.numel() >= kwargs.get("kv_bound_start", 0) + num_seqs + assert pos_offsets.numel() >= num_helix or not kwargs["has_overlap"] + + +k3_spec_decode_head = K3SpecDecodeHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_splitk_dense.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_splitk_dense.py new file mode 100644 index 000000000000..726039b727dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_splitk_dense.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M BF16 / FP16 GEMM with an in-kernel cluster split-K reduction (CuTe DSL), at the +tactics the DSpark drafter pins per projection.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.low_m_bf16_splitk import SplitKTactic, run_splitk_dense + + +class K3SplitkDense(OpWrapper): + """``out = a [M, K] @ b [K, N] (+ bias [N])`` in ``out``'s dtype (BF16 / FP16), FP32 + accumulation: each cluster rank sums an exact K slice, rank 0 reduces the partials through + DSMEM, casts and stores once. Returns ``out``. + + ``b`` is column-major (``weight.t()`` of a contiguous ``[N, K]`` weight); ``a`` row-major, + or with ``strided_a`` a row-major view with a 16-byte-aligned row stride. All 32-byte + aligned; ``tactic`` is ``(mma_m, mma_n, split_k, ab_stages)``. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The helix16 ULL drafter's projections at its block row counts (1 / 2 / 4 requests x 8), + # each with the tactic ``dspark_draft_gemm`` pins for that shape and row bound. + CELLS = ( + Cell( + why="helix16 ULL c4 [q_a | kv_a] 2112 x 7168 at 32 rows", + spec=dict(n=2112, k=7168, rows=32, tactic=(64, 32, 4, 8)), + ), + Cell( + why="helix16 ULL c1 [q_a | kv_a] 2112 x 7168 at 8 rows", + spec=dict(n=2112, k=7168, rows=8, tactic=(64, 8, 4, 12)), + ), + Cell( + why="helix16 ULL c2 context latents of the 5 layers, 2880 x 7168 at 16 rows", + spec=dict(n=2880, k=7168, rows=16, tactic=(64, 16, 4, 10)), + ), + Cell( + why="helix16 ULL c4 q_b, 64 heads x 192 = 12288 x 1536 at 32 rows", + spec=dict(n=12288, k=1536, rows=32, tactic=(64, 32, 1, 9)), + ), + Cell( + why="helix16 ULL c4 MLP gate_up shard of the TP16 group, 1792 x 7168 at 32 rows", + spec=dict(n=1792, k=7168, rows=32, tactic=(64, 32, 4, 8)), + ), + Cell( + why="helix16 ULL c2 MLP down shard, 7168 x 896 at 16 rows", + spec=dict(n=7168, k=896, rows=16, tactic=(64, 16, 1, 11)), + ), + Cell( + why="helix16 ULL c1 draft LM head shard, 10240 x 7168 at 8 rows", + spec=dict(n=10240, k=7168, rows=8, tactic=(128, 8, 2, 6)), + ), + ) + + note = """ + ``dspark_draft_gemm.draft_gemm`` pins it for the drafter's projections at up to 32 rows (the + tactics per (N, K) and row bound in ``_PINS`` / ``_PINS_SM107``). Changes the drafter's + accumulation order only. The launcher validates tensors and tactic and raises ValueError. A + process-wide compile cache holds one kernel per (dtype, tactic, PDL, bias, leading dims); + layouts are dynamic. + """ + + def __call__( + self, + a: torch.Tensor, + b: torch.Tensor, + bias: Optional[torch.Tensor], + out: torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + strided_a: bool = False, + ) -> torch.Tensor: + return run_splitk_dense(a, b, bias, out, pdl, tactic, strided_a=strided_a) + + def reference( + self, + a: torch.Tensor, + b: torch.Tensor, + bias: Optional[torch.Tensor], + out: torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + strided_a: bool = False, + ) -> torch.Tensor: + """``out`` after the call (new): FP64 product, rounded once to ``out``'s dtype.""" + y = a.double() @ b.double() + if bias is not None: + y = y + bias.double() + return y.to(out.dtype) + + +k3_splitk_dense = K3SplitkDense() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_store_context_latents.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_store_context_latents.py new file mode 100644 index 000000000000..cffbb2ebb95e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_store_context_latents.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Masked store of a step's context latents of every drafter layer into the paged pools +(Triton, one program per token and layer).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import store_context_latents, store_mla_context_latents + + +def to_e4m3_saturated(x: torch.Tensor) -> torch.Tensor: + """``to_draft_fp8``: E4M3 at unit scale, saturated to +-448, round to nearest even.""" + return x.float().clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + + +class K3StoreContextLatents(OpWrapper): + """``pools[l][pages[t], 0, 0, offsets[t]] = latent[t, l] * mask[t]`` for every token ``t`` + and layer ``l``. Returns None. + + ``latent`` is ``[N, L, D]`` with a unit inner stride; ``pools`` the ``L`` per-layer + ``[pages, 1, 1, page_size, D]`` views (identical strides), ``pool_ptrs`` their ``[L]`` int64 + base addresses, ``pages`` / ``offsets`` ``[N]`` and ``mask`` ``[N]`` int8. ``mask`` 0 stores a + zero (the signed zero of ``x * 0``). With ``mla`` (the MLA drafter's store), ``mask`` -1 leaves + the slot untouched and an FP8 E4M3 pool stores ``to_draft_fp8`` of the masked latent; + without it the pools have ``latent``'s dtype. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="dep16 (c383): 16 requests x 6 latents of 5 layers into FP8 pools, 0 / 1 masks", + spec=dict(num_gens=16, width=6, cp_size=1, pool="fp8", mla=True), + ), + Cell( + why="helix8 c75/c174: 8 requests x 8 latents, -1 on the other 7 ranks' positions", + spec=dict(num_gens=8, width=8, cp_size=8, pool="fp8", mla=True), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4): 4 requests x 8 latents over 16 ranks", + spec=dict(num_gens=4, width=8, cp_size=16, pool="fp8", mla=True), + ), + Cell( + why="the unfused preparation's store into BF16 pools (no shipped profile takes it)", + spec=dict(num_gens=16, width=6, cp_size=1, pool="bf16", mla=False), + ), + ) + + note = """ + The multiply by the mask runs in FP32 and rounds back to the latent dtype, which is exact for + 0 and 1. The MLA store is the one the tail-fused generation step runs (possibly on the + drafter's side stream); the plain one belongs to the unfused preparation, which no shipped + K3 generation profile takes. + """ + + def __call__( + self, + latent: torch.Tensor, + pools: list, + pool_ptrs: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + mask: torch.Tensor, + mla: bool = True, + ) -> None: + if mla: + store_mla_context_latents(latent, pools, pool_ptrs, pages, offsets, mask) + else: + store_context_latents(latent, pools, pool_ptrs, pages, offsets, mask) + + def reference( + self, + latent: torch.Tensor, + pools: list, + pool_ptrs: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + mask: torch.Tensor, + mla: bool = True, + ) -> list[torch.Tensor]: + """The ``L`` pools after the call (clones); ``pool_ptrs`` is not read.""" + out = [] + pages_i, offsets_i = pages.to(torch.int64), offsets.to(torch.int64) + write = mask >= 0 if mla else torch.ones_like(mask, dtype=torch.bool) + keep = (mask > 0).to(torch.float32)[:, None] + for layer, pool in enumerate(pools): + new = pool.clone() + vals = (latent[:, layer].float() * keep).to(latent.dtype) + target = new[:, 0, 0] + if new.dtype == torch.float8_e4m3fn: + # Assigned as bytes: the bits are the same. + vals = to_e4m3_saturated(vals).view(torch.uint8) + target = target.view(torch.uint8) + target[pages_i[write], offsets_i[write]] = vals[write].to(target.dtype) + out.append(new) + return out + + def is_valid( + self, + latent: torch.Tensor, + pools: list, + pool_ptrs: torch.Tensor, + pages: torch.Tensor, + offsets: torch.Tensor, + mask: torch.Tensor, + mla: bool = True, + ) -> None: + # The kernel stores through pool_ptrs, not through the views it is handed. + expected = [p.data_ptr() for p in pools] + assert pool_ptrs.tolist() == expected, "pool_ptrs must hold the pools' base addresses" + assert mask.dtype == torch.int8, "mask must be int8" + if not mla: + # The plain kernel multiplies by the mask: -1 would store a negated latent. + assert bool((mask >= 0).all()), "the plain store takes a 0 / 1 mask" + + +k3_store_context_latents = K3StoreContextLatents() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_local_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_local_select.py new file mode 100644 index 000000000000..37b15ec8adc7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_local_select.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-rank half of vocab-parallel token selection: each split of this rank's logits shard +reduced to a ``(value, global index)`` pair in the exchange buffer (Triton).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.vocab_parallel_kernels import _PairLayout, local_select +from ._reference import gumbel, philox4x32, publish_pairs + + +def sampled_scores(local: torch.Tensor, col_base: int, sampling: tuple) -> torch.Tensor: + """FP64 ``[rows, shard]`` scores whose argmax is the draw: the logits for rows at or below + the greedy threshold, else ``logits / T`` plus the Philox Gumbel noise of the global column.""" + temperatures, seeds, offsets, threshold = sampling + rows, n_cols = local.shape + device = local.device + t = temperatures[:rows].double() + greedy = t <= threshold + inv_t = 1.0 / torch.where(greedy, torch.ones_like(t), t) + cols = col_base + torch.arange(n_cols, device=device, dtype=torch.int64) + seed = seeds[:rows].to(torch.int64)[:, None] + offset = offsets[:rows].to(torch.int64)[:, None] + zeros = torch.zeros((rows, n_cols), dtype=torch.int64, device=device) + c0 = zeros + (cols // 4)[None, :] + c1 = zeros + torch.arange(rows, device=device)[:, None] + c2 = zeros + (offset & 0xFFFFFFFF) + c3 = zeros + ((offset >> 32) & 0xFFFFFFFF) + words = torch.stack(philox4x32(seed, c0, c1, c2, c3), dim=-1) + noise = gumbel(words.gather(-1, (cols % 4).expand(rows, -1).unsqueeze(-1)).squeeze(-1)) + x = local.double() + return torch.where(greedy[:, None], x, x * inv_t[:, None] + noise) + + +class K3VocabLocalSelect(OpWrapper): + """Publish, for every row of ``local`` ``[rows, shard]`` (unit inner stride) and every split + ``s`` of ``layout`` (columns ``[s * chunk, (s + 1) * chunk)``), the split's winner as float32 + ``(value, col_base + column)`` at columns ``2 (slot0 + s)`` and ``2 (slot0 + s) + 1`` of the + row of ``layout.buffer``; with ``zero_width > 0`` the split-0 program also zeroes the row's + columns outside ``[zero_lo, zero_hi)``, so that an all-reduce sum gathers. Returns None. + + The winner is the ``torch.argmax`` of the FP32 logits, or with ``sampling = (temperatures, + seeds, offsets, greedy_threshold)`` the Gumbel-max draw ``argmax(logits / T + G)`` for rows + with ``T > greedy_threshold``. The noise of global column ``v`` of row ``r`` is + ``-log(-log(u))`` of output ``v % 4`` of ``Philox4x32-10(seed[r], (v // 4, r, offset[r] + lo, offset[r] hi))``; every rank draws the same noise for a column. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # One rank of the vocab-sharded drafter / target head (163840 columns over the helix group) + # with the MNNVL exchange's layout; rows are requests x (draft length + 1). + CELLS = ( + Cell( + why="helix16 ULL c1, rank 0 of 16: 8 rows over a 10240-column shard, greedy, 5 splits", + spec=dict(tp=16, rank=0, rows=8, sampling=False, ties=False), + ), + Cell( + why="helix16 ULL c4, last rank: 32 rows, greedy", + spec=dict(tp=16, rank=15, rows=32, sampling=False, ties=False), + ), + Cell( + why="helix16 ULL c4, rank 5: 32 rows of Gumbel-max sampling, some rows greedy", + spec=dict(tp=16, rank=5, rows=32, sampling=True, ties=False), + ), + Cell( + why="helix8 c24/c42, rank 0 of 8: 16 rows over a 20480-column shard, 10 splits", + spec=dict(tp=8, rank=0, rows=16, sampling=False, ties=False), + ), + Cell( + why="helix8 c75/c174, last rank: 64 rows, 4 splits of 6144 columns, the last partial", + spec=dict(tp=8, rank=7, rows=64, sampling=False, ties=False), + ), + Cell( + why="helix8 c75/c174, rank 2: maxima tied across splits and NaN rows (argmax order)", + spec=dict(tp=8, rank=2, rows=64, sampling=False, ties=True), + ), + Cell( + why="helix8 c240, rank 3: 120 rows (20 x 6) of sampling, 3 splits of 8192", + spec=dict(tp=8, rank=3, rows=120, sampling=True, ties=False), + ), + ) + + note = """ + ``layout`` is a ``_PairLayout`` of ``_PairExchange.layout``: its buffer is new each call, and + the exchange (MNNVL all-reduce, NCCL all-gather, or none on one rank) follows. Grid (rows, + splits), 8 warps, 2048 columns per iteration. ``col_base`` and ``chunk`` are multiples of 4, + so Philox groups never straddle splits. The draws follow ``softmax(logits / T)``; they are + not the gathered sampler's draws. + """ + + def __call__( + self, + local: torch.Tensor, + col_base: int, + layout: _PairLayout, + sampling: Optional[tuple] = None, + ) -> None: + local_select(local, col_base, layout, sampling) + + def reference( + self, + local: torch.Tensor, + col_base: int, + layout: _PairLayout, + sampling: Optional[tuple] = None, + ) -> torch.Tensor: + """``layout.buffer`` after the call (a clone). Sampled values are computed in FP64.""" + if sampling is None: + return publish_pairs(layout, local.float(), col_base) + return publish_pairs(layout, sampled_scores(local, col_base, sampling), col_base) + + def is_valid( + self, + local: torch.Tensor, + col_base: int, + layout: _PairLayout, + sampling: Optional[tuple] = None, + ) -> None: + assert local.dim() == 2 and local.stride(-1) == 1, "local must be [rows, shard], unit inner" + assert layout.buffer.dtype == torch.float32, "pairs are float32" + assert layout.chunk % 4 == 0 and col_base % 4 == 0, "Philox groups straddle splits" + assert layout.splits * layout.chunk >= local.shape[1], "the splits miss columns" + # Global indices travel as float32: exact only below 2**24. + assert col_base + local.shape[1] <= 2**24, "vocab too large for float32 indices" + + +k3_vocab_local_select = K3VocabLocalSelect() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_reduce_pairs.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_reduce_pairs.py new file mode 100644 index 000000000000..8d3286f94ac4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_reduce_pairs.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Second half of vocab-parallel token selection: the exchanged ``(value, global index)`` pairs +of every rank reduced to one token per row (Triton).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.vocab_parallel_kernels import reduce_pairs +from ._reference import argmax_pairs, interleaved_pairs + + +class K3VocabReducePairs(OpWrapper): + """``out[r]`` = the global index of the pair the ``torch.argmax`` order selects among slots + ``s < n_slots`` of row ``r`` (value at ``pairs.view(-1)[r * stride_row + s * stride_slot]``, + index right after it). ``out`` is ``[rows]`` of any integer dtype; returns ``out``.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The MNNVL exchange of the whole helix group (every rank's per-split pairs, interleaved); + # int64 for the target tokens, int32 for the draft tokens of a drafter without a Markov head. + CELLS = ( + Cell( + why="helix16 ULL c1: 8 rows, 16 ranks x 5 splits = 80 slots", + spec=dict(tp=16, rows=8, out_dtype="int64", ties=False), + ), + Cell( + why="helix16 ULL c4: 32 rows, 80 slots, int32 draft tokens", + spec=dict(tp=16, rows=32, out_dtype="int32", ties=False), + ), + Cell( + why="helix8 c24/c42: 16 rows, 8 ranks x 10 splits = 80 slots", + spec=dict(tp=8, rows=16, out_dtype="int64", ties=False), + ), + Cell( + why="helix8 c75/c174: 64 rows (two programs), ties across ranks and NaN rows", + spec=dict(tp=8, rows=64, out_dtype="int64", ties=True), + ), + Cell( + why="helix8 c240: 120 rows (20 x 6, a partial program), 8 ranks x 3 splits", + spec=dict(tp=8, rows=120, out_dtype="int64", ties=False), + ), + ) + + note = """ + The pairs come from ``k3_vocab_local_select`` through the exchange: MNNVL / one rank lay + them out interleaved per row (slot stride 2), NCCL as rank blocks (slot stride 2 * rows). + A NaN wins; equal values go to the lower index, which is the lower rank's lower column, so + the result equals ``torch.argmax`` over the gathered logits. 32 rows per program. + """ + + def __call__( + self, + pairs: torch.Tensor, + stride_row: int, + stride_slot: int, + n_slots: int, + out: torch.Tensor, + ) -> torch.Tensor: + return reduce_pairs(pairs, stride_row, stride_slot, n_slots, out) + + def reference( + self, + pairs: torch.Tensor, + stride_row: int, + stride_slot: int, + n_slots: int, + out: torch.Tensor, + ) -> torch.Tensor: + """The new ``out`` (a new tensor of its shape and dtype); ``out`` is not written.""" + vals, idx = interleaved_pairs(pairs, stride_row, stride_slot, n_slots, out.shape[0]) + _, best = argmax_pairs(vals, idx) + return best.to(out.dtype) + + +k3_vocab_reduce_pairs = K3VocabReducePairs() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_shard_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_shard_embedding.py new file mode 100644 index 000000000000..1ba7e2e1be58 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/k3_vocab_shard_embedding.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This rank's part of a vocab-sharded embedding lookup (Triton), before the all-reduce.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.spec.spec_tail_kernels import vocab_shard_lookup + + +class K3VocabShardEmbedding(OpWrapper): + """``out[t] = weight[ids[t] - vocab_start]`` where ``vocab_start <= ids[t] < vocab_end``, + zero elsewhere: ``[len(ids), H]`` in ``weight``'s dtype, new. ``weight`` is this rank's + contiguous ``[shard, H]`` rows, ``ids`` 1-D of any integer dtype. The all-reduce (sum) of the + group's results is the full-vocab lookup; the caller runs it.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + + CELLS = ( + Cell( + why="helix8 c75/c174: 8 bonus ids and the mask id on rank 3 of 8 (20480-row shard)", + spec=dict(tp=8, rank=3, num_ids=9), + ), + Cell( + why="helix8 c24/c42: the last of 8 ranks, which holds the mask id", + spec=dict(tp=8, rank=7, num_ids=3), + ), + Cell( + why="helix16_moetp4ep4 (ULL c4): rank 0 of 16 (10240-row shard)", + spec=dict(tp=16, rank=0, num_ids=5), + ), + ) + + note = """ + The DFlash generation step's lookup of ``[bonus..., mask]`` in a vocab-parallel drafter + embedding where the owner-multicast site (``k3_draft_embed``) does not apply. Grid + (len(ids), cdiv(H, 1024)); rows outside the shard are stored as zeros, not skipped. + """ + + def __call__( + self, ids: torch.Tensor, weight: torch.Tensor, vocab_start: int, vocab_end: int + ) -> torch.Tensor: + return vocab_shard_lookup(ids, weight, vocab_start, vocab_end) + + def reference( + self, ids: torch.Tensor, weight: torch.Tensor, vocab_start: int, vocab_end: int + ) -> torch.Tensor: + ids = ids.to(torch.int64) + local = (ids >= vocab_start) & (ids < vocab_end) + rows = weight[torch.where(local, ids - vocab_start, 0)] + return torch.where(local[:, None], rows, torch.zeros_like(rows)) + + def is_valid( + self, ids: torch.Tensor, weight: torch.Tensor, vocab_start: int, vocab_end: int + ) -> None: + # The kernel addresses rows as weight + row * H. + assert weight.is_contiguous(), "weight must be contiguous" + assert vocab_end - vocab_start <= weight.shape[0], "the shard is wider than the table" + + +k3_vocab_shard_embedding = K3VocabShardEmbedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/_fla_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/_fla_index.py new file mode 100644 index 000000000000..0338937302b2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/_fla_index.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Adapted from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/utils/index.py +# and https://github.com/fla-org/flash-linear-attention/blob/main/fla/utils.py +# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/fla/index.py +"""FLA's varlen chunk indices for the KDA prefill: ``(sequence, chunk)`` of every chunk.""" + +import functools +from typing import Any, Callable, Dict, Optional, Tuple + +import torch +import triton + + +def tensor_cache(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + """ + A decorator that caches the most recent results of a function with tensor inputs. + This decorator will store the output of the decorated function for the most recent set of input tensors. + The cache is limited to a fixed size (default is 4). When the cache is full, the oldest entry will be removed. + Args: + fn (Callable[..., torch.Tensor]): + The function to be decorated. It should take tensor inputs and return tensor outputs. + Returns: + Callable[..., torch.Tensor]: + A wrapped version of the input function with single-entry caching. + """ + + cache_entries: Tuple[Optional[Tuple], Optional[Dict], Any] = [] + cache_size = 4 + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal cache_entries, cache_size + for i, entry in enumerate(cache_entries): + last_args, last_kwargs, last_result = entry + if len(args) == len(last_args) and len(kwargs) == len(last_kwargs): + if all(a is b for a, b in zip(args, last_args)) and all( + k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items() + ): + cache_entries = ( + cache_entries[:i] + cache_entries[i + 1 :] + [(args, kwargs, last_result)] + ) + return last_result + + result = fn(*args, **kwargs) + + if len(cache_entries) >= cache_size: + cache_entries = cache_entries[1:] + cache_entries.append((args, kwargs, result)) + return result + + return wrapper + + +@tensor_cache +def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + return cu_seqlens[1:] - cu_seqlens[:-1] + + +@tensor_cache +def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: + indices = torch.cat( + [torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()] + ) + return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/_kda_prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/_kda_prefill.py new file mode 100644 index 000000000000..cb34b2bb55e1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/_kda_prefill.py @@ -0,0 +1,1382 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL host launcher of the Kimi K3 KDA prefill. + +The implementation source-integrates the production orchestration from the +KDA ``chunk_kda_fwd`` benchmark. The default pipeline launches fused K123, +Akk inverse, and persistent K4 kernels. + +Supported: + - Equal-length sequences (B ≥ 1, with the existing B=1 padding path) + - Variable-length sequences (B=1 with cu_seqlens) + - safe_gate mode (sigmoid + lower_bound) and softplus mode + - dt_bias + - use_gate_in_kernel=True with A_log + - optional in-kernel beta sigmoid + - chunk_size=64 + +``kda_prefill`` updates selected recurrent-state pool rows in place. Intermediate matrices +remain private runner workspace. It runs the private kernels in ``kernels/kda`` and is not a +registered torch op (main's ``custom_ops`` registers a different ``trtllm::kda_prefill``, +FlashInfer BT16 on raw Q/K/V); the model calls it through the ``ops/ssm/kda_prefill`` entry. +``kernels/kda/_kda_kernels.py`` imports that entry (and through it this module) on first use; importing +it raises ``ImportError`` without CuTe DSL or FlashInfer. +""" + +import time +import weakref +from typing import Optional + +import torch + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm.logger import logger + +if IS_CUTLASS_DSL_AVAILABLE and IS_FLASHINFER_AVAILABLE: + import cuda.bindings.driver as cuda_driver + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + from ...kernels.kda.akk_inverse import akk_inv_host as _akk_inv_host + from ...kernels.kda.fused_k123 import make_host_function as _fused_make_host + from ...kernels.kda.k4_persistent import BYTES_PER_TENSORMAP as _K4P_BTM + from ...kernels.kda.k4_persistent import NUM_TENSORMAPS as _K4P_NTM + from ...kernels.kda.k4_persistent import make_host_fn as _k4p_make_host + from ._fla_index import prepare_chunk_indices +else: + raise ImportError("Kimi K3 KDA prefill requires NVIDIA CUTLASS DSL and FlashInfer") + + +def _ct(t, etype): + """Create a CuTe tensor from PyTorch tensor.""" + r = from_dlpack(t, assumed_align=16) + r.element_type = etype + return r + + +def _fits_32bit_stride(tensor: torch.Tensor) -> bool: + int32_max = 2**31 - 1 + max_offset = int(tensor.storage_offset()) + for size, stride in zip(tensor.shape, tensor.stride()): + stride = abs(int(stride)) + if stride > int32_max: + return False + if size: + max_offset += (int(size) - 1) * stride + if max_offset > int32_max: + return False + return True + + +def _dynamic_dlpack_arg(tensor: torch.Tensor): + wrapper = from_dlpack( + tensor, + assumed_align=16, + use_32bit_stride=_fits_32bit_stride(tensor), + ) + for dim, stride in enumerate(tensor.stride()): + if stride == 1: + return wrapper.mark_layout_dynamic(dim) + return wrapper + + +def _dynamic_vector_layout_key(tensor: torch.Tensor): + """Key a 1D CuTe argument whose sole extent is runtime dynamic.""" + assert tensor.ndim == 1 + return ( + tensor.dtype, + tuple(tensor.stride()), + _fits_32bit_stride(tensor), + ) + + +def _dynamic_rows_layout_key(tensor: torch.Tensor): + """Key a ``_dynamic_dlpack_arg`` tensor whose leading (row) extent is runtime dynamic. + + The inner extents, all strides and the stride width stay in the key. + """ + return ( + tensor.dtype, + tuple(tensor.shape[1:]), + tuple(tensor.stride()), + _fits_32bit_stride(tensor), + ) + + +def _current_cu_stream(device): + """CUstream handle of torch's CURRENT stream, queried per call. + + The executor runs the model on a dedicated non-blocking + ``torch.cuda.Stream`` (``py_executor.execution_stream``), not the default + stream. Every kernel launch must go to this stream: launching on the DSL + default stream (what ``.launch()`` does without a ``stream`` argument) + races with the projections that produce q/k/v/g/beta and with the + consumers of O/final_state — silent, intermittent corruption in the + runtime while single-stream unit tests pass. This is also why the value + must never be cached alongside the scratch buffers: the buffers may be + created under a different stream (e.g. warmup) than later forwards. + """ + return cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + + +# Cached eqlen dummy cu/ci cute wrappers — shared by K123 and akk_inv. +# Avoids per-call torch.empty + from_dlpack overhead (~10-12us each). +_eqlen_dummy_cache = {} + + +def _get_eqlen_dummies(device, idx_dtype=torch.int64): + """Returns cached (cu_ct, ci_ct) cute wrappers for eqlen (B+1=2, NT+1=2).""" + key = (device.index if device.index is not None else 0, idx_dtype) + if key not in _eqlen_dummy_cache: + cu_t = torch.empty(2, dtype=idx_dtype, device=device) + ci_t = torch.empty(1, 2, dtype=idx_dtype, device=device) + cu_etype = cutlass.Int64 if idx_dtype == torch.int64 else cutlass.Int32 + _eqlen_dummy_cache[key] = (_ct(cu_t, cu_etype), _ct(ci_t, cu_etype)) + return _eqlen_dummy_cache[key] + + +def _cute_int_type(dtype): + """Map PyTorch integer dtype to CUTLASS element type.""" + if dtype == torch.int32: + return cutlass.Int32 + elif dtype == torch.int64: + return cutlass.Int64 + else: + raise ValueError(f"Unsupported integer dtype: {dtype}") + + +# ========== Fused K1+K2+K3 compilation cache ========== +_fused_k123_cache = {} + + +def _kda_persistent_grid_blocks(device: torch.device) -> int: + """Persistent-grid size for the KDA prefill kernels on ``device``. + + One resident block per SM, minus the SMs ``cute_dsl_utils.sm_reserve()`` + keeps free for a kernel resident for the whole step (the DWDP weight + gather on context workers). Launching more blocks than free SMs would + leave the surplus blocks queued behind the whole first wave. + """ + from tensorrt_llm._torch.cute_dsl_utils import sm_reserve + + num_sms = torch.cuda.get_device_properties(device).multi_processor_count + return max(1, num_sms - sm_reserve()) + + +# id(cu_seqlens) -> bool. Skips per-call GPU->CPU sync on subsequent calls +# when the same cu_seqlens tensor is reused (typical training/inference loop). +_varlen_pure_cache = {} +# id(cu_seqlens) -> int seqlen, populated alongside _varlen_pure_cache for +# single-seq cu_seqlens. +_varlen_single_seqlen_cache = {} + +# ===== Caching rule for this module (read before adding a cache) ===== +# +# A cute wrapper built by ``from_dlpack`` holds a STRONG reference to the +# torch tensor it was built from. So a module-level dict that stores a +# wrapper keeps that tensor -- and its GPU storage -- alive for the life of +# the process. id()-keyed weakref pruning cannot save you here: the cache +# entry is itself what keeps the key object alive, so the finalizer never +# runs. Such an entry is immortal by construction. +# +# Therefore: +# * Never put a wrapper over a per-call activation in a module-level dict +# (see the note in _launch_k4_persistent). +# * Never put a wrapper over LRU-managed scratch in a module-level dict +# either -- it defeats the eviction. Store it next to the buffers whose +# lifetime it must share (see ``cute_wrappers`` in _get_buffers). +# * A wrapper over a private COPY the module intends to own forever is +# fine (see _get_dt_bias_ct / _get_k4_varlen_cu_co, which clone on +# purpose so the cached value never aliases the keyed object). +# +# Every shape-keyed dict below must also be LRU-bounded via _lru_put: T is +# the packed context-token count of one iteration (attn_metadata +# .num_ctx_tokens), which the scheduler does not quantize, so an unbounded +# shape-keyed cache grows without limit under any real serving load. + + +def _lru_put(cache, key, value, max_entries): + """Insert ``value`` into ``cache``, evicting in insertion order first.""" + while len(cache) >= max_entries: + cache.pop(next(iter(cache))) + cache[key] = value + return value + + +def _lru_touch(cache, key): + """Move ``key`` to the MRU end so hot shapes survive eviction.""" + cache[key] = cache.pop(key) + return cache[key] + + +def _prune_on_gc(cache, key, *keyobjs): + """Drop ``cache[key]`` when any of ``keyobjs`` is garbage-collected. + + The id()-keyed caches in this module are only sound while the keyed + Python objects stay alive — CPython reuses an object's address (its id) + the moment it is freed, so a dead entry could otherwise be served for an + unrelated new tensor with recycled id but different contents/storage. + A finalizer runs at dealloc, before the address can be recycled, so a + pruned entry can never alias a new object. This also bounds cache growth + when callers pass fresh tensors every call. + """ + for o in keyobjs: + weakref.finalize(o, cache.pop, key, None) + + +# Cache for dt_bias `.float().contiguous().view(H, K)` + cute wrapper. +# dt_bias is typically a nn.Parameter — same object across iterations. +_dt_bias_cache = {} + + +def _get_dt_bias_ct(dt_bias, H, K): + """Returns cached cute wrapper for dt_bias.float().view(H, K).""" + key = (id(dt_bias), H, K) + entry = _dt_bias_cache.get(key) + if entry is None: + # Copy (never view) the keyed object: a view would pin it, making + # the entry immortal (weakref pruning could never fire). Callers + # pass a fresh .detach() per call, so an aliasing entry would leak. + bias_t = dt_bias.detach().float().reshape(H, K).clone() + entry = (bias_t, _ct(bias_t, cutlass.Float32)) + _dt_bias_cache[key] = entry + _prune_on_gc(_dt_bias_cache, key, dt_bias) + return entry[1] + + +# Cache for the empty 1x1 fp32 bias tensor used when dt_bias is None. +_empty_bias_cache = {} + + +def _get_empty_bias_ct(device): + idx = device.index if device.index is not None else 0 + if idx not in _empty_bias_cache: + t = torch.empty(1, 1, dtype=torch.float32, device=device) + _empty_bias_cache[idx] = _ct(t, cutlass.Float32) + return _empty_bias_cache[idx] + + +# K4 varlen cu_seqlens / chunk_offsets cute wrappers. Caches the int32-cast +# tensor and its mark_layout_dynamic wrapper so they survive multiple calls +# with the same input objects. +_k4_varlen_cu_co_cache = {} + + +def _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets): + key = (id(cu_seqlens), id(chunk_offsets)) + entry = _k4_varlen_cu_co_cache.get(key) + if entry is None: + # Always materialize copies: if the cached value aliased the keyed + # object, the entry would pin it and the weakref pruning could + # never fire (immortal entry). + cu_int32 = cu_seqlens.to(torch.int32).contiguous() + if cu_int32 is cu_seqlens: + cu_int32 = cu_seqlens.clone() + co_int32 = chunk_offsets.to(torch.int32).contiguous() + if co_int32 is chunk_offsets: + co_int32 = chunk_offsets.clone() + cu_ct = from_dlpack(cu_int32, assumed_align=4).mark_layout_dynamic() + cu_ct.element_type = cutlass.Int32 + co_ct = from_dlpack(co_int32, assumed_align=4).mark_layout_dynamic() + co_ct.element_type = cutlass.Int32 + # Hold refs to the int32 tensors so they don't get GC'd and the + # underlying storage remain valid as long as cu_ct / co_ct live. + entry = (cu_int32, co_int32, cu_ct, co_ct) + _k4_varlen_cu_co_cache[key] = entry + _prune_on_gc(_k4_varlen_cu_co_cache, key, cu_seqlens, chunk_offsets) + return entry[2], entry[3] + + +# Cache the (cu_for_k4, chunk_offsets_for_k4) pair keyed by id(cu_seqlens). +# Both tensors are computed on-GPU with no host sync — replaces the previous +# `cu_seqlens.cpu().tolist() + Python loop + torch.tensor` chain that forced +# a 50-200us GPU->CPU stall on the K4 prep path every varlen call. +_varlen_k4_input_cache = {} + + +def _get_varlen_k4_inputs(cu_seqlens, BT): + key = id(cu_seqlens) + entry = _varlen_k4_input_cache.get(key) + if entry is None: + # Copy (never alias) the keyed object — see _get_k4_varlen_cu_co. + cu_int32 = cu_seqlens.to(torch.int32).contiguous() + if cu_int32 is cu_seqlens: + cu_int32 = cu_seqlens.clone() + seq_lens = cu_int32[1:] - cu_int32[:-1] + chunk_counts = (seq_lens + (BT - 1)) // BT + zero = torch.zeros(1, dtype=torch.int32, device=cu_int32.device) + co_int32 = torch.cat([zero, torch.cumsum(chunk_counts, dim=0).to(torch.int32)]) + co_int32 = co_int32.contiguous() + _varlen_k4_input_cache[key] = (cu_int32, co_int32) + _prune_on_gc(_varlen_k4_input_cache, key, cu_seqlens) + return _varlen_k4_input_cache[key] + + +# ========== BF16 akk_inv compilation cache ========== +_akk_inv_cache = {} +# ========== K4 persistent (varlen via cu_seqlens) compilation cache ========== +_k4p_cache = {} +_k4p_tm_ws = {} + +# Buffer cache: avoid re-allocating ~67us of intermediate tensors per call. +# Also caches cute.Tensor wrappers (saves ~7us each call from from_dlpack). +# +# Sized in bytes, not in "entries": at K3's H=96, K=V=128 one entry owns +# ~0.12 MiB per token, i.e. ~1 GiB at max_num_tokens=8192. This scratch is +# allocated AFTER the KV pool has been sized from the warmup peak (warmup +# compiles a single 257-token context), so every resident entry is memory +# the KV-cache accounting never saw -- keep the cap small. +# +# A small cap costs almost nothing: T is the packed context-token count of +# ONE iteration, shared by all 69 KDA layers, so a 1-entry cache already +# hits 68/69 within a forward. Cross-iteration hits are incidental -- +# T = attn_metadata.num_ctx_tokens is not quantized by the scheduler +# (only budget-trimmed chunks get rounded to a 64-multiple), so the shape +# usually changes from one iteration to the next. +_buf_cache = {} +_BUF_CACHE_MAX_ENTRIES = 2 + +# Padded-input scratch cache for the eqlen partial-chunk path. Keyed by +# (B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T). +# real_T is part of the key so the g sentinel tail [real_T:T_padded] = -1e3 +# is set once and reused across calls with the same shape. +_padded_input_cache = {} +_PADDED_INPUT_CACHE_MAX_ENTRIES = _BUF_CACHE_MAX_ENTRIES + +# Sentinel-padded g scratch for varlen single-seq Phase 2.1 path. Keyed by +# (B, T_padded, H, K, dtype, device, real_T). The tail [real_T:T_padded] is +# pre-set to -1e3 once at cache init; subsequent calls only overwrite the +# valid prefix [0, real_T). +# +# real_T is the request's true prompt length, so the key space is as wide as +# the prompt-length distribution -- one ~200 MiB buffer per distinct length +# if left unbounded. +_g_sentinel_cache = {} +_G_SENTINEL_CACHE_MAX_ENTRIES = _BUF_CACHE_MAX_ENTRIES + + +def _get_g_sentinel_buffer(B, T_padded, H, K, dtype_g, device, real_T): + key = (B, T_padded, H, K, dtype_g, device.index if device.index is not None else 0, real_T) + if key in _g_sentinel_cache: + return _lru_touch(_g_sentinel_cache, key) + e = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device) + if real_T < T_padded: + e[:, real_T:] = -1000.0 + return _lru_put(_g_sentinel_cache, key, e, _G_SENTINEL_CACHE_MAX_ENTRIES) + + +def _get_padded_input_buffers(B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T): + key = ( + B, + T_padded, + H, + K, + dtype_qkv, + dtype_g, + dtype_beta, + device.index if device.index is not None else 0, + real_T, + ) + if key in _padded_input_cache: + return _lru_touch(_padded_input_cache, key) + q_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + k_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + v_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + beta_pad = torch.zeros(B, T_padded, H, dtype=dtype_beta, device=device) + # g: zero in [0, real_T), sentinel -1e3 in [real_T, T_padded). Caller's + # data overwrites the prefix each call; the sentinel tail never moves. + g_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device) + if real_T < T_padded: + g_pad[:, real_T:] = -1000.0 + e = (q_pad, k_pad, v_pad, g_pad, beta_pad) + return _lru_put(_padded_input_cache, key, e, _PADDED_INPUT_CACHE_MAX_ENTRIES) + + +def _get_buffers(dev, dtype_k, B, T, H, K_dim, V_dim, NT, N_seqs, BT, varlen=False): + """Return cached intermediate tensors for the KDA prefill pipeline.""" + key = (dev.index or 0, B, T, H, K_dim, V_dim, NT, N_seqs, varlen) + if key not in _buf_cache: + bf16 = cutlass.BFloat16 + fp32 = cutlass.Float32 + # Varlen chunk-tile kernels (akk_inv's Stage-0 cp.async A-tile load, + # K4's A-tile loads) transfer the full BT-row tile of every chunk + # and neutralize invalid rows only after the access; when the + # batch's FINAL chunk is partial they read up to BT-1 rows past the + # logical T. These are driver-owned scratch buffers, so honor the + # kernel's boundary-at-the-data contract by allocating one chunk of + # zeroed slack past T (the eqlen path gets the same guarantee from + # its 256-multiple input padding). The slack rows are never + # consumed — OOB rows are zeroed in SMEM or masked at stores. + # Input tensors (beta) get the opposite treatment: the K123 kernel + # bounds-checks those loads, since the driver doesn't own them. + if varlen: + assert B == 1, f"varlen expects packed B=1 input, got B={B}" + T_alloc = T + BT if varlen else T + + def _with_slack(ctor, *shape, dtype): + full = ctor(*shape, device=dev, dtype=dtype) + return full[:, :T] + + k_scaled = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k) # raw, no beta + kg = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k) + q_scaled = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k) + beta_activated = _with_slack(torch.empty, B, T_alloc, H, dtype=torch.bfloat16) + gk_last_exp = torch.empty(B, NT, H, K_dim, device=dev, dtype=torch.float32) + A_qk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) + A_kk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) + O_flat = _with_slack(torch.empty, B, T_alloc, H, V_dim, dtype=dtype_k) + cu_eqlen = torch.arange(0, (B + 1) * T, T, dtype=torch.int32, device=dev) + # co_eqlen is consumed only on the eqlen path (varlen K4 derives its + # chunk offsets from cu_seqlens). Eqlen inputs are pre-padded to a + # 256-multiple upstream, so T // BT >= 4 there; varlen calls can + # carry T < BT (short-prompt batches), where arange(step=T // BT) + # would raise "step must be nonzero" building this dead buffer. + # Clamp the step so the buffer stays constructible. + nt_eq = max(T // BT, 1) + co_eqlen = torch.arange(0, (B + 1) * nt_eq, nt_eq, dtype=torch.int32, device=dev) + + T_total = B * T + A_kk_flat = A_kk.reshape(T_total, H, BT) + A_qk_flat = A_qk.reshape(T_total, H, BT) + KS_flat = k_scaled.reshape(T_total, H, K_dim) + QS_flat = q_scaled.reshape(T_total, H, K_dim) + KG_flat = kg.reshape(T_total, H, K_dim) + O_token = O_flat.reshape(T_total, H, V_dim) + gk_flat = gk_last_exp.reshape(-1, H, K_dim) + + def _wrap(t, etype): + r = from_dlpack(t, assumed_align=16).mark_layout_dynamic() + r.element_type = etype + return r + + a_ct = _wrap(A_kk_flat, bf16) + aqc_ct = _wrap(A_qk_flat, bf16) + ks_ct = _wrap(KS_flat, bf16) # raw k_scaled (beta absorbed in akk_inv) + qs_ct = _wrap(QS_flat, bf16) + kg_ct = _wrap(KG_flat, bf16) + o_ct = _wrap(O_token, bf16) + gk_ct = _wrap(gk_flat, fp32) + + # K123 consumes the un-flattened [B, T, H, *] views of the same + # scratch (K4 consumes the token-flat ones above), so it needs its + # own wrappers. They MUST be built here and stored in the cache + # entry: a wrapper holds a strong reference to its torch tensor, so + # parking these in a module-level dict makes the whole buffer set + # immortal and the LRU below frees nothing. See the caching rule at + # the top of this file. + k123_ks_ct = _ct(k_scaled, bf16) + k123_kg_ct = _ct(kg, bf16) + k123_qs_ct = _ct(q_scaled, bf16) + k123_beta_activated_ct = _ct(beta_activated, bf16) + k123_gk_ct = _ct(gk_last_exp, fp32) + k123_aqk_ct = _ct(A_qk, bf16) + k123_akk_ct = _ct(A_kk, bf16) + cu_eqlen_ct = from_dlpack(cu_eqlen, assumed_align=4).mark_layout_dynamic() + cu_eqlen_ct.element_type = cutlass.Int32 + co_eqlen_ct = from_dlpack(co_eqlen, assumed_align=4).mark_layout_dynamic() + co_eqlen_ct.element_type = cutlass.Int32 + + # akk_inv views: bf16 storage reinterpreted as fp32 (packed 2x bf16 -> 1x fp32). + # Layout-dynamic so the single compiled akk_inv (shape-independent + # key) accepts them regardless of T; the host rebuilds its own + # runtime-shaped views from the iterator anyway. + akk_in_view = from_dlpack(A_kk, assumed_align=16).mark_layout_dynamic() + akk_in_view.element_type = fp32 + akk_out_view = from_dlpack(A_kk, assumed_align=16).mark_layout_dynamic() + akk_out_view.element_type = fp32 + + cute_wrappers = dict( + a_ct=a_ct, + aqc_ct=aqc_ct, + ks_ct=ks_ct, + qs_ct=qs_ct, + kg_ct=kg_ct, + o_ct=o_ct, + gk_ct=gk_ct, + cu_eqlen_ct=cu_eqlen_ct, + co_eqlen_ct=co_eqlen_ct, + akk_in_view=akk_in_view, + akk_out_view=akk_out_view, + k123_ks_ct=k123_ks_ct, + k123_kg_ct=k123_kg_ct, + k123_qs_ct=k123_qs_ct, + k123_beta_activated_ct=k123_beta_activated_ct, + k123_gk_ct=k123_gk_ct, + k123_aqk_ct=k123_aqk_ct, + k123_akk_ct=k123_akk_ct, + # Filled lazily on first launch — saves cache_key tuple build + + # outer dict lookup on subsequent calls. + _k123_fns={}, + _akk_inv_fn=None, + ) + + # Eviction is only real because nothing outside this tuple holds a + # reference to the buffers or their wrappers. + return _lru_put( + _buf_cache, + key, + ( + k_scaled, + kg, + q_scaled, + beta_activated, + gk_last_exp, + A_qk, + A_kk, + O_flat, + cu_eqlen, + co_eqlen, + cute_wrappers, + ), + _BUF_CACHE_MAX_ENTRIES, + ) + # LRU refresh so hot shapes survive eviction. + return _lru_touch(_buf_cache, key) + + +def _launch_k4_persistent( + cute_wrappers, + v_beta, + state_pool, + state_indices, + cu_seqlens, + chunk_offsets, + cu_eqlen_passed=False, + num_sm=148, + H=None, + V_dim=None, + snapshot_rows=None, + snapshot_chunks=None, +): + """Launch persistent K4 with cached CuTe wrappers. + + ``snapshot_rows`` / ``snapshot_chunks`` (int32 ``[N_seqs]``, both or + neither) select the K4 variant that also stores sequence ``b``'s state + after its first ``64 * snapshot_chunks[b]`` tokens into + ``state_pool[snapshot_rows[b]]`` (``-1``: no snapshot for ``b``). + + No fast-launch (args-tuple) cache here: such a cache pins the per-call + activation and state-pool tensors via their CuTe wrappers (the wrapper holds the + storage, so the keyed object never dies and weakref pruning never + fires) — a per-call GPU memory leak in the executor runtime. The + per-call cost is re-wrapping a handful of tensors (~10us each). + """ + bf16 = cutlass.BFloat16 + N_seqs = cu_seqlens.shape[0] - 1 + dev = v_beta.device + dev_idx = dev.index or 0 + + s_ct = _dynamic_dlpack_arg(state_pool) + a_ct = cute_wrappers["a_ct"] + b_ct = cute_wrappers["ks_ct"] # raw k_scaled (beta absorbed in akk_inv) + q_ct = cute_wrappers["qs_ct"] + aqc_ct = cute_wrappers["aqc_ct"] + kg_ct = cute_wrappers["kg_ct"] + o_ct = cute_wrappers["o_ct"] + gk_ct = cute_wrappers["gk_ct"] + + # v is a per-call activation — wrap fresh every call (never cache; see + # the caching rule at the top of this file). + v_view = v_beta.reshape(-1, H, V_dim) if v_beta.dim() == 4 else v_beta + v_ct = from_dlpack(v_view, assumed_align=16).mark_layout_dynamic() + v_ct.element_type = bf16 + + if cu_eqlen_passed: + cu_ct = cute_wrappers["cu_eqlen_ct"] + co_ct = cute_wrappers["co_eqlen_ct"] + else: + cu_ct, co_ct = _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets) + + state_indices_ct = _dynamic_dlpack_arg(state_indices) + has_snapshot = snapshot_rows is not None + snapshot_args = () + snapshot_key = None + if has_snapshot: + snapshot_args = ( + _dynamic_dlpack_arg(snapshot_rows), + _dynamic_dlpack_arg(snapshot_chunks), + ) + snapshot_key = ( + _dynamic_vector_layout_key(snapshot_rows), + _dynamic_vector_layout_key(snapshot_chunks), + ) + + # Allocate for the maximum persistent grid once. The scheduler applies + # min(N_seqs * H, num_sm) at runtime, so pre-minimizing here would turn + # every distinct request count into another compiled host function. + tm_key = (dev_idx, num_sm) + if tm_key not in _k4p_tm_ws: + tm_ws_t = torch.zeros( + num_sm * _K4P_NTM * _K4P_BTM, + dtype=torch.uint8, + device=dev, + ) + tm_ct = from_dlpack(tm_ws_t, assumed_align=16) + tm_ct.element_type = cutlass.Uint8 + _k4p_tm_ws[tm_key] = (tm_ws_t, tm_ct) + else: + tm_ws_t, tm_ct = _k4p_tm_ws[tm_key] + + # Launch on torch's CURRENT stream — see _current_cu_stream. + stream = _current_cu_stream(dev) + + # H and the state V/K dims are in the key, as are the pool's inner + # extents and strides: the compiled function is head-count- and + # layout-specific. The pool's row count is not: the kernel addresses + # rows through state_indices / snapshot_rows only, and a bf16 pool + # reaches here as a dense fp32 copy with one row per context sequence + # (plus snapshot rows), so keying on it would compile once per context + # batch size while serving. N_seqs and the token/chunk counts remain + # runtime values. + cache_key = ( + dev_idx, + num_sm, + H, + state_pool.shape[-1], + V_dim, + _dynamic_rows_layout_key(state_pool), + # The sole dimension of this vector is runtime dynamic. Excluding + # its extent lets warmup serve full and underfilled context batches. + _dynamic_vector_layout_key(state_indices), + ) + if has_snapshot: + # The snapshot variant is a different host function and kernel. + cache_key = cache_key + (("snapshot",) + snapshot_key,) + args = ( + a_ct, + b_ct, + v_ct, + q_ct, + aqc_ct, + kg_ct, + o_ct, + gk_ct, + s_ct, + state_indices_ct, + cu_ct, + co_ct, + *snapshot_args, + tm_ct, + N_seqs, + stream, + ) + k4_fn = _k4p_cache.get(cache_key) + if k4_fn is None: + host_fn = _k4p_make_host(num_sm=num_sm, has_snapshot=has_snapshot) + compile_start = time.perf_counter() + k4_fn = cute.compile(host_fn, *args) + _k4p_cache[cache_key] = k4_fn + # A compile blocks the forward pass that triggers it; after warmup, + # each of these lines is a serving stall of that length. + logger.info( + f"KDA prefill K4 compiled in {time.perf_counter() - compile_start:.2f} s " + f"(snapshot={has_snapshot}, state rows={state_pool.shape[0]}, " + f"sequences={N_seqs}, cache entries={len(_k4p_cache)})" + ) + k4_fn(*args) + + +def _launch_fused_k123_inv( + q, + k, + g, + A_log, + beta, + beta_activated, + scale, + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk_inv, + cu_seqlens, + chunk_indices, + is_varlen, + NT, + valid_tokens, + dt_bias=None, + safe_gate=False, + lower_bound=None, + use_beta_sigmoid_in_kernel=True, + akk_in_view=None, + akk_out_view=None, + cute_wrappers=None, + varlen_pure_override=False, + varlen_is_aligned=None, +): + """Persistent K1+K2 (writes A_kk in I+L format with diag=1) chained with + BF16 akk_inv (in-place inversion). Final A_kk_inv = (I+L)^-1.""" + + # No fast-launch (args-tuple) cache: it would pin the per-call q/k/g/beta + # activations via their cute wrappers (the wrapper holds the storage, so + # the keyed objects never die and weakref pruning never fires) — a + # per-call GPU memory leak in the executor runtime. Wrapper gathering + # below costs ~10us per tensor. + B, T, H, K = q.shape + BT = 64 + dev = q.device.index or 0 + T_padded = T if is_varlen else None + has_bias = dt_bias is not None + + # Prefer the host-derived alignment property supplied by the model runtime. + # Direct low-level callers that omit it retain the cached device inference + # fallback for backward compatibility. + varlen_pure = False + if is_varlen and cu_seqlens is not None: + if varlen_pure_override: + # Caller (Phase 2.1 single-seq path) sentinel-padded the data for + # this call; the cached per-object verdict stays untouched. + varlen_pure = True + elif varlen_is_aligned is not None: + varlen_pure = varlen_is_aligned + else: + _vp_key = id(cu_seqlens) + if _vp_key not in _varlen_pure_cache: + cu_cpu = cu_seqlens.cpu().tolist() + seq_lens = [cu_cpu[i + 1] - cu_cpu[i] for i in range(len(cu_cpu) - 1)] + _varlen_pure_cache[_vp_key] = all((sl % BT) == 0 for sl in seq_lens) + _prune_on_gc(_varlen_pure_cache, _vp_key, cu_seqlens) + varlen_pure = _varlen_pure_cache[_vp_key] + # B/NT/T are runtime args of the compiled host_fn (see make_host_function), + # avoiding shape-dependent grid and view specialization. Varlen remains + # fully shape-independent; eqlen retains T only as a cache discriminator + # for the raw tensor batch strides described below. + # + # cu/ci DTYPE must be part of the key: the kernel reads + # mCuSeqlens/mChunkIndices with the element type baked at compile time + # (int64 elements are addressed with stride 8, int32 with stride 4). + # Reusing an int64-compiled kernel on int32 tensors (or vice versa) + # misaddresses every cu/ci element — garbage seq ids/chunk starts -> + # OOB reads (cudaErrorIllegalAddress) or silent corruption. Observed as + # a crash when a dump-replay batch (int64 cu/ci) preceded a synthetic + # batch (int32 cu/ci) in one process. + cu_ci_dtypes = (cu_seqlens.dtype, chunk_indices.dtype) if is_varlen else None + # Eqlen kernels directly index mBeta/mAqk/mAkk using the input wrappers' + # compile-time batch stride, which depends on T. Varlen fixes B=1, so that + # batch stride is never used and all T shapes can keep sharing one compile. + eqlen_t_key = T if not is_varlen else -1 + grid_blocks = _kda_persistent_grid_blocks(q.device) + cache_key = ( + H, + is_varlen, + dev, + has_bias, + safe_gate, + varlen_pure, + cu_ci_dtypes, + eqlen_t_key, + use_beta_sigmoid_in_kernel, + beta.dtype, + grid_blocks, + ) + + # Inputs are guaranteed contiguous by upstream linear projections. + # A_log is fp32 model param; .float() is no-op when dtype already matches. + # q/k/g/beta/cu/ci are per-call activations: plain _ct, never cached + # (see the caching rule at the top of this file). + q_ct = _ct(q, cutlass.BFloat16) + k_ct = _ct(k, cutlass.BFloat16) + g_ct = _ct(g, cutlass.BFloat16) + alog_ct = _ct(A_log if A_log.dtype == torch.float32 else A_log.float(), cutlass.Float32) + if beta.dtype == torch.float32: + beta_etype = cutlass.Float32 + elif beta.dtype == torch.bfloat16: + beta_etype = cutlass.BFloat16 + else: + raise ValueError(f"Kimi K3 KDA prefill beta must be float32 or bfloat16, got {beta.dtype}") + beta_ct = _ct(beta, beta_etype) + beta_activated_ct = ( + cute_wrappers["k123_beta_activated_ct"] + if cute_wrappers is not None + else _ct(beta_activated, cutlass.BFloat16) + ) + + # Scratch wrappers come from the buffer-cache entry that owns the scratch + # itself, so wrapper and buffer die together and evicting the entry + # actually returns the ~1 GiB it holds. Caching them in a module-level + # id()-keyed dict instead is what leaked one full buffer set per distinct + # prefill token count. CONTRACT: cute_wrappers must be the dict returned + # by the same _get_buffers call that produced k_scaled..A_kk_inv -- + # _chunk_kda_fwd unpacks both from one return value. Direct low-level + # callers may omit it and pay ~10us per tensor to rebuild. + if cute_wrappers is not None: + ks_ct = cute_wrappers["k123_ks_ct"] + kg_ct = cute_wrappers["k123_kg_ct"] + qs_ct = cute_wrappers["k123_qs_ct"] + gk_ct = cute_wrappers["k123_gk_ct"] + aqk_ct = cute_wrappers["k123_aqk_ct"] + akk_ct = cute_wrappers["k123_akk_ct"] + else: + ks_ct = _ct(k_scaled, cutlass.BFloat16) + kg_ct = _ct(kg, cutlass.BFloat16) + qs_ct = _ct(q_scaled, cutlass.BFloat16) + gk_ct = _ct(gk_last_exp, cutlass.Float32) + aqk_ct = _ct(A_qk, cutlass.BFloat16) + akk_ct = _ct(A_kk_inv, cutlass.BFloat16) + + if is_varlen: + cu_ct = _ct(cu_seqlens, _cute_int_type(cu_seqlens.dtype)) + ci_ct = _ct(chunk_indices, _cute_int_type(chunk_indices.dtype)) + else: + cu_ct, ci_ct = _get_eqlen_dummies(q.device, torch.int64) + + if dt_bias is not None: + bias_ct = _get_dt_bias_ct(dt_bias, H, K) + else: + bias_ct = _get_empty_bias_ct(q.device) + lb_val = float(lower_bound) if lower_bound is not None else 0.0 + + # Launch on torch's CURRENT stream — see _current_cu_stream. + stream = _current_cu_stream(q.device) + + ct_args = ( + q_ct, + k_ct, + g_ct, + alog_ct, + beta_ct, + beta_activated_ct, + scale, + ks_ct, + kg_ct, + qs_ct, + gk_ct, + aqk_ct, + akk_ct, + cu_ct, + ci_ct, + bias_ct, + lb_val, + # Runtime shape scalars (rt_nt / rt_b / rt_t_total in host_fn). + NT, + B, + B * (T_padded if is_varlen else NT * BT), + valid_tokens, + stream, + ) + + if cache_key not in _fused_k123_cache: + host_fn = _fused_make_host( + B, + NT, + H, + is_varlen=is_varlen, + T_padded=T_padded, + has_bias=has_bias, + use_safe_gate=safe_gate, + use_beta_sigmoid=use_beta_sigmoid_in_kernel, + varlen_pure=varlen_pure, + grid_blocks=grid_blocks, + ) + _fused_k123_cache[cache_key] = cute.compile(host_fn, *ct_args) + k123_fn = _fused_k123_cache[cache_key] + k123_fn(*ct_args) + + # ===== Chained BF16 akk_inv (in-place: A_kk_inv = (I+L)^-1) ===== + # Views are layout-dynamic so one compiled akk_inv serves every batch + # shape (the host builds its own runtime-shaped views from the iterator; + # the wrapper layout only matters for the call signature). + if akk_in_view is None: + akk_in_view = from_dlpack(A_kk_inv, assumed_align=16).mark_layout_dynamic() + akk_in_view.element_type = cutlass.Float32 + akk_out_view = from_dlpack(A_kk_inv, assumed_align=16).mark_layout_dynamic() + akk_out_view.element_type = cutlass.Float32 + + if is_varlen: + # Reuse the cached cute wrappers from K123 (same tensor objects). + akk_cu_ct = cu_ct + akk_ci_ct = ci_ct + is_varlen_int = 1 + T_val = T + else: + akk_cu_ct, akk_ci_ct = cu_ct, ci_ct + is_varlen_int = 0 + T_val = NT * BT + + beta_for_akk = beta_activated if use_beta_sigmoid_in_kernel else beta + if beta_for_akk.dtype == torch.float32: + akk_beta_etype = cutlass.Float32 + else: + akk_beta_etype = cutlass.BFloat16 + akk_beta_ct = ( + beta_activated_ct if use_beta_sigmoid_in_kernel else _ct(beta_for_akk, akk_beta_etype) + ) + + # B/NT/T_val remain runtime args of akk_inv_host. Eqlen T is still part of + # the cache key because the beta wrapper bakes its batch stride. + # cu/ci dtype is a specializer for the same reason as in the K123 key above + # (element type baked into the compiled reader). + akk_cache_key = (H, is_varlen, dev, cu_ci_dtypes, eqlen_t_key, beta_for_akk.dtype) + if akk_cache_key not in _akk_inv_cache: + _akk_inv_cache[akk_cache_key] = cute.compile( + _akk_inv_host, + akk_in_view, + akk_out_view, + akk_beta_ct, + B, + NT, + H, + akk_cu_ct, + akk_ci_ct, + is_varlen_int, + T_val, + stream, + ) + akk_fn = _akk_inv_cache[akk_cache_key] + akk_args = (akk_in_view, akk_out_view, akk_beta_ct, B, NT, akk_cu_ct, akk_ci_ct, T_val, stream) + akk_fn(*akk_args) + + +def _chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = True, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + varlen_is_aligned: bool | None = None, + single_sequence_length: int | None = None, + snapshot_rows: torch.Tensor | None = None, + snapshot_chunks: torch.Tensor | None = None, +) -> torch.Tensor: + """Run optimized KDA prefill against indexed recurrent-state rows. + + ``snapshot_rows`` / ``snapshot_chunks`` additionally store sequence + ``b``'s state after its first ``64 * snapshot_chunks[b]`` tokens into + ``state_pool[snapshot_rows[b]]`` (see ``_launch_k4_persistent``). + """ + if safe_gate and lower_bound is None: + lower_bound = -5.0 + + is_varlen = cu_seqlens is not None + + if q.shape[1] == 0: + raise ValueError("Indexed KDA prefill does not support an empty token batch.") + + B, T, H, K = q.shape + V_dim = v.shape[-1] + device = q.device + BT = 64 + + # Phase 1: handle eqlen with T % 64 != 0 entirely on the eqlen path — + # never borrow varlen's mask code. Pad inputs to a CHUNKS_PER_BLOCK*BT + # = 256 multiple so the persistent scheduler's `cgs_per_head = NT // 4` + # divides cleanly and every chunk has 64 valid rows of (zero-padded / + # sentinel-padded) data. K123 eqlen kernel runs unchanged — it doesn't + # know partial chunks exist, the boundary is handled at the data + # boundary (K4's bounded-TMA principle, just expressed via host pad + # since K123 stores via autovec_copy not TMA). + # + # Varlen with non-aligned seq lengths is a separate problem (Phase 2): + # multi-seq varlen can't be host-padded without repacking memory. + real_T = T + CPB_BT = 4 * BT # CHUNKS_PER_BLOCK * BT, the cgs_per_head divisibility unit + # The trigger is CPB_BT, not BT: the eqlen scheduler floors `cgs_per_head` + # and has no `chunk_idx < num_chunks` guard, so a BT-aligned T whose chunk + # count is not a multiple of CHUNKS_PER_BLOCK (T = 64, 128, 192, 320, ...) + # would silently drop its trailing chunks. + needs_eqlen_pad = (not is_varlen) and (T % CPB_BT != 0) + if needs_eqlen_pad: + if B != 1: + raise NotImplementedError( + f"eqlen with B>1 and T % {CPB_BT} != 0 not supported (got B={B}, T={T})." + ) + T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT + # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache + # key). torch.cat would reallocate + copy the full 200MB q tensor every + # call — caching the destination buffer drops that to a single slice + # copy of the valid prefix (caller already lives in our buffer for + # subsequent calls reusing the same id, but we re-copy unconditionally + # since the caller may have updated the data in-place). + q_pad, k_pad, v_pad, g_pad, beta_pad = _get_padded_input_buffers( + B, T_padded, H, K, q.dtype, g.dtype, beta.dtype, q.device, real_T + ) + # q/k/v/beta zero-padded → K1/K2 MMAs naturally produce 0 for OOB rows. + # Tail [real_T:] of q_pad/k_pad/v_pad/beta_pad is pre-zeroed at cache + # init and never written, so we only copy the valid prefix. + q_pad[:, :real_T].copy_(q) + k_pad[:, :real_T].copy_(k) + v_pad[:, :real_T].copy_(v) + beta_pad[:, :real_T].copy_(beta) + # g uses a -1e3 sentinel so the gate activation saturates to 0 for OOB + # rows (both safe_gate sigmoid and softplus paths). Plain g=0 gives + # nonzero activation that would corrupt the cumsum past seq end. The + # tail is set to -1e3 once at cache init; we only copy the valid prefix. + g_pad[:, :real_T].copy_(g) + q, k, v, g, beta = q_pad, k_pad, v_pad, g_pad, beta_pad + T = T_padded # downstream buffer alloc + kernel layout use T_padded + + # Phase 2.1: varlen with a SINGLE non-aligned sequence — caller already + # zero-padded q/k/v/beta to a 64-multiple (FLA convention), but g's tail + # is also zero, which causes the gate activation to be non-zero past seq + # end and corrupts the cumsum / GkLast. We sentinel-pad g (cheap: ~5MB + # copy) and force VARLEN_PURE=1 so all 4 mask sites compile-elide. Same + # K4 "boundary at the data" principle as the eqlen path. + # + # Multi-seq varlen is NOT handled here — its OOB regions overlap with + # adjacent seqs' data, so sentinel-pad on g would corrupt the next seq. + # Multi-seq optimization needs per-seq dynamic tensormap (Phase 2.2). + needs_varlen_single_pad = False + if is_varlen and cu_seqlens is not None and cu_seqlens.shape[0] == 2 and B == 1: + if single_sequence_length is not None: + single_sequence_is_aligned = single_sequence_length % BT == 0 + if varlen_is_aligned is not None and varlen_is_aligned != single_sequence_is_aligned: + raise ValueError("varlen_is_aligned disagrees with single_sequence_length") + varlen_is_aligned = single_sequence_is_aligned + if not single_sequence_is_aligned: + real_T = single_sequence_length + needs_varlen_single_pad = True + else: + _vl_key = id(cu_seqlens) + if _vl_key not in _varlen_pure_cache: + cu_cpu = cu_seqlens.cpu().tolist() + sl = cu_cpu[1] - cu_cpu[0] + _varlen_pure_cache[_vl_key] = sl % BT == 0 + _varlen_single_seqlen_cache[_vl_key] = sl + _prune_on_gc(_varlen_pure_cache, _vl_key, cu_seqlens) + _prune_on_gc(_varlen_single_seqlen_cache, _vl_key, cu_seqlens) + varlen_is_aligned = _varlen_pure_cache[_vl_key] + if not varlen_is_aligned: + real_T = _varlen_single_seqlen_cache[_vl_key] + needs_varlen_single_pad = True + varlen_pure_override = False + if needs_varlen_single_pad: + # q/k/v/beta already zero-padded by caller (FLA convention). Re-build g + # with -1000 sentinel in the tail so VARLEN_PURE=1 path is correct. + # Cache the resulting g buffer so repeated calls with same input ids + # don't re-allocate. + cur_T = q.shape[1] # caller's padded T + assert cur_T % BT == 0 and cur_T >= real_T, ( + f"varlen single-seq path expects caller-padded input " + f"(T={cur_T}, seqlen={real_T}); see " + "K3KDAKernelDispatch.prefill_chunk_kda" + ) + g_pad = _get_g_sentinel_buffer(B, cur_T, H, K, g.dtype, g.device, real_T) + g_pad[:, :real_T].copy_(g[:, :real_T]) + g = g_pad + # Force VARLEN_PURE=1 for THIS call only (the data is sentinel-padded + # now). Never write the override into _varlen_pure_cache: the cached + # verdict must keep saying "not pure" so a later call with the SAME + # cu_seqlens object re-runs this sentinel pad — poisoning the cache + # made the second call skip the pad while still compiling the + # mask-free variant (silent tail corruption). + varlen_pure_override = True + + # Phase 2.2 (multi-seq via host repack) was attempted but isn't net positive: + # the scatter/gather memcpy cost (~800us GPU bandwidth per call) exceeds + # the kernel mask-elision savings (~250us). Keeping multi-seq non-pure on + # the original masked path. The right fix is kernel-level dynamic + # tensormap (K4-style per-tile bounded TMA) but that requires a major + # K123 kernel refactor — left as future work. + # multiseq_info = None + + if is_varlen: + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) + # Any chunk count is launchable: the varlen scheduler rounds up + # (`total_cgs_per_head = ceil(num_chunks / CHUNKS_PER_BLOCK)`), launches + # the constant `NUM_SMS` grid, and guards every chunk loop with + # `chunk_idx < num_chunks`. Only eqlen needs divisibility, and it gets + # that from its own `CPB_BT` pad above. + N_seqs = len(cu_seqlens) - 1 + else: + NT = T // BT + N_seqs = B + + # ===== Cached buffers + cute wrappers (avoid alloc + from_dlpack overhead per call) ===== + ( + k_scaled, + kg, + q_scaled, + beta_activated, + gk_last_exp, + A_qk, + A_kk, + O_flat, + cu_eqlen, + co_eqlen, + cute_wrappers, + ) = _get_buffers(device, k.dtype, B, T, H, K, V_dim, NT, N_seqs, BT, varlen=is_varlen) + + # Beta is fused entirely in akk_inv kernel epilogue (post-inv column-scale). + # No host v*beta and no K1 k_scaled*beta any more. + _launch_fused_k123_inv( + q, + k, + g, + A_log, + beta, + beta_activated, + scale, + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk, + cu_seqlens, + chunk_indices, + is_varlen, + NT, + real_T, + dt_bias=dt_bias, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + akk_in_view=cute_wrappers["akk_in_view"], + akk_out_view=cute_wrappers["akk_out_view"], + cute_wrappers=cute_wrappers, + varlen_pure_override=varlen_pure_override, + varlen_is_aligned=varlen_is_aligned, + ) + + # ===== K4: persistent kernel (eqlen + varlen via cu_seqlens) ===== + if is_varlen: + # GPU-side cumsum + cache by id(cu_seqlens). No host sync. + cu_for_k4, chunk_offsets_for_k4 = _get_varlen_k4_inputs(cu_seqlens, BT) + else: + cu_for_k4 = cu_eqlen + chunk_offsets_for_k4 = co_eqlen + + _launch_k4_persistent( + cute_wrappers, + v, + state_pool, + state_indices, + cu_for_k4, + chunk_offsets_for_k4, + cu_eqlen_passed=(not is_varlen), + num_sm=_kda_persistent_grid_blocks(v.device), + H=H, + V_dim=V_dim, + snapshot_rows=snapshot_rows, + snapshot_chunks=snapshot_chunks, + ) + + o = O_flat + if needs_eqlen_pad: + # Caller called with original T = real_T; their downstream code expects + # outputs at that shape. Slice the padded scratch tail back off. + o = o[:, :real_T] + return o + + +class KdaPrefillRunner: + """Runs the source-integrated Kimi K3 KDA prefill pipeline.""" + + @staticmethod + def forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + scale: float, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = True, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + snapshot_rows: Optional[torch.Tensor] = None, + snapshot_chunks: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Run KDA prefill and update selected recurrent-state rows.""" + if not IS_CUTLASS_DSL_AVAILABLE: + raise RuntimeError("Kimi K3 KDA prefill requires NVIDIA CUTLASS DSL") + if chunk_size != 64: + raise ValueError(f"Kimi K3 KDA prefill requires chunk_size=64, got {chunk_size}") + if A_log is None and q.shape[1] != 0: + raise ValueError("Kimi K3 KDA prefill requires A_log") + + return _chunk_kda_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + state_pool=state_pool, + state_indices=state_indices, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_gate_in_kernel=use_gate_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, + snapshot_rows=snapshot_rows, + snapshot_chunks=snapshot_chunks, + ) + + +def _validate_snapshot_args( + num_sequences: int, + state_pool: torch.Tensor, + snapshot_rows: Optional[torch.Tensor], + snapshot_chunks: Optional[torch.Tensor], +) -> None: + if (snapshot_rows is None) != (snapshot_chunks is None): + raise ValueError("snapshot_rows and snapshot_chunks must be given together.") + if snapshot_rows is None: + return + for name, tensor in (("snapshot_rows", snapshot_rows), ("snapshot_chunks", snapshot_chunks)): + if tensor.dtype != torch.int32 or tensor.ndim != 1 or tensor.shape[0] != num_sequences: + raise ValueError( + f"Expected {name} to be an int32 tensor of shape [{num_sequences}], " + f"got {tensor.dtype} {tuple(tensor.shape)}." + ) + if not tensor.is_contiguous() or tensor.data_ptr() % 16 != 0: + raise ValueError(f"{name} must be contiguous and 16-byte aligned.") + if tensor.device != state_pool.device: + raise ValueError(f"{name} and state_pool must be on the same device.") + + +def _validate_indexed_state_pool( + q: torch.Tensor, + v: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], +) -> None: + num_sequences = q.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1 + H, K, V = q.shape[2], q.shape[3], v.shape[3] + if state_pool.dtype != torch.float32 or state_pool.ndim != 4: + raise ValueError( + "Expected state_pool to be a rank-4 fp32 tensor with shape [slots, H, V, K]." + ) + if state_pool.shape[1:] != (H, V, K): + raise ValueError( + f"Expected state_pool shape [slots, {H}, {V}, {K}], got {tuple(state_pool.shape)}." + ) + expected_inner_strides = (V * K, K, 1) + if state_pool.stride()[1:] != expected_inner_strides: + raise ValueError( + "Indexed KDA prefill requires dense H/V/K inner dimensions; " + f"expected strides {expected_inner_strides}, got " + f"{state_pool.stride()[1:]}." + ) + if state_pool.stride(0) < H * V * K or state_pool.stride(0) % 4 != 0: + raise ValueError( + "Indexed KDA prefill requires a non-overlapping, 16-byte-aligned slot stride." + ) + if state_pool.data_ptr() % 16 != 0: + raise ValueError("Indexed KDA prefill requires a 16-byte-aligned state pool.") + if state_indices.ndim != 1 or state_indices.shape[0] != num_sequences: + raise ValueError( + f"Expected state_indices shape [{num_sequences}], got {tuple(state_indices.shape)}." + ) + if state_indices.dtype not in (torch.int32, torch.int64): + raise ValueError("state_indices must have dtype int32 or int64.") + if not state_indices.is_contiguous(): + raise ValueError("state_indices must be contiguous.") + if state_indices.data_ptr() % 16 != 0: + raise ValueError("Indexed KDA prefill requires 16-byte-aligned state_indices.") + for name, tensor in ( + ("q", q), + ("v", v), + ("state_indices", state_indices), + ): + if tensor.device != state_pool.device: + raise ValueError(f"{name} and state_pool must be on the same device.") + + +def kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + scale: float, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = True, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + snapshot_rows: Optional[torch.Tensor] = None, + snapshot_chunks: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Run KDA prefill directly against a V-first recurrent-state pool. + + With ``snapshot_rows`` / ``snapshot_chunks`` (int32 ``[num_sequences]``), + sequence ``b`` with ``snapshot_chunks[b] = c >= 0`` also writes its state + after its first ``64 * c`` tokens (``0 <= c < ceil(len_b / 64)``) into + ``state_pool[snapshot_rows[b]]``, a row no sequence reads or updates. + """ + _validate_indexed_state_pool( + q, + v, + state_pool, + state_indices, + cu_seqlens, + ) + _validate_snapshot_args(state_indices.shape[0], state_pool, snapshot_rows, snapshot_chunks) + return KdaPrefillRunner.forward( + q=q, + k=k, + v=v, + g=g, + beta=beta, + state_pool=state_pool, + state_indices=state_indices, + A_log=A_log, + scale=scale, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_gate_in_kernel=use_gate_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, + snapshot_rows=snapshot_rows, + snapshot_chunks=snapshot_chunks, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/fused_kda_conv_prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/fused_kda_conv_prefill.py new file mode 100644 index 000000000000..6ada828d3199 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/fused_kda_conv_prefill.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA prefill's short convolution, SiLU and Q / K L2 norm in one Triton pass over the token-major +QKV projection, reading and writing back each sequence's conv state.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.kda import _kda_kernels +from .fused_kda_post_conv import l2_normalize_heads, pad_tokens + +PAD_SLOT_ID = _kda_kernels.PAD_SLOT_ID +TP1_HEADS = 96 # the tp1 context server runs every KDA head + + +class FusedKdaConvPrefill(OpWrapper): + """Return ``(q, k, v)`` (and ``beta``) of packed sequences, writing their conv states. + + ``qkv`` ``[T, 3 * H * D]`` holds the Q, K, V sections of the projection; sequence ``b`` is + rows ``query_start_loc[b] : query_start_loc[b + 1]``. Each channel is convolved causally with + ``conv_weight`` ``[3 * H * D, W]`` over the sequence, preceded by ``conv_states[s]`` + (``s = cache_indices[b]``) where ``has_initial_state[b]`` else zeros, then SiLU and rounded + to ``qkv.dtype``; Q and K are then ``t / sqrt(sum(t^2) + l2_norm_eps)`` per head (with + ``normalize_qk``). Outputs are contiguous ``[1, T', H, D]``, ``T'`` = ``padded_tokens`` or + ``T``, rows past ``T`` zero; with ``beta`` also its fp32 ``[1, T', H]`` copy. In place: each + sequence's last ``W - 1`` inputs (after its initial state) go to ``conv_states[s]``; a + sequence on ``pad_slot_id`` reads no state and writes none. + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell( + why="tp1, one fresh 300-token sequence: Q/K/V and beta zero-padded to the 64 chunk", + spec=dict(seq_lens=(300,), has_init=(False,), prepad=True), + ), + Cell( + why="tp1, a continuation chunk of 2048 tokens from its conv state, chunk-aligned", + spec=dict(seq_lens=(2048,), has_init=(True,), prepad=False), + ), + Cell( + why="tp1, four mixed fresh / continuing sequences, one with fewer tokens than the width", + spec=dict(seq_lens=(64, 2, 129, 63), has_init=(True, True, False, True), prepad=False), + ), + ) + note = """ + Kimi K3's KDA prefill calls it where features.kda_fused_conv_prefill is on (the tp1 context + target): qkv the [q | k | v] column view of the QKVG output (unit channel stride, row stride + 4 * H * D), the [W, 3 * H * D] weight copy transposed (unit channel stride), the bf16 + [slots, 3 * H * D, 3] conv pool, the metadata's int32 query_start_loc and cache_indices, + 96 heads of 128, width 4, normalize_qk; padded_tokens and beta (a [f_a | b] column view) + for a single sequence whose length is not a chunk multiple (features.kda_prepad_prefill). + + Two launches: the token kernel, then the sequence-edge kernel, which rewrites the first + W - 1 tokens of each sequence and updates its conv state after reading it. The conv + accumulates in fp32 (bf16 operands with mixed-precision FMA on SM100+), SiLU and the norm + use approximate exp / reciprocal: results are the unfused causal_conv1d + post-conv chain's + to rounding. Outputs are new. Triton compiles once per shape class. + """ + + def __call__( + self, + qkv: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float = 1e-6, + padded_tokens: Optional[int] = None, + beta: Optional[torch.Tensor] = None, + pad_slot_id: int = PAD_SLOT_ID, + normalize_qk: bool = True, + ) -> tuple[torch.Tensor, ...]: + return _kda_kernels.fused_kda_conv_prefill( + qkv, + conv_weight, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + num_heads, + head_dim, + l2_norm_eps, + padded_tokens, + beta, + pad_slot_id, + normalize_qk, + ) + + def reference( + self, + qkv: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float = 1e-6, + padded_tokens: Optional[int] = None, + beta: Optional[torch.Tensor] = None, + pad_slot_id: int = PAD_SLOT_ID, + normalize_qk: bool = True, + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: + """fp64 convolution: ``(outputs, conv_states after the call)``; nothing is written.""" + tokens = qkv.shape[0] + width = conv_weight.shape[1] + weight = conv_weight.double() + states = conv_states.clone() + activations = torch.zeros(qkv.shape[1], tokens, dtype=torch.float64, device=qkv.device) + starts = query_start_loc.tolist() + for (lo, hi), slot, init in zip( + zip(starts, starts[1:]), cache_indices.tolist(), has_initial_state.tolist() + ): + seq = qkv[lo:hi].double().t() + padding = slot == pad_slot_id + if init and not padding: + hist = conv_states[slot].double() + else: + hist = seq.new_zeros(seq.shape[0], width - 1) + full = torch.cat([hist, seq], dim=1) + y = sum(weight[:, tap : tap + 1] * full[:, tap : tap + hi - lo] for tap in range(width)) + activations[:, lo:hi] = y * torch.sigmoid(y) + if not padding: + states[slot] = full[:, -(width - 1) :].to(states.dtype) + rounded = activations.to(qkv.dtype).t().reshape(1, tokens, 3, num_heads, head_dim) + q, k, v = rounded.unbind(2) + if normalize_qk: + q = l2_normalize_heads(q, l2_norm_eps).to(qkv.dtype) + k = l2_normalize_heads(k, l2_norm_eps).to(qkv.dtype) + out_tokens = tokens if padded_tokens is None else padded_tokens + outputs = tuple(pad_tokens(t, out_tokens).contiguous() for t in (q, k, v)) + if beta is not None: + beta_rows = beta.reshape(1, tokens, num_heads).float() + outputs = (*outputs, pad_tokens(beta_rows, out_tokens).contiguous()) + return outputs, states + + def is_valid( + self, + qkv: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float = 1e-6, + padded_tokens: Optional[int] = None, + beta: Optional[torch.Tensor] = None, + pad_slot_id: int = PAD_SLOT_ID, + normalize_qk: bool = True, + ) -> None: + starts = query_start_loc.tolist() + # The kernels take the sequence bounds as given: overlapping or out-of-range spans + # convolve the wrong rows without an error. + assert starts[0] == 0 and starts[-1] == qkv.shape[0], "query_start_loc must span qkv" + assert all(lo <= hi for lo, hi in zip(starts, starts[1:])), "query_start_loc must ascend" + slots = [s for s in cache_indices.tolist() if s != pad_slot_id] + assert all(0 <= s < conv_states.shape[0] for s in slots), "slot outside conv_states" + assert len(set(slots)) == len(slots), "two sequences on one conv-state slot race" + + +fused_kda_conv_prefill = FusedKdaConvPrefill() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/fused_kda_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/fused_kda_post_conv.py new file mode 100644 index 000000000000..0d2e0843650a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/fused_kda_post_conv.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA prefill's channel-major convolution output to token-major Q / K / V (Triton), Q and K +L2-normalized per head, optionally zero-padded to the prefill chunk with beta alongside.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.kda import _kda_kernels + +# The prefill token counts the cells drive: a single token, a partial final chunk, and whole chunks. +_TOKENS = (1, 127, 1024) +# Generation targets: (target, KDA heads per rank). +_GENERATION = (("helix8", 12), ("helix16_moetp4ep4", 6), ("dep16", 96)) + + +def l2_normalize_heads(t: torch.Tensor, eps: float) -> torch.Tensor: + """``t / sqrt(sum(t^2) + eps)`` over the last dimension, in fp64.""" + td = t.double() + return td / torch.sqrt(td.square().sum(-1, keepdim=True) + eps) + + +def pad_tokens(t: torch.Tensor, tokens: int) -> torch.Tensor: + """``t`` ``[1, T, ...]`` zero-padded along dimension 1 to ``tokens`` rows.""" + if t.shape[1] == tokens: + return t + pad = t.new_zeros((t.shape[0], tokens - t.shape[1], *t.shape[2:])) + return torch.cat([t, pad], dim=1) + + +class FusedKdaPostConv(OpWrapper): + """Split ``packed`` ``[3 * H * D, T]`` (the convolution's Q, K, V sections, channel-major) into + contiguous ``[1, T', H, D]`` Q, K, V in ``packed.dtype``; Q and K are ``t / sqrt(sum(t^2) + + l2_norm_eps)`` per head. ``T'`` is ``padded_tokens`` when given, else ``T``; rows past ``T`` + are zeros. With ``beta`` (``[T, H]`` or ``[1, T, H]``, unit head stride) a fourth output is + its fp32 ``[1, T', H]`` copy, zero-padded the same way.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"{target} prefill ({heads} heads), {tokens} tokens: the unfused convolution path", + spec=dict(heads=heads, tokens=tokens), + ) + for target, heads in _GENERATION + for tokens in _TOKENS + ) + note = """ + Kimi K3's KDA prefill calls it after the causal convolution where + features.kda_fused_conv_prefill is off (every generation target), with head_dim 128 and the + default l2_norm_eps, on the convolution's contiguous output. features.kda_prepad_prefill + asks for padded_tokens and beta (a single sequence padded to the prefill chunk), but only the + tp1 target sets it, and that target runs fused_kda_conv_prefill instead: the padded form is + accepted and tested by its kernel test, not certified here. + + No state and no metadata; inputs are not mutated, outputs are new. Triton compiles once per + shape class (the token counts are not specialization keys). A non-contiguous packed takes + the portable torch path, which takes no padding. + """ + + def __call__( + self, + packed: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float = 1e-6, + padded_tokens: Optional[int] = None, + beta: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, ...]: + return _kda_kernels.fused_kda_post_conv( + packed, num_heads, head_dim, l2_norm_eps, padded_tokens, beta + ) + + def reference( + self, + packed: torch.Tensor, + num_heads: int, + head_dim: int, + l2_norm_eps: float = 1e-6, + padded_tokens: Optional[int] = None, + beta: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, ...]: + tokens = packed.shape[1] + out_tokens = tokens if padded_tokens is None else padded_tokens + q, k, v = packed.reshape(3, num_heads, head_dim, tokens).permute(0, 3, 1, 2).unsqueeze(1) + outputs = ( + l2_normalize_heads(q, l2_norm_eps).to(packed.dtype), + l2_normalize_heads(k, l2_norm_eps).to(packed.dtype), + v, + ) + outputs = tuple(pad_tokens(t, out_tokens).contiguous() for t in outputs) + if beta is None: + return outputs + beta_rows = beta.reshape(1, tokens, num_heads).float() + return (*outputs, pad_tokens(beta_rows, out_tokens).contiguous()) + + +fused_kda_post_conv = FusedKdaPostConv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_kda_fab_fa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_kda_fab_fa.py new file mode 100644 index 000000000000..a09fb0c24bc4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_kda_fab_fa.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KDA verify's ``[f_a | b]`` projection (CuTe DSL, the fab kernel's GEMM 1 and reduction), for +a verify kernel that computes the forget gate from ``f_a`` itself.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_kda_fab +from .k3_kda_fab_proj import fab_cells + + +class K3KdaFabFa(OpWrapper): + """Return ``bfa = bf16(x @ w_bfa^T)`` (bf16 ``[M, N1]``, fp32 accumulation): columns + ``[0, 128)`` are ``f_a``, ``[128, 128 + num_heads)`` beta, bitwise those of + ``k3_kda_fab_proj`` for the same ``w_fb`` row count. ``w_fb`` is only prefetched into L2.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS: tuple[Cell, ...] = fab_cells("f_a and beta for the verify's in-kernel forget gate") + note = """ + Kimi K3's KDA verify calls it on the projection aux stream next to the QKVG GEMM, where the + verify kernel runs in one wave (or above 224 rows) and can_use_k3_kda_fab_fa accepts the + operands (as for k3_kda_fab_proj); the verify kernel then computes bf16(f_a @ w_fb^T). Columns + past 128 + num_heads are the padding rows' products (zero weights). The output is a new + contiguous tensor. + + Same per-device hand-off workspace and ticket counter as k3_kda_fab_proj (process-wide; + launches on a device must not overlap); every CTA becomes a worker, so all of them (at most + 28) must become co-resident. Compiled kernels are cached per process. The launch uses PDL. + Importing the entry needs CuTe DSL (SM100 / SM103 / SM107). + """ + + def __call__( + self, + x: torch.Tensor, + w_bfa: torch.Tensor, + num_heads: int, + w_fb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return k3_kda_fab.k3_kda_fab_fa(x, w_bfa, num_heads, w_fb) + + def reference( + self, + x: torch.Tensor, + w_bfa: torch.Tensor, + num_heads: int, + w_fb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``w_fb`` and ``num_heads`` do not enter the result.""" + return (x.double() @ w_bfa.double().t()).to(torch.bfloat16) + + +k3_kda_fab_fa = K3KdaFabFa() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_kda_fab_proj.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_kda_fab_proj.py new file mode 100644 index 000000000000..6f3a921249c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_kda_fab_proj.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KDA verify's beta and forget gate in one launch (CuTe DSL): ``[f_a | b]`` GEMM, then the +``f_b`` GEMM of its ``f_a`` columns.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_kda_fab + +HIDDEN = 7168 # Kimi K3's hidden size, the GEMM's K +FA_DIM = 128 # f_a width + +# The verify row counts (requests x (1 + drafts)) of the targets whose KDA head shard the fab +# kernels take (at most 16 heads): (target, heads per rank, rows, what). +FAB_SHAPES = ( + ("helix8", 12, 96, "c240: 16 requests x 6 rows"), + ("helix8", 12, 64, "c75 / c174: 8 requests x 8 rows"), + ("helix8", 12, 16, "c24 / c42: 2 requests x 8 rows"), + ("helix16_moetp4ep4", 6, 8, "ULL c1: one request x 8 rows"), + ("helix16_moetp4ep4", 6, 32, "ULL c4: 4 requests x 8 rows"), +) + + +def bfa_rows(heads: int) -> int: + """Rows of the fused ``[f_a | b]`` weight: ``128 + heads`` padded to a multiple of 8.""" + return -(-(FA_DIM + heads) // 8) * 8 + + +def fab_cells(why_suffix: str) -> tuple[Cell, ...]: + """One cell per verify shape of :data:`FAB_SHAPES`.""" + return tuple( + Cell( + why=f"{target}, {heads} heads, {rows} rows ({what}): {why_suffix}", + spec=dict(heads=heads, rows=rows), + ) + for target, heads, rows, what in FAB_SHAPES + ) + + +class K3KdaFabProj(OpWrapper): + """Return ``(beta, forget_gate)``: ``bfa = bf16(x @ w_bfa^T)``, ``beta = bfa[:, 128 : 128 + H]`` + (bf16 ``[M, H]``) and ``forget_gate = bf16(bfa[:, :128] @ w_fb^T)`` (bf16 ``[M, N2]``), both + GEMMs accumulated in fp32.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS: tuple[Cell, ...] = fab_cells("beta and the forget gate for the BF16 verify core") + note = """ + Kimi K3's KDA verify calls it on the projection aux stream next to the QKVG GEMM, for + M <= 256 rows where can_use_k3_kda_fab accepts the operands (contiguous bf16 x [M, 7168], + w_bfa = [f_a_proj; b_proj] padded to a multiple of 8 rows, at most 144; w_fb = f_b_proj + [H * 128, 128]; at most 16 heads). Outputs are new contiguous tensors. + + The CTAs hand their partial sums over through a per-device workspace and ticket counter + (process-wide, created on the first call outside graph capture; can_use_k3_kda_fab rejects a + capturing caller until then). Launches sharing a device must not overlap: stream order, and + with PDL the ticket and the workspace refill come after griddepcontrol.wait. Compiled kernels + are cached per process by tiling and shape. The launch uses PDL. Importing the entry needs + CuTe DSL (SM100 / SM103 / SM107). + """ + + def __call__( + self, x: torch.Tensor, w_bfa: torch.Tensor, w_fb: torch.Tensor, num_heads: int + ) -> tuple[torch.Tensor, torch.Tensor]: + return k3_kda_fab.k3_kda_fab_proj(x, w_bfa, w_fb, num_heads) + + def reference( + self, x: torch.Tensor, w_bfa: torch.Tensor, w_fb: torch.Tensor, num_heads: int + ) -> tuple[torch.Tensor, torch.Tensor]: + bfa = (x.double() @ w_bfa.double().t()).to(torch.bfloat16) + gate = (bfa[:, :FA_DIM].double() @ w_fb.double().t()).to(torch.bfloat16) + return bfa[:, FA_DIM : FA_DIM + num_heads].contiguous(), gate + + +k3_kda_fab_proj = K3KdaFabProj() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_l2_prefetch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_l2_prefetch.py new file mode 100644 index 000000000000..7499779113c1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/k3_l2_prefetch.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""L2 prefetch of byte ranges (CuTe DSL ``cp.async.bulk.prefetch.L2``): the next KDA layer's +projection weights, pulled into L2 on a side stream during a verify step.""" + +from __future__ import annotations + +from typing import Sequence + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels import k3_l2_prefetch as _kernel + +DEFAULT_MAX_CTAS = _kernel.DEFAULT_MAX_CTAS + + +class K3L2Prefetch(OpWrapper): + """Issue L2 prefetches for ``ranges`` (``(address, bytes)``, 16-byte aligned, at most eight; + empty ones dropped) on the current stream. Returns None and writes nothing.""" + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell( + why=( + "helix8 (12 heads), up to 16 verify rows: the next KDA layer's QKVG and o_proj " + "FP8 weights and scales, [f_a | b] and f_b" + ), + spec=dict( + nbytes=( + 4 * 12 * 128 * 7168, # qkvg_proj weight, fp8 + 4 * 12 * 7168 // 128 * 4, # its 128 x 128 block scales, fp32 + 7168 * 12 * 128, # o_proj weight, fp8 + 7168 // 128 * 12 * 4, # its block scales, fp32 + 144 * 7168 * 2, # [f_a | b], bf16, 128 + 12 rows padded to 144 + 12 * 128 * 128 * 2, # f_b, bf16 + ) + ), + ), + ) + note = """ + The only effect is on the L2 contents: no memory is written and nothing waits for the + prefetches, so the reference is a no-op and the test proves the ranges' bytes and their + neighbours unchanged. Kimi K3's KDA verify (features.kda_l2_prefetch, the helix8 target, up to + 16 rows) launches it on the projection aux stream behind an event recorded after this + layer's QKVG and [f_a | b] kernels, without PDL; it triggers its programmatic dependents at + once. CTAs: one warp each, about one 8 KiB chunk per thread, at most max_ctas. Compiled + kernels are cached per (range count, CTAs). Importing the entry needs CuTe DSL. + """ + + def __call__(self, ranges: Sequence[tuple[int, int]], max_ctas: int = DEFAULT_MAX_CTAS) -> None: + _kernel.run_k3_l2_prefetch(ranges, max_ctas=max_ctas) + + def reference( + self, ranges: Sequence[tuple[int, int]], max_ctas: int = DEFAULT_MAX_CTAS + ) -> None: + """Nothing: prefetches have no result.""" + return None + + +k3_l2_prefetch = K3L2Prefetch() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/kda_mtp_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/kda_mtp_decode.py new file mode 100644 index 000000000000..9dc2771a3b37 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/kda_mtp_decode.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused KDA multi-token verify (CuTe DSL): replay of the accepted drafts, short convolution, +gates, delta rule over the new tokens, and the gated output norm, committing the state in place.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.kda import mtp_ops + +HEAD_DIM = 128 # K = V of the kernel +CONV_WIDTH = 4 +MXFP8_BLOCK = 128 # elements per UE8M0 scale of the MXFP8 output +FP8_MAX = 448.0 + +# The deployments' verify shapes: (target, KDA heads per rank, requests, drafts). +_SHAPES = ( + ("helix8", 12, 16, 5, "c240 generation: 16 requests with five drafts"), + ("helix8", 12, 8, 7, "c75 / c174 generation: 8 requests with seven drafts"), + ("helix8", 12, 2, 7, "c24 / c42 generation: 2 requests with seven drafts"), + ("helix16_moetp4ep4", 6, 1, 7, "ULL c1: one request with seven drafts"), + ("helix16_moetp4ep4", 6, 4, 7, "ULL c4: 4 requests with seven drafts"), + ("dep16", 96, 16, 5, "c383 generation: all 96 heads, 16 requests with five drafts"), +) + + +def _variants(heads: int) -> tuple[tuple[str, str], ...]: + """The verify variants a layer of ``heads`` heads runs (the f_a gate needs at most 16).""" + found = [ + ("bf16", "BF16 core after the fab / F.linear gate, PDL launch"), + ("mxfp8", "MXFP8 core for the prequantized o_proj"), + ] + if heads <= 16: + found.append(("fused_fb", "forget gate computed from f_a in the kernel, PDL launch")) + return tuple(found) + + +def _ceil_log2_exponent(x: torch.Tensor) -> torch.Tensor: + """``ceil(log2(x))`` of positive fp32 ``x``, exactly (the UE8M0 round-up).""" + mantissa, exponent = torch.frexp(x) + return exponent - (mantissa == 0.5).to(exponent.dtype) + + +def _mxfp8_rows(y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """E4M3 ``[rows, heads * V]`` and the flat R128c4 UE8M0 scales of bf16 ``y`` ``[rows, heads, + V]``: one power-of-two scale per (row, head), ``2 ** ceil(log2(amax / 448))``, replicated four + times. Scale bytes of the 128-row block's padding rows are zero here (never written).""" + rows, heads, width = y.shape + yf = y.float() + amax = yf.abs().amax(dim=-1).clamp_min(1e-10) + exponent = _ceil_log2_exponent(amax * (1.0 / FP8_MAX)) + quantized = (yf * torch.exp2(-exponent.float())[..., None]).to(torch.float8_e4m3fn) + padded_rows = -(-rows // MXFP8_BLOCK) * MXFP8_BLOCK + scales = torch.zeros(padded_rows * heads * 4, dtype=torch.uint8, device=y.device) + row = torch.arange(rows, device=y.device)[:, None] + head = torch.arange(heads, device=y.device)[None, :] + base = ( + (row // MXFP8_BLOCK) * heads * 512 + head * 512 + (row % 32) * 16 + ((row % 128) // 32) * 4 + ) + code = (exponent + 127).to(torch.uint8) + for replica in range(4): + scales[(base + replica).flatten()] = code.flatten() + return quantized.reshape(rows, heads * width), scales + + +class KdaMtpDecode(OpWrapper): + """Verify ``1 + num_spec`` new tokens per request against its state slot; return the core. + + Request ``n`` owns slot ``ssm_state_indices[n]``. It first replays its ``a`` accepted drafts + of the previous round from the replay caches (``k_cache`` / ``g_cache`` / ``v_cache`` / + ``beta_cache``, the conv inputs in ``cs_*[..., W - 1:]``), then runs the new tokens: depthwise + convolution over the conv window, SiLU, Q/K L2 norm (Q scaled by ``scale``), the gate + ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` (``g`` the raw gate, or + ``bf16(fa @ w_fb^T)`` when ``fa`` / ``w_fb`` are given), ``sigmoid(beta)`` and the delta rule. + With ``packed_token_layout`` the new tokens are rows ``n * (1 + num_spec) + j`` and ``a`` is + ``num_accepted_tokens[slot]``; otherwise they start at ``cu_seqlens[n] + a`` with + ``a = num_accepted_tokens[n]``. + + Returns the core ``[1, T, H, V]`` (bf16, ``rmsnorm(bf16(o)) * onorm_weight * + sigmoid(onorm_g)`` with ``fuse_output_norm``), or with ``quantize_output`` that normalized + core as E4M3 ``[T, H * V]`` and its UE8M0 1x128 scales in the flat R128c4 layout. + + In place: after the first new token the slot's state goes to ``recurrent_state`` and its + conv window to ``cs_*[..., :W - 1]``; the other new tokens are cached as the next round's + drafts (``k`` / ``g`` / ``v`` / ``beta`` post-activation, raw conv inputs in ``cs_*``). + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"{target}, {heads} heads: {why}; {variant_why}", + spec=dict(heads=heads, requests=requests, num_spec=num_spec, variant=variant), + ) + for target, heads, requests, num_spec, why in _SHAPES + for variant, variant_why in _variants(heads) + ) + note = """ + Kimi K3 calls it from the KDA layer's verify (K3KDAKernelDispatch.mtp_verify / + mtp_verify_fused_fb) with packed_token_layout, fuse_output_norm, K = V = 128, conv width 4, + lower_bound -5, scale 128**-0.5, onorm_eps 1e-5 and the cache manager's pools: fp32 V-first + recurrent_state, dim-contiguous bf16 conv caches [slots, H * 128, 3 + num_spec], replay caches + k / v / beta bf16 and g fp32, int32 slot-indexed accepted counts. x_q / x_k / x_v / onorm_g + are row-strided sections of the QKVG output, beta (and fa) of the [f_a | b] output. The + variants: fa / w_fb (the forget gate computed in the kernel, BF16 core, use_pdl), BF16 core + (use_pdl), quantize_output (MXFP8 core, no PDL). The legacy cu_seqlens layout and the + zero_accepted / regular_metadata hints are accepted but only K3's packed form is certified. + + Kernel selection is per call: the v2 kernel for the packed fused-norm BF16 contract (unless + TRTLLM_K3_KDA_MTP_V2=0), on SM107 the two-CTA kernel for the MXFP8 output, else the v1 + kernel (SM107: its compact layout on multi-wave grids with at most three drafts); all compute + the same arithmetic. TRTLLM_KDA_MTP_BF16_MATH=1 switches the recurrence to bf16 operands. + + A process-wide cache holds the compiled kernels, keyed by layouts and constexpr flags; warm + each key before graph capture. Outputs are new tensors. Launches go to the current stream; + with use_pdl the kernel loads state, replay entries and weights before griddepcontrol.wait. + Importing the entry needs CuTe DSL. + """ + + def __call__( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: Optional[torch.Tensor], + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + fuse_output_norm: bool = False, + quantize_output: bool = False, + fa: Optional[torch.Tensor] = None, + w_fb: Optional[torch.Tensor] = None, + use_pdl: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return mtp_ops.kda_mtp_decode_impl( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + k_cache=k_cache, + g_cache=g_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=ssm_state_indices, + cu_seqlens=cu_seqlens, + num_spec=num_spec, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + zero_accepted_hint=zero_accepted_hint, + regular_metadata_hint=regular_metadata_hint, + packed_token_layout=packed_token_layout, + onorm_g=onorm_g, + onorm_weight=onorm_weight, + onorm_eps=onorm_eps, + fuse_output_norm=fuse_output_norm, + quantize_output=quantize_output, + fa=fa, + w_fb=w_fb, + use_pdl=use_pdl, + ) + + def reference( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: Optional[torch.Tensor], + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + fuse_output_norm: bool = False, + quantize_output: bool = False, + fa: Optional[torch.Tensor] = None, + w_fb: Optional[torch.Tensor] = None, + use_pdl: bool = False, + ) -> tuple[torch.Tensor | tuple[torch.Tensor, torch.Tensor], dict[str, torch.Tensor]]: + """fp64 token-by-token verify: ``(output, post)``, nothing written. + + ``output`` is what the op returns; ``post`` maps ``cs_q`` / ``cs_k`` / ``cs_v`` / + ``recurrent_state`` / ``k_cache`` / ``g_cache`` / ``v_cache`` / ``beta_cache`` to their + expected contents after the call. Rows the op does not write are zero (the legacy + layout's replay holes; the scale bytes of the padding rows). The hints and ``use_pdl`` + do not enter the result. + """ + f64 = torch.float64 + heads = A_log.shape[0] + key_dim, value_dim = x_q.shape[-1], x_v.shape[-1] + width = w_q.shape[1] + rows = x_q.shape[1] + scale = key_dim**-0.5 if scale is None else scale + names = ("cs_q", "cs_k", "cs_v", "recurrent_state", "k_cache", "g_cache", "v_cache") + sources = (cs_q, cs_k, cs_v, recurrent_state, k_cache, g_cache, v_cache) + post = {name: t.clone() for name, t in zip(names, sources)} + post["beta_cache"] = beta_cache.clone() + convs = (cs_q, cs_k, cs_v) + weights = tuple(w.double() for w in (w_q, w_k, w_v)) + inputs = tuple(x[0].reshape(rows, -1) for x in (x_q, x_k, x_v)) + if fa is not None: + raw_gate = (fa[0].double() @ w_fb.double().t()).to(torch.bfloat16).double() + else: + raw_gate = g[0].reshape(rows, -1).double() + raw_gate = raw_gate.view(rows, heads, key_dim) + dt_bias.double().view(heads, key_dim) + rate = A_log.double().exp()[:, None] + out = torch.zeros(rows, heads, value_dim, dtype=f64, device=x_q.device) + + def delta_rule(s, k_t, g_t, v_t, b_t): + s = s * g_t.exp()[:, None, :] + err = v_t - torch.einsum("hvk,hk->hv", s, k_t) + return s + b_t[:, None, None] * err[:, :, None] * k_t[:, None, :] + + slots = ssm_state_indices.tolist() + accepted = num_accepted_tokens.tolist() + starts = None if cu_seqlens is None else cu_seqlens.tolist() + for n, slot in enumerate(slots): + if packed_token_layout: + drafts, first = accepted[slot], n * (1 + num_spec) + else: + drafts = accepted[n] + first = starts[n] + drafts + window = [c[slot, :, : width - 1].double() for c in convs] + s = recurrent_state[slot].double() + for i in range(drafts): + window = [ + torch.cat([w[:, 1:], c[slot, :, width - 1 + i].double()[:, None]], dim=1) + for w, c in zip(window, convs) + ] + s = delta_rule( + s, + k_cache[slot, i].double().view(heads, key_dim), + g_cache[slot, i].double().view(heads, key_dim), + v_cache[slot, i].double().view(heads, value_dim), + beta_cache[slot, i].double().view(heads), + ) + for j in range(1 + num_spec): + row = first + j + raw = [x[row].double() for x in inputs] + window = [torch.cat([w, r[:, None]], dim=1) for w, r in zip(window, raw)] + act = [(w * wt).sum(-1) for w, wt in zip(window, weights)] + act = [a * torch.sigmoid(a) for a in act] + window = [w[:, 1:] for w in window] + q_t = torch.nn.functional.normalize(act[0].view(heads, key_dim), dim=-1) * scale + k_t = torch.nn.functional.normalize(act[1].view(heads, key_dim), dim=-1) + v_t = act[2].view(heads, value_dim) + g_t = lower_bound * torch.sigmoid(rate * raw_gate[row]) + b_t = torch.sigmoid(beta[0, row].double()) + s = delta_rule(s, k_t, g_t, v_t, b_t) + out[row] = torch.einsum("hvk,hk->hv", s, q_t) + if j == 0: + post["recurrent_state"][slot] = s.to(recurrent_state.dtype) + for name, w in zip(("cs_q", "cs_k", "cs_v"), window): + post[name][slot, :, : width - 1] = w.to(cs_q.dtype) + continue + draft = j - 1 + post["k_cache"][slot, draft] = k_t.reshape(k_cache[slot, draft].shape) + post["g_cache"][slot, draft] = g_t.reshape(g_cache[slot, draft].shape) + post["v_cache"][slot, draft] = v_t.reshape(v_cache[slot, draft].shape) + post["beta_cache"][slot, draft] = b_t.reshape(beta_cache[slot, draft].shape) + for name, r in zip(("cs_q", "cs_k", "cs_v"), raw): + post[name][slot, :, width - 1 + draft] = r.to(cs_q.dtype) + + if not (fuse_output_norm or quantize_output): + return out.to(x_q.dtype).unsqueeze(0), post + core = out.to(torch.bfloat16).double() + rms = torch.rsqrt(core.square().mean(-1, keepdim=True) + onorm_eps) + gate = torch.sigmoid(onorm_g[0].double().view(rows, heads, value_dim)) + normed = (core * rms * onorm_weight.double() * gate).to(torch.bfloat16) + if quantize_output: + return _mxfp8_rows(normed), post + return normed.unsqueeze(0), post + + def is_valid( + self, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: Optional[torch.Tensor], + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + fuse_output_norm: bool = False, + quantize_output: bool = False, + fa: Optional[torch.Tensor] = None, + w_fb: Optional[torch.Tensor] = None, + use_pdl: bool = False, + ) -> None: + slots = ssm_state_indices.tolist() + accepted = num_accepted_tokens.tolist() + counts = [accepted[s] for s in slots] if packed_token_layout else accepted[: len(slots)] + # Two requests on one slot race on its state, conv window and replay caches. + assert len(set(slots)) == len(slots), "ssm_state_indices must name one slot per request" + # Only num_spec drafts are cached: a larger count replays entries no round wrote. + assert all(0 <= a <= num_spec for a in counts), "accepted counts must be in [0, num_spec]" + # The hints select specialized kernels that trust them. + assert not zero_accepted_hint or not any(counts), "zero_accepted_hint with accepted drafts" + if regular_metadata_hint: + requests = len(slots) + assert slots == list(range(requests)), "regular_metadata_hint needs identity slots" + assert cu_seqlens is not None and cu_seqlens.tolist() == [ + n * (2 * num_spec + 1) for n in range(requests + 1) + ], "regular_metadata_hint needs cu_seqlens = arange * (2 * num_spec + 1)" + + +kda_mtp_decode = KdaMtpDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/kda_prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/kda_prefill.py new file mode 100644 index 000000000000..3ec12808be29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/kda_prefill.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Chunked KDA prefill (CuTe DSL: fused K1-K3, A_kk inverse, persistent K4) on L2-normalized Q/K, +updating indexed fp32 V-first state rows in place.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from . import _kda_prefill as prefill_ops + +#: Kimi K3's gate_lower_bound, the launcher's default under ``safe_gate``. +_DEFAULT_LOWER_BOUND = -5.0 + + +def _sequences(q: torch.Tensor, cu_seqlens: Optional[torch.Tensor]) -> list[tuple[int, int, int]]: + """``(row, start, end)`` of every sequence: equal-length rows, or the ``cu_seqlens`` spans of + row 0.""" + if cu_seqlens is None: + return [(row, 0, q.shape[1]) for row in range(q.shape[0])] + starts = cu_seqlens.tolist() + return [(0, lo, hi) for lo, hi in zip(starts, starts[1:])] + + +class KdaPrefill(OpWrapper): + """Run the gated delta rule over each sequence from ``state_pool[state_indices[b]]``; return + ``o`` ``[B, T, H, V]``. + + Sequences are the ``B`` rows of equal length, or (``cu_seqlens``, ``B == 1``) the packed + spans. ``q`` / ``k`` arrive L2-normalized; ``q`` is scaled by ``scale``. With + ``use_gate_in_kernel`` the log-decay is ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` + (``safe_gate``) or ``-exp(A_log) * softplus(g + dt_bias)``, else ``g`` itself; with + ``use_beta_sigmoid_in_kernel`` beta is ``sigmoid(beta)``. The final state of sequence ``b`` is + written to ``state_pool[state_indices[b]]``; with ``snapshot_rows`` / ``snapshot_chunks`` the + state after its first ``64 * snapshot_chunks[b]`` tokens also goes to + ``state_pool[snapshot_rows[b]]`` (``-1`` skips). + """ + + ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + # The tp1 context server holds all 96 KDA heads; generation servers do not prefill. + CELLS = ( + Cell( + "tp1 context: one fresh 300-token prompt, Q/K/V/beta chunk-padded by the post-conv " + "step and g not", + dict(heads=96, seq_lens=(300,), fresh=(True,)), + ), + Cell( + "tp1 context: a 1024-token chunk continuing a long prompt from its slot", + dict(heads=96, seq_lens=(1024,), fresh=(False,)), + ), + Cell( + "tp1 context: four packed prompts, fresh and continued, with partial chunks", + dict(heads=96, seq_lens=(64, 2, 129, 63), fresh=(True, False, True, False)), + ), + Cell( + "tp1 context: three packed prompts with in-forward state snapshots (one skipped)", + dict( + heads=96, + seq_lens=(256, 200, 130), + fresh=(False, True, False), + snapshot_chunks=(2, -1, 1), + ), + ), + ) + note = """ + Kimi K3 calls it from the KDA layer's prefill (K3KDAKernelDispatch.prefill_chunk_kda) with + K = V = 128, chunk_size 64, safe_gate, lower_bound -5, use_gate_in_kernel and + use_beta_sigmoid_in_kernel, bf16 Q/K/V/g, fp32 beta, fp32 A_log / dt_bias, and the + metadata's int64 cu_seqlens, chunk_indices, varlen_is_aligned and single_sequence_length + (so no layer reads cu_seqlens back from the GPU). A single unaligned sequence arrives + zero-padded to a chunk multiple (Q/K/V/beta; g may stay unpadded): the op rebuilds g's + sentinel tail itself, and output rows past cu_seqlens[-1] are unspecified (the caller crops + them). Equal-length inputs with T not a multiple of 256 take a padded copy (B == 1 only). + + In place: the state_pool rows named by state_indices (and snapshot_rows). The caller zeroes + fresh rows first; continuation chunks reuse the slot in order. The returned o is a view of a + process-wide scratch buffer keyed by shape: the next call with the same shape overwrites it. + Process-wide caches also hold the compiled kernels (keyed by layout and constexpr flags), the + intermediate workspaces and padded-input buffers; warm each shape before graph capture. + Launches go to the current stream. Importing the entry needs CuTe DSL and FlashInfer. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + scale: float, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = True, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + snapshot_rows: Optional[torch.Tensor] = None, + snapshot_chunks: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return prefill_ops.kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + state_pool=state_pool, + state_indices=state_indices, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_gate_in_kernel=use_gate_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, + snapshot_rows=snapshot_rows, + snapshot_chunks=snapshot_chunks, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + scale: float, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = True, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + snapshot_rows: Optional[torch.Tensor] = None, + snapshot_chunks: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """fp64 token-by-token recurrence: ``(o, state_pool after the call)``; nothing is written. + + ``o`` rows past ``cu_seqlens[-1]`` are zero (the op leaves them unspecified). + ``chunk_indices``, ``chunk_size``, ``varlen_is_aligned`` and ``single_sequence_length`` + are scheduling hints and do not enter the result. + """ + heads, key_dim = q.shape[2], q.shape[3] + out = torch.zeros(v.shape, dtype=torch.float64, device=v.device) + pool = state_pool.clone() + slots = state_indices.tolist() + snaps = ( + [-1] * len(slots) + if snapshot_chunks is None + else list(zip(snapshot_rows.tolist(), snapshot_chunks.tolist())) + ) + bound = _DEFAULT_LOWER_BOUND if lower_bound is None else lower_bound + for b, (row, lo, hi) in enumerate(_sequences(q, cu_seqlens)): + qd = q[row, lo:hi].double() * scale + kd = k[row, lo:hi].double() + vd = v[row, lo:hi].double() + gd = g[row, lo:hi].double() + bd = beta[row, lo:hi].double() + if use_gate_in_kernel: + if dt_bias is not None: + gd = gd + dt_bias.double().view(heads, key_dim) + rate = A_log.double().exp()[:, None] + if safe_gate: + gd = bound * torch.sigmoid(rate * gd) + else: + gd = -rate * torch.nn.functional.softplus(gd) + if use_beta_sigmoid_in_kernel: + bd = torch.sigmoid(bd) + decay = gd.exp() + s = state_pool[slots[b]].double() + snap = snaps[b] + for t in range(hi - lo): + if snap != -1 and snap[1] >= 0 and t == chunk_size * snap[1]: + pool[snap[0]] = s.to(pool.dtype) + s = s * decay[t][:, None, :] + err = vd[t] - torch.einsum("hvk,hk->hv", s, kd[t]) + s = s + bd[t][:, None, None] * err[:, :, None] * kd[t][:, None, :] + out[row, lo + t] = torch.einsum("hvk,hk->hv", s, qd[t]) + pool[slots[b]] = s.to(pool.dtype) + return out.to(v.dtype), pool + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + scale: float, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = True, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + snapshot_rows: Optional[torch.Tensor] = None, + snapshot_chunks: Optional[torch.Tensor] = None, + ) -> None: + spans = _sequences(q, cu_seqlens) + lengths = [hi - lo for _, lo, hi in spans] + slots = state_indices.tolist() + # Two sequences on one slot race on its row. + assert len(set(slots)) == len(slots), "state_indices must name one slot per sequence" + if cu_seqlens is not None: + # The hints replace device reads of cu_seqlens: a wrong one runs the mask-free + # kernel on a partial chunk, or pads at the wrong length, without an error. + if single_sequence_length is not None: + assert len(lengths) == 1 and lengths[0] == single_sequence_length, ( + "single_sequence_length must be the one sequence's cu_seqlens length" + ) + if varlen_is_aligned is not None: + assert varlen_is_aligned == all(n % chunk_size == 0 for n in lengths), ( + "varlen_is_aligned must say whether every sequence is a chunk multiple" + ) + if chunk_indices is not None: + expected = [ + (s, c) for s, n in enumerate(lengths) for c in range(-(-n // chunk_size)) + ] + assert chunk_indices.tolist() == [list(e) for e in expected], ( + "chunk_indices must list (sequence, chunk) of every chunk of cu_seqlens" + ) + if snapshot_chunks is not None: + rows = snapshot_rows.tolist() + for row, chunks, n in zip(rows, snapshot_chunks.tolist(), lengths): + if chunks < 0: + continue + assert chunks < -(-n // chunk_size), "a snapshot must fall inside its sequence" + assert row not in slots, "a snapshot row must not be a live state row" + + +kda_prefill = KdaPrefill() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/text_config.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/text_config.py new file mode 100644 index 000000000000..f5f5c0da8ee7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/text_config.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 on Rubin (SM107): reading the text config. + +The text config of a vision-language config, the layer schedule (which layers are KDA and +which MLA), and the SiTU activation's betas. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import PretrainedConfig + + +def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: + """Return the finite SiTu betas required by the routed-expert kernels.""" + config_situ_beta = getattr(cfg, "activation_situ_beta", None) + situ_beta = 1.0 if config_situ_beta is None else config_situ_beta + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + if situ_linear_beta is None: + raise ValueError( + "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " + "None means an identity linear branch that the fused kernels cannot represent." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + ) + return float(situ_beta), float(situ_linear_beta) + + +# --------------------------------------------------------------------------- +# Config helpers. +# --------------------------------------------------------------------------- + + +def _get_text_config(pretrained_config: "PretrainedConfig"): + """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" + if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( + not hasattr(pretrained_config, "linear_attn_config") + and hasattr(pretrained_config, "text_config") + ): + return pretrained_config.text_config + return pretrained_config + + +def _is_kda_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] + + +def _is_mla_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/weights.py new file mode 100644 index 000000000000..1bd53290aab1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/weights.py @@ -0,0 +1,1172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 on Rubin (SM107): checkpoint mapping, loading and sharding. + +The model loads through two mixins, each the first base of the class it serves: + +* ``KimiK3DefaultWeights`` (of ``ModelingV2KimiK3Nvfp4Sm107Default``): builds the vision tower + if it was deferred, loads it, and hands the ``language_model.*`` keys, prefix stripped, to + the text model. +* ``KimiK3DefaultTextWeights`` (of ``KimiK3DefaultForCausalLM``): the checkpoint name plan and + its key check, the trunk parameters, the checkpoint FP8 attention projections, and the + routed experts. The checkpoint is streamed, and every rank keeps only its share: its + attention heads (under Helix CP its chunk of the ``v_b_proj`` heads), its slice of the + tensor-parallel MLPs, and its slice of each MoE layer's expert bank (whole experts under + MoE EP; under MoE TP the intermediate shard of every expert, so every expert tensor is + read). + +The module-level helpers name the checkpoint keys of fused projections, cut a rank's slice +of a tensor, and read FP8 block-scaled projections (dequantized for a BF16 parameter). +""" + +from __future__ import annotations + +import copy +import json +import math +import os +import threading +import weakref +from contextlib import ExitStack +from typing import Dict, List, Optional, Set, Tuple + +import torch +from safetensors import safe_open +from torch import nn + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_kimi_k25 import DISAGG, _has_meta_tensors +from tensorrt_llm._torch.models.modeling_utils import filter_weights, run_concurrently +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._utils import is_sm_100f +from tensorrt_llm.logger import logger + +from .fp8_weight_read import ( + _convert_moe_mlps_to_fp8_weight_read, + _Fp8BlockScaleWeightReadLinear, + _has_weights, + _is_fp8_block_projection, + _swap_linear_to_fp8_weight_read, +) +from .kda import K3KDALinearAttention +from .layers import Bf16Linear, Fp8BlockScaleLinear, shard_range +from .moe_runtime import K3MoERuntime + +_KIMI_K3_MLA_DERIVED_PARAM_SUFFIXES = ( + ".self_attn.mixer.k_b_proj_trans", + ".self_attn.mixer.v_b_proj", + ".self_attn.mixer.k_b_proj_trans_scale", + ".self_attn.mixer.v_b_proj_scale", + ".self_attn.mixer.k_b_proj_trans_dequant", + ".self_attn.mixer.v_b_proj_dequant", +) + + +# --------------------------------------------------------------------------- +# Dense / shared-expert MLP: fused [gate | up] layout (``GatedMlp``). +# +# The HF checkpoint stores separate ``gate_proj`` / ``up_proj`` tensors; +# ``load_weights`` row-concatenates them into ``gate_up_proj`` (see +# ``_gate_up_ckpt_keys``), replacing two GEMMs + torch.cat with one GEMM. +# --------------------------------------------------------------------------- + + +_GATE_UP_FUSED_SUFFIX = ".gate_up_proj.weight" + + +def _gate_up_ckpt_keys(fused_key: str) -> Tuple[str, str]: + """Checkpoint ``(gate_proj, up_proj)`` keys whose row-concat loads the + fused ``gate_up_proj`` parameter named by ``fused_key``.""" + return ( + fused_key.replace(_GATE_UP_FUSED_SUFFIX, ".gate_proj.weight"), + fused_key.replace(_GATE_UP_FUSED_SUFFIX, ".up_proj.weight"), + ) + + +def _shard_dim( + tensor: torch.Tensor, dim: int, tp_size: int, tp_rank: int, device: torch.device +) -> torch.Tensor: + """Rank ``tp_rank``'s contiguous ``1 / tp_size`` of ``tensor`` along ``dim`` on ``device`` + (``layers.shard_range``); the whole tensor for one rank, a missing dim (a 1-D tensor's + input columns) or a width-1 dim. The width must divide by ``tp_size``.""" + if tp_size <= 1 or dim >= tensor.dim() or tensor.shape[dim] == 1: + return tensor.to(device) + width = tensor.shape[dim] + if width % tp_size: + raise ValueError(f"dim {dim} of width {width} does not split over {tp_size} ranks") + start, end = shard_range(width, tp_size, tp_rank) + return tensor.narrow(dim, start, end - start).to(device) + + +def _shard_head_major_param( + name: str, + src: torch.Tensor, + param: torch.nn.Parameter, + *, + kda_tp_size: int, + kda_tp_rank: int, + model_tp_rank: int, +) -> torch.Tensor: + """TP-shard a KDA ``linear_attn`` or shared-expert / dense-MLP + ``down_proj`` checkpoint tensor down to this rank's local slice. + + Returns ``src`` unchanged when it already matches ``param`` (replicated or + single-rank tensors) and for every other name — the caller then routes + MLA head-shards and the shape-match copy. Slices are contiguous + (:func:`_shard_dim`). + + - ``.linear_attn.`` names are KDA head-major tensors (``linear_attn`` + exists only on KDA layers): ``o_proj`` shards its input columns (dim 1), + every other projection its output rows (dim 0), by ``kda_tp_size``. + - a shared-expert / dense-MLP ``down_proj`` is ROW-sharded on its input + columns; the fused ``gate_up_proj`` is row-concatenated and returns + before this helper, so ``down_proj`` is the only half that reaches here. + The TP factor comes from the checkpoint-vs-param shapes; a subgroup + smaller than model TP repeats, so the shard index is ``model_tp_rank`` + modulo the parameter's shard count. + """ + if src.shape == param.shape: + return src + if ".linear_attn." in name: + dim = 1 if name.endswith(".o_proj.weight") else 0 + return _shard_dim(src, dim, kda_tp_size, kda_tp_rank, param.device) + if name.endswith(".down_proj.weight") and (".shared_experts." in name or ".mlp." in name): + assert src.shape[1] % param.shape[1] == 0, ( + f"{name}: checkpoint input dim {src.shape[1]} is not " + f"divisible by param input dim {param.shape[1]}" + ) + tp = src.shape[1] // param.shape[1] + return _shard_dim(src, 1, tp, model_tp_rank % tp, param.device) + return src + + +def _helix_cp_v_b_shard( + v_weight: torch.Tensor, + *, + num_heads_tp_cp: int, + cp_rank: int, +) -> torch.Tensor: + """Select this CP rank's ``v_b_proj`` head chunk from a tp-local KV-B split. + + Under Helix context-parallel (cp_size > 1) ``v_b_proj`` holds only this + rank's 1/cp post-all-to-all head chunk, while ``kv_b_proj`` and + ``k_b_proj_trans`` keep every tp-local head. When cp_size == 1 + (``num_heads_tp_cp`` equals the full tp-local head count) the tensor is + returned unchanged. ``v_weight`` is ``[num_heads_tp, v_head_dim, + kv_lora_rank]``; only its leading head axis is sliced. + """ + if num_heads_tp_cp != v_weight.shape[0]: + lo = cp_rank * num_heads_tp_cp + v_weight = v_weight[lo : lo + num_heads_tp_cp] + return v_weight + + +_FP8_BLOCK_SCALE_SUFFIX = "_scale" +_FP8_BLOCK = 128 + + +def _fp8_block_scale_key(weight_key: str) -> str: + return weight_key + _FP8_BLOCK_SCALE_SUFFIX + + +def _checkpoint_fp8_pair(ckpt_key: str, src: torch.Tensor, weights): + """``(fp8 weight, fp32 128x128 block scale)`` if this tensor is FP8_PB_WO. + + ``None`` for an ordinary (BF16) checkpoint tensor. Raises when an FP8 + tensor has no companion scale rather than letting the caller fall through + to a cast that would silently drop it. + """ + if src.dtype != torch.float8_e4m3fn: + return None + scale_key = _fp8_block_scale_key(ckpt_key) + if scale_key not in weights and scale_key + "_inv" in weights: + scale_key += "_inv" + if scale_key not in weights: + raise KeyError( + f"Kimi K3: {ckpt_key} is FP8 E4M3 but has no {scale_key} or {scale_key}_inv; refusing " + "to load it as if it were unquantized." + ) + scale = _materialize(weights[scale_key]).float() + # The checkpoint stores the block scale 4-D as + # [ceil(N/128), 1, ceil(K/128), 1]. Normalize here, once, so every + # consumer sees the plain 2-D [n_blocks_m, n_blocks_k] that both the + # dequantization below and deep_gemm's transform_sf_into_required_layout + # expect -- the latter asserts on rank and gave an unhelpful + # ``assert sf.dim() == ...`` when handed the raw 4-D tensor. + if scale.dim() == 4: + scale = scale.reshape(scale.shape[0], scale.shape[2]) + return src, scale + + +def _dequantize_fp8_block_scaled(ckpt_key: str, src: torch.Tensor, weights) -> torch.Tensor: + """Apply checkpoint scales for projections explicitly configured as BF16. + + Checkpoint-native FP8 attention bypasses this loader. This handles an + excluded projection or another unquantized trunk parameter whose source + tensor is quantized; casting the codes alone would discard its scales. + """ + pair = _checkpoint_fp8_pair(ckpt_key, src, weights) + if pair is None: + return src + src, scale = pair + out = src.to(torch.float32) + if scale.numel() == 1: + return (out * scale.reshape(())).to(torch.bfloat16) + expanded = scale.repeat_interleave(_FP8_BLOCK, dim=0).repeat_interleave(_FP8_BLOCK, dim=1) + if expanded.shape[0] < out.shape[0] or expanded.shape[1] < out.shape[1]: + raise ValueError( + f"Kimi K3: {_fp8_block_scale_key(ckpt_key)} covers " + f"{tuple(expanded.shape)} but {ckpt_key} is {tuple(out.shape)}." + ) + return (out * expanded[: out.shape[0], : out.shape[1]]).to(torch.bfloat16) + + +def _shards_itself(linear: nn.Module) -> bool: + """A ``layers`` projection that takes the full checkpoint tensors and keeps its own TP shard.""" + return ( + isinstance(linear, Fp8BlockScaleLinear) + and linear.tp_mode is not None + and linear.tp_size > 1 + ) + + +def _full_projection_shape(linear: nn.Module) -> Tuple[int, int]: + """``(out, in)`` of the whole checkpoint projection a TP shard ``linear`` loads from.""" + if isinstance(linear, Fp8BlockScaleLinear): + return (linear.full_out_features, linear.full_in_features) + # Main's Linear (KDA / MLA projections not yet on core): the shape its tp_mode splits. + shape = [linear.out_features, linear.in_features] + if linear.tp_mode is not None: + shape[0 if linear.tp_mode == TensorParallelMode.COLUMN else 1] *= linear.tp_size + return tuple(shape) + + +def _load_fp8_projection(linear: nn.Module, weight: torch.Tensor, scale: torch.Tensor) -> None: + """Copy checkpoint E4M3 codes and their FP32 128 x 128 scale grid into an FP8 projection + (its own TP shard of them where it shards).""" + if isinstance(linear, Fp8BlockScaleLinear): + linear.load(weight, scale) + else: + linear.load_weights([{"weight": weight, "weight_scale": scale}]) + + +def _projection_shard( + linear: nn.Module, tensor: torch.Tensor, scale_span: Optional[int] = None +) -> torch.Tensor: + """This rank's slice of a full checkpoint tensor (or with ``scale_span`` its block-scale + grid) of a TP-sharded projection, on the projection's device.""" + if isinstance(linear, (Bf16Linear, Fp8BlockScaleLinear)): + return linear.shard(tensor, scale_span=scale_span) + kwargs = {} if scale_span is None else {"scale_span": scale_span} + return linear.load_shard(tensor, device=linear.weight.device, **kwargs) + + +def _materialize(value) -> torch.Tensor: + """Materialize a (possibly lazy safetensors slice) weight value.""" + if isinstance(value, torch.Tensor): + return value + # ``[:]`` is how a lazy slice is realized, but it is invalid on a 0-dim + # entry. The NVFP4 checkpoint stores weight_scale_2 / input_scale as + # scalars, so realize those entries with ``[()]`` instead. + get_shape = getattr(value, "get_shape", None) + if get_shape is not None and len(get_shape()) == 0: + return value[()] + return value[:] + + +class KimiK3DefaultTextWeights: + """Checkpoint loading of the text model: name plan, trunk, FP8 projections, expert slices. + + The first base of ``KimiK3DefaultForCausalLM``, so its ``load_weights`` resolves before + ``SpecDecOneEngineForCausalLM``'s. + """ + + # ------------------------------------------------------------------ + # Weight loading (streams the 1.5TB checkpoint; only the rank-local + # expert slice of each MoE layer is kept: whole experts under MoE EP, + # the intra-expert intermediate shard of ALL experts under MoE TP — + # in the TP case every expert tensor is read and sliced, so expect a + # correspondingly longer load). + # ------------------------------------------------------------------ + + def _attention_fp8_linears(self) -> list[tuple[str, nn.Module, nn.Module]]: + """Collect checkpoint-backed FP8 attention projections for weight loading. + + The loader uses this list to enumerate required weight/scale keys, exclude + these parameters from ordinary trunk loading, and load FP8 codes and scales + together. Derived QKVG fusion is excluded because it is built after loading. + + Returns: + Tuples of runtime module path, linear, and owning attention module. + """ + projections = [] + for index, layer in enumerate(self.model.layers): + if not _has_weights(layer): + continue + if layer.is_kda: + attention = layer.linear_attn + scope = f"model.layers.{index}.linear_attn" + else: + attention = layer.self_attn.mixer + scope = f"model.layers.{index}.self_attn.mixer" + for attr, module in attention.named_children(): + if attr == "qkvg_proj": + # Decode fusion is derived from separately loaded checkpoint projections. + continue + if _is_fp8_block_projection(module): + projections.append((f"{scope}.{attr}", module, attention)) + return projections + + def _load_attention_fp8(self, weights: Dict[str, torch.Tensor], prefix: str) -> int: + """Load checkpoint FP8 codes/scales without a BF16 projection intermediate.""" + projections = self._attention_fp8_linears() + for name, linear, attention in projections: + # Both attention implementations use self_attn keys in the checkpoint. + key = ( + prefix + + name.replace(".linear_attn.", ".self_attn.").replace( + ".self_attn.mixer.", ".self_attn." + ) + + ".weight" + ) + attr = name.rsplit(".", 1)[1] + if attr == "kv_a_proj_with_mqa" and attention.fuse_qkv_a_proj: + # The runtime fuses Q-A/KV-A, but the checkpoint stores them separately. + # Construction checks that the fusion boundary aligns with FP8 blocks. + q_key = key.replace("kv_a_proj_with_mqa.weight", "q_a_proj.weight") + pairs = [] + for part_key, rows in ( + (q_key, attention.q_lora_rank), + (key, attention.kv_lora_rank + attention.qk_rope_head_dim), + ): + part = _materialize(weights[part_key]) + pair = _checkpoint_fp8_pair(part_key, part, weights) + if pair is None: + raise ValueError( + f"{part_key}: quant config declares FP8 but tensor is not E4M3" + ) + codes, scales = pair + expected = (rows, attention.hidden_size) + if tuple(codes.shape) != expected: + raise ValueError( + f"{part_key}: checkpoint shape {tuple(codes.shape)} != {expected}" + ) + if tuple(scales.shape) != tuple(math.ceil(d / _FP8_BLOCK) for d in expected): + raise ValueError(f"{part_key}: scale shape does not cover {expected}") + pairs.append(pair) + _load_fp8_projection( + linear, + torch.cat([pair[0] for pair in pairs], dim=0), + torch.cat([pair[1] for pair in pairs], dim=0), + ) + continue + source = weights[key] + shape = ( + tuple(source.shape) + if isinstance(source, torch.Tensor) + else tuple(source.get_shape()) + ) + local_shape = (linear.out_features, linear.in_features) + scale_slice = None + if isinstance(attention, K3KDALinearAttention) and not _shards_itself(linear): + # KDA linears already have rank-local dimensions. Slice the lazy source + # before materializing it, and apply the same shard to its block scales. + split_dim = 1 if attr == "o_proj" else 0 + tp_size = attention._kda_tp_size + tp_rank = attention._kda_tp_rank + expected_shape = list(local_shape) + expected_shape[split_dim] *= tp_size + if shape != tuple(expected_shape): + raise ValueError(f"{key}: checkpoint shape {shape} != {tuple(expected_shape)}") + if tp_size > 1: + width = local_shape[split_dim] + start, end = tp_rank * width, (tp_rank + 1) * width + if start % _FP8_BLOCK or end % _FP8_BLOCK: + raise ValueError(f"{key}: TP slice [{start}:{end}] is not block-aligned") + indices = [slice(None), slice(None)] + indices[split_dim] = slice(start, end) + source = source[tuple(indices)] + scale_indices = [slice(None), slice(None)] + scale_indices[split_dim] = slice( + start // _FP8_BLOCK, math.ceil(end / _FP8_BLOCK) + ) + scale_slice = tuple(scale_indices) + pair = _checkpoint_fp8_pair(key, _materialize(source), weights) + if pair is None: + raise ValueError( + f"{key}: quant config declares FP8 but checkpoint tensor is not E4M3" + ) + weight, scale = pair + full_scale_shape = tuple(math.ceil(dim / _FP8_BLOCK) for dim in shape) + if tuple(scale.shape) != full_scale_shape: + raise ValueError(f"{key}: scale shape {tuple(scale.shape)} != {full_scale_shape}") + if scale_slice is not None: + scale = scale[scale_slice] + if isinstance(attention, K3KDALinearAttention) and not _shards_itself(linear): + if tuple(scale.shape) != tuple(math.ceil(dim / _FP8_BLOCK) for dim in local_shape): + raise ValueError(f"{key}: scale shard does not cover local shape {local_shape}") + _load_fp8_projection(linear, weight, scale) + elif attr == "kv_b_proj": + # KV-B also supplies decode absorption weights; reorder codes and scales + # together and retain raw FP8 copies for the absorption kernels. + self._load_mla_fp8_kv_b(attention, weight, scale) + else: + # The projection shards the full checkpoint tensors itself. + expected_shape = _full_projection_shape(linear) + if shape != expected_shape: + raise ValueError(f"{key}: checkpoint shape {shape} != {expected_shape}") + _load_fp8_projection(linear, weight, scale) + return len(projections) + + @staticmethod + def _load_mla_fp8_kv_b(attention: nn.Module, weight: torch.Tensor, scale: torch.Tensor) -> None: + """Reorder checkpoint KV-B heads and their scale grid together.""" + linear = attention.kv_b_proj + h, n, v, k = ( + attention.num_heads_tp, + attention.qk_nope_head_dim, + attention.v_head_dim, + attention.kv_lora_rank, + ) + if any(dim % _FP8_BLOCK for dim in (n, v, k)): + raise ValueError( + "Kimi MLA FP8 absorption requires block-aligned head and latent dimensions" + ) + expected_shape = (h * linear.tp_size * (n + v), k) + if tuple(weight.shape) != expected_shape: + raise ValueError(f"Kimi MLA KV-B shape {tuple(weight.shape)} != {expected_shape}") + local = _projection_shard(linear, weight).reshape(h, n + v, k) + local_scale = _projection_shard(linear, scale, scale_span=_FP8_BLOCK).reshape( + h, (n + v) // _FP8_BLOCK, k // _FP8_BLOCK + ) + key_weight, value_weight = local.split((n, v), dim=1) + key_scale, value_scale = local_scale.split((n // _FP8_BLOCK, v // _FP8_BLOCK), dim=1) + linear.weight.data.copy_( + torch.cat((key_weight.reshape(h * n, k), value_weight.reshape(h * v, k))) + ) + linear.weight_scale.data.copy_( + torch.cat((key_scale.flatten(0, 1), value_scale.flatten(0, 1))) + ) + attention.k_b_proj_trans.data.copy_(key_weight.transpose(1, 2)) + attention.k_b_proj_trans_scale.data.copy_(key_scale.transpose(1, 2)) + if attention.mapping.has_cp_helix(): + expanded = value_scale.repeat_interleave(_FP8_BLOCK, -2).repeat_interleave( + _FP8_BLOCK, -1 + ) + # The MLA layer owns the V-up-before-exchange flag; bind through it (its + # transform_weights rebuilds the same weight from kv_b_proj). + attention._bind_helix_v_up_weight((value_weight.float() * expanded).to(torch.bfloat16)) + value_weight = _helix_cp_v_b_shard( + value_weight, + num_heads_tp_cp=attention.num_heads_tp_cp, + cp_rank=attention.mapping.cp_rank, + ) + value_scale = _helix_cp_v_b_shard( + value_scale, + num_heads_tp_cp=attention.num_heads_tp_cp, + cp_rank=attention.mapping.cp_rank, + ) + # Absorption must keep raw checkpoint codes independent of Linear's + # backend-specific weight transformations. + attention.v_b_proj = nn.Parameter(value_weight.clone().contiguous(), requires_grad=False) + attention.v_b_proj_scale.data.copy_(value_scale) + for target, codes, scales in ( + ( + attention.k_b_proj_trans_dequant, + key_weight.transpose(1, 2), + key_scale.transpose(1, 2), + ), + (attention.v_b_proj_dequant, value_weight, value_scale), + ): + if target is not None: + expanded = scales.repeat_interleave(_FP8_BLOCK, -2).repeat_interleave( + _FP8_BLOCK, -1 + ) + target.data.copy_((codes.float() * expanded).to(target.dtype)) + + def _trunk_parameters(self) -> Dict[str, torch.nn.Parameter]: + """Named parameters of the trunk only. Spec-dec draft modules + (e.g. the DFlash drafter attached by SpecDecOneEngineForCausalLM) + live in a separate checkpoint loaded by + ModelLoader.load_draft_weights, not in the target checkpoint. MLA + K/V absorb Parameters are derived by the KV-B loader and are likewise + excluded from checkpoint jobs.""" + fp8_prefixes = tuple(name + "." for name, _, _ in self._attention_fp8_linears()) + return { + name: param + for name, param in self.named_parameters() + if not name.startswith("draft_model.") + and not name.startswith(fp8_prefixes) + and ".linear_attn.qkvg_proj." not in name + and not name.endswith(_KIMI_K3_MLA_DERIVED_PARAM_SUFFIXES) + } + + def checkpoint_name_plan( + self, prefix: str + ) -> Tuple[Dict[str, str], Set[str], List[Tuple[int, K3MoERuntime, str]]]: + """Return ``(name_map, expected_keys, expert_jobs)``. + + ``name_map`` maps every model parameter name to its checkpoint key + (for fused ``gate_up_proj`` parameters the mapped key is virtual; + the two real per-half keys come from ``_gate_up_ckpt_keys``); + ``expected_keys`` additionally covers the rank-local per-expert MXFP4 + tensors; ``expert_jobs`` lists ``(layer_idx, moe_module, key_base)`` + for backend-owned expert slots. Exposed separately so the weight-name + mapping can be dry-run without touching any tensor data. + """ + params = self._trunk_parameters() + expected_keys = set() + name_map: Dict[str, str] = {} + for name in params: + # The routed experts' weights are loaded expert by expert (expert_jobs); + # they do not have one-to-one checkpoint parameter names. + if ".routed_experts.backend." in name: + continue + if name == "lm_head.weight": + ckpt_key = prefix + "lm_head.weight" + else: + if ".linear_attn." in name: + ckpt_key = prefix + name.replace(".linear_attn.", ".self_attn.") + else: + # MLA retains its runtime/mixer hierarchy; the checkpoint + # has no intermediate ``mixer`` scope. + ckpt_key = prefix + name.replace(".self_attn.mixer.", ".self_attn.") + name_map[name] = ckpt_key + if name.endswith(_GATE_UP_FUSED_SUFFIX): + # Fused [gate | up] MLP layout (dense mlp / shared_experts): + # the checkpoint stores two separate tensors. + expected_keys.update(_gate_up_ckpt_keys(ckpt_key)) + else: + expected_keys.add(ckpt_key) + + # Backend-owned expert slots (per-expert checkpoint tensors; the + # rank-local id range — an EP slice of whole experts, or ALL experts + # when the routed MoE is TP-sharded (moe_ep=1 -> ids 0..num_experts)). + expert_jobs = [] + for layer_idx, layer in enumerate(self.model.layers): + if not getattr(layer, "is_moe", False) or not _has_weights(layer): + continue + moe = layer.block_sparse_moe + base = f"{prefix}model.layers.{layer_idx}.block_sparse_moe.experts" + # Required suffixes only: ``optional_suffixes`` are per-checkpoint + # (an AWQ pre-quant scale), and demanding them here would reject + # every checkpoint of the same layout that does not carry them. + # Unmatched expert keys only warn, and this validator already + # exempts expert keys from that, so the present ones stay quiet. + kinds = moe.expert_ckpt_spec.kinds + for expert_idx in moe.local_expert_ids: + for w in ("w1", "w2", "w3"): + for kind in kinds: + expected_keys.add(f"{base}.{expert_idx}.{w}.{kind}") + expert_jobs.append((layer_idx, moe, base)) + for name, _, _ in self._attention_fp8_linears(): + key = ( + prefix + + name.replace(".linear_attn.", ".self_attn.").replace( + ".self_attn.mixer.", ".self_attn." + ) + + ".weight" + ) + expected_keys.update((key, _fp8_block_scale_key(key))) + for index, layer in enumerate(self.model.layers): + if getattr(layer, "is_kda", True) or not _has_weights(layer): + continue + attention = layer.self_attn.mixer + if getattr(attention, "fuse_qkv_a_proj", False): + key = f"{prefix}model.layers.{index}.self_attn.q_a_proj.weight" + expected_keys.add(key) + if _is_fp8_block_projection(attention.kv_a_proj_with_mqa): + expected_keys.add(_fp8_block_scale_key(key)) + return name_map, expected_keys, expert_jobs + + def load_weights(self, weights: Dict[str, torch.Tensor]) -> None: + prefix = "language_model." if any(k.startswith("language_model.") for k in weights) else "" + params = self._trunk_parameters() + name_map, expected_keys, expert_jobs = self.checkpoint_name_plan(prefix) + + self._validate_checkpoint_keys(weights, expected_keys, prefix) + num_params = self._load_trunk_params(weights, params, name_map) + num_params += self._load_attention_fp8(weights, prefix) + self._load_expert_slices(weights, expert_jobs) + self._finalize_weight_load(num_params, len(expert_jobs)) + device = next(self.parameters()).device + if device.type == "cuda": + # Lazy source mappings are load-scoped; finish nonblocking H2D + # work before the caller can release the weights container. + torch.cuda.synchronize(device) + + def _validate_checkpoint_keys( + self, weights: Dict[str, torch.Tensor], expected_keys: Set[str], prefix: str + ) -> None: + """Key-set validation (both directions): every expected key must be + present; unmatched checkpoint keys (beyond the expected leftovers) + only warn.""" + ckpt_keys = set(weights.keys()) + relevant_ckpt_keys = { + k + for k in ckpt_keys + if not (k.startswith("vision_tower.") or k.startswith("mm_projector.")) + } + missing = sorted( + key + for key in expected_keys - ckpt_keys + if not (key.endswith(".weight_scale") and key + "_inv" in ckpt_keys) + ) + if missing: + raise KeyError( + f"Kimi K3 load_weights: {len(missing)} expected checkpoint " + f"keys are missing, e.g. {missing[:10]}" + ) + unexpected = relevant_ckpt_keys - expected_keys + # Non-local experts and (in layer-truncated debug mode) extra layers + # are expected leftovers. + surprising = sorted( + k + for k in unexpected + if ".block_sparse_moe.experts." not in k and not k.startswith(f"{prefix}model.layers.") + ) + if surprising: + logger.warning( + f"Kimi K3 load_weights: {len(surprising)} unmatched " + f"checkpoint keys, e.g. {surprising[:10]}" + ) + + def _load_trunk_params( + self, + weights: Dict[str, torch.Tensor], + params: Dict[str, torch.nn.Parameter], + name_map: Dict[str, str], + ) -> int: + """Load every non-expert trunk parameter concurrently (with the + per-parameter TP-shard / pad / fuse conversions) and return the + number of parameters loaded.""" + # The checkpoint stores every MLA KV-B head as interleaved [K | V] + # rows. Runtime keeps one DeepSeek-style [all K | all V] parameter + # instead, so context can project directly into the FMHA layout and + # absorbed decode can take zero-copy K/V views. + mla_mixers = [ + layer.self_attn.mixer + for layer in self.model.layers + if not getattr(layer, "is_kda", True) and _has_weights(layer) + ] + mla_kv_b_mixers = {id(mixer.kv_b_proj.weight): mixer for mixer in mla_mixers} + mla_fused_a_mixers = { + id(mixer.kv_a_proj_with_mqa.weight): mixer + for mixer in mla_mixers + if getattr(mixer, "fuse_qkv_a_proj", False) + } + mla_head_shard_linears = {} + for mixer in mla_mixers: + mla_head_shard_linears[id(mixer.q_b_proj.weight)] = mixer.q_b_proj + mla_head_shard_linears[id(mixer.o_proj.weight)] = mixer.o_proj + g_proj = getattr(mixer, "g_proj", None) + if g_proj is not None: + mla_head_shard_linears[id(g_proj.weight)] = g_proj + + device = next(self.parameters()).device + + # MLP TP shard index. A dense MLP whose intermediate size does not + # divide model TP uses a smaller repeated TP subgroup, so its local + # shard rank is model tp_rank modulo the parameter's shard count. + # Under helix the modules were sharded against the repurposed + # mapping; the restored CP original's tp_rank is 0 on every rank. + model_tp_rank = ( + self._repurposed_tp_mapping.tp_rank + if self._repurposed_tp_mapping is not None + else self.model_config.mapping.tp_rank + ) + # KDA head-shard (attention-DP off): rank r loads head rows/cols + # [r*local : (r+1)*local] of every head-major KDA tensor. + kda_tp_size, kda_tp_rank = 1, 0 + for layer in self.model.layers: + if getattr(layer, "is_kda", False): + kda_tp_size = layer.linear_attn._kda_tp_size + kda_tp_rank = layer.linear_attn._kda_tp_rank + break + + def load_param(name: str, param: torch.nn.Parameter) -> None: + if device.type == "cuda": + torch.cuda.set_device(device) + if name.endswith(_GATE_UP_FUSED_SUFFIX): + # Row-concat the checkpoint's separate gate_proj / up_proj + # tensors into the fused [gate | up] parameter. + gate_key, up_key = _gate_up_ckpt_keys(name_map[name]) + inter = param.shape[0] // 2 + # Materialize + FP8-block-dequant each half (this branch reads + # its own sources, so it needs the same FP8 handling as the + # single-tensor path below), then column(TP)-shard and + # row-concat. TP factor from the checkpoint-vs-param shapes; a + # subgroup smaller than model TP repeats, so the shard index is + # model tp_rank modulo the half's shard count. + gate_full = _dequantize_fp8_block_scaled( + gate_key, _materialize(weights[gate_key]), weights + ) + tp = gate_full.shape[0] // inter + # _shard_dim returns the whole tensor when tp <= 1, so the rank needs no + # tp > 1 guard. + rk = model_tp_rank % tp + gate = _shard_dim(gate_full, 0, tp, rk, param.device) + up = _shard_dim( + _dequantize_fp8_block_scaled(up_key, _materialize(weights[up_key]), weights), + 0, + tp, + rk, + param.device, + ) + if gate.shape != (inter, param.shape[1]) or up.shape != gate.shape: + raise ValueError( + f"{name}: checkpoint gate/up shapes " + f"{tuple(gate.shape)} / {tuple(up.shape)} do not " + f"concat to param shape {tuple(param.shape)}" + ) + param.data[:inter].copy_(gate.to(param.dtype)) + param.data[inter:].copy_(up.to(param.dtype)) + return + src = _materialize(weights[name_map[name]]) + src = _dequantize_fp8_block_scaled(name_map[name], src, weights) + fused_a = mla_fused_a_mixers.get(id(param)) + if fused_a is not None: + q_key = name_map[name].replace("kv_a_proj_with_mqa.weight", "q_a_proj.weight") + q_weight = _dequantize_fp8_block_scaled( + q_key, _materialize(weights[q_key]), weights + ) + if q_weight.shape != (fused_a.q_lora_rank, fused_a.hidden_size) or src.shape != ( + fused_a.kv_lora_rank + fused_a.qk_rope_head_dim, + fused_a.hidden_size, + ): + raise ValueError( + f"{name}: checkpoint Q-A/KV-A shapes do not match MLA dimensions" + ) + src = torch.cat((q_weight, src), dim=0) + if name == "lm_head.weight": + # LMHead is vocab-sharded (TP column) + gathered; its + # load_weights shards the full checkpoint tensor. + self.lm_head.load_weights(weights=[{"weight": src}]) + return + mla_mixer = mla_kv_b_mixers.get(id(param)) + if mla_mixer is not None: + h = mla_mixer.num_heads_tp + n = mla_mixer.qk_nope_head_dim + v = mla_mixer.v_head_dim + kv = mla_mixer.kv_lora_rank + local = _projection_shard(mla_mixer.kv_b_proj, src).view(h, n + v, kv) + k_weight, v_weight = local.split([n, v], dim=1) + param.data.copy_( + torch.cat( + [ + k_weight.reshape(h * n, kv), + v_weight.reshape(h * v, kv), + ], + dim=0, + ).to(param.dtype) + ) + mla_mixer.k_b_proj_trans.data.copy_(k_weight.transpose(1, 2)) + # Share the loaded BF16 V rows with context GEMM instead of keeping a + # separate copy. Wrap the view to retain v_b_proj's parameter registration. + v_weight = _helix_cp_v_b_shard( + param[h * n :].view(h, v, kv), + num_heads_tp_cp=mla_mixer.num_heads_tp_cp, + cp_rank=mla_mixer.mapping.cp_rank, + ) + mla_mixer.v_b_proj = nn.Parameter(v_weight, requires_grad=False) + if mla_mixer.mapping.has_cp_helix(): + mla_mixer._bind_helix_v_up_weight(param[h * n :].view(h, v, kv)) + return + if name.endswith(".A_log") and src.numel() != param.numel(): + # The checkpoint pads A_log from [num_heads] to [head_dim] + # (e.g. [96] -> [128]); the tail must be zeros. Under KDA + # head-shard TP the param holds this rank's head range + # instead of the full [num_heads]. + assert src.numel() > param.numel(), (name, src.shape) + if kda_tp_size > 1: + lo = kda_tp_rank * param.numel() + src = src[lo : lo + param.numel()] + else: + tail = src[param.numel() :] + if tail.abs().max().item() != 0.0: + raise ValueError( + f"{name}: expected zero padding beyond " + f"{param.numel()} entries, got nonzero tail" + ) + src = src[: param.numel()] + # KDA ``linear_attn`` (head-major) and shared-expert / dense-MLP + # ``down_proj`` shards are shape-derived; replicated tensors and + # every other name pass through unchanged. MLA head-shards and the + # shape-match copy are handled by the block below. + src = _shard_head_major_param( + name, + src, + param, + kda_tp_size=kda_tp_size, + kda_tp_rank=kda_tp_rank, + model_tp_rank=model_tp_rank, + ) + + if src.shape != param.shape: + # MLA q_b/g/o head-shard: delegate to the same Linear modules + # that own their COLUMN/ROW sharding policy (#17684 removed the + # 96->128 head padding). KV-B is handled above; KDA and + # shared-expert/MLP shards were resolved in the pre-block above. + mla_sharded_linear = mla_head_shard_linears.get(id(param)) + if mla_sharded_linear is not None: + shard = _projection_shard(mla_sharded_linear, src) + if shard.shape != param.shape: + raise ValueError( + f"{name}: MLA shard shape {tuple(shard.shape)} does " + f"not match param shape {tuple(param.shape)}" + ) + param.data.copy_(shard.to(param.dtype)) + return + raise ValueError( + f"{name}: shard/pad result {tuple(src.shape)} != param " + f"shape {tuple(param.shape)}" + ) + param.data.copy_(src.to(param.dtype)) + + param_jobs = [(name, params[name]) for name in name_map] + run_concurrently(load_param, param_jobs, num_workers=8) + + logger.info( + f"Kimi K3: loaded {len(mla_mixers)} MLA KV-B projections in grouped runtime layout" + ) + return len(param_jobs) + + def _load_expert_slices( + self, + weights: Dict[str, torch.Tensor], + expert_jobs: List[Tuple[int, K3MoERuntime, str]], + ) -> None: + """Load the rank-local MXFP4 expert slices of every MoE layer into + the backend expert slots, then verify every slot was filled.""" + device = next(self.parameters()).device + + # Layouts whose per-expert loader only stages its input must have the + # staging containers created before any thread runs, and must be + # finalized once the layer's last slot lands: nothing else on this path + # calls process_weights_after_loading. Which thread completes a layer is + # a race, so the completion test and the "already finalized" bookkeeping + # are one critical section. + finalize_lock = threading.Lock() + finalized_backends = set() + prepared_backends = set() + + def ensure_prepared(moe: K3MoERuntime): + """Prepare a layer's streaming state on its FIRST expert, not up front. + + Preparing every layer before the load starts is what OOM-ed the + MegaMoE CuteDSL backend: it keeps its raw NVFP4 source params as + 0-element placeholders and rematerializes them at full shape here, + so preparing all 92 layers held 92 layers of raw weights at once + instead of the handful actually being filled. Cutlass did not care + because its parameters are allocated either way. + + Bounded lazily instead. The paired shrink already happens per layer + in process_weights_after_loading, so the live set is whatever is + genuinely in flight -- measured at 1 layer per shard file for this + checkpoint (its rank-local experts are 1:1 with files), so ~4 with + 4 loader threads. Correctness does not depend on that layout + though; a checkpoint that split a layer across files would only + raise the peak, not break this. + + The membership add happens AFTER preparing, so the lock-free fast + path can only ever be stale in the safe direction. + """ + spec = moe.expert_ckpt_spec + if not spec.needs_layer_finalize: + return + backend = moe.routed_experts.backend + if id(backend) in prepared_backends: + return + with finalize_lock: + if id(backend) in prepared_backends: + return + moe.prepare_expert_load() + prepared_backends.add(id(backend)) + + def maybe_finalize_layer(moe: K3MoERuntime): + spec = moe.expert_ckpt_spec + if not spec.needs_layer_finalize: + return + backend = moe.routed_experts.backend + with finalize_lock: + loaded = len(getattr(backend, spec.loaded_slots_attr, ())) + if loaded != backend.expert_size_per_partition: + return + if id(backend) in finalized_backends: + return + finalized_backends.add(id(backend)) + # Computes the alphas, interleaves the w2 scales, and (MegaMoE) + # packs into the mega buffers and shrinks the raw source params + # back to placeholders. Paired with ensure_prepared above, that + # pairing is what bounds the per-layer footprint. The per-expert + # drain inside load_streaming_nvfp4_expert separately bounds the + # Cutlass w3_w1 staging, which is per expert rather than per layer. + moe.finalize_expert_load() + + def load_expert( + moe: K3MoERuntime, + base: str, + local_slot_id: int, + expert_idx: int, + get_tensor, + get_optional_tensor, + ): + if device.type == "cuda": + torch.cuda.set_device(device) + backend = moe.routed_experts.backend + # Every route into an expert goes through here (file-grouped, + # split-file and the shared-dict fallback alike), so this is the + # one place preparation has to be hooked. + ensure_prepared(moe) + moe.expert_ckpt_spec.loader( + backend, base, expert_idx, local_slot_id, get_tensor, get_optional_tensor + ) + maybe_finalize_layer(moe) + + def load_experts_from_weights(layer_idx: int, moe: K3MoERuntime, base: str): + del layer_idx + for local_slot_id, expert_idx in enumerate(moe.local_expert_ids): + load_expert( + moe, + base, + local_slot_id, + expert_idx, + lambda key: _materialize(weights[key]), + lambda key: _materialize(weights[key]) if key in weights else None, + ) + + # ---- backend expert slots: file-grouped streaming ---- + # The shared lazy ``weights`` dict keeps every shard mmapped for the + # whole load, so pages it touches cannot be dropped until the load + # ends (fadvise skips mapped pages). The expert slices are ~90 GB of + # DISTINCT pages per rank — with 4 ranks/node that overruns the job + # cgroup and OOM-kills the step (observed repeatedly on GB300 + # trays). Instead, group the rank-local expert tensors by shard file + # and stream each file through a short-lived handle: + # open -> copy -> close (unmap) -> fadvise(DONTNEED). + # The lazy loader records the directory it opened; prefer it over + # ``_name_or_path``, which transformers no longer populates (it is + # empty on transformers 5.x, which silently sent the whole load down + # the fallback below and OOM-killed the step). + ckpt_dir = getattr(weights, "checkpoint_dir", None) or getattr( + self.model_config.pretrained_config, "_name_or_path", None + ) + index_path = os.path.join(ckpt_dir or "", "model.safetensors.index.json") + checkpoint_prefix = getattr(weights, "checkpoint_prefix", "") + + def checkpoint_key(key: str) -> str: + return f"{checkpoint_prefix}{key}" + + if expert_jobs and ckpt_dir and os.path.isfile(index_path): + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + per_file: Dict[str, list] = {} + split_file_jobs = [] + for layer_idx, moe, base in expert_jobs: + del layer_idx + spec = moe.expert_ckpt_spec + for local_slot_id, expert_idx in enumerate(moe.local_expert_ids): + keys = [ + f"{base}.{expert_idx}.{w}.{kind}" + for w in ("w1", "w2", "w3") + for kind in spec.kinds + ] + # An optional tensor this checkpoint does store has to be + # counted here: it decides which files the job opens, and + # ``get_optional_tensor`` below can only read a file that + # ``files`` named. + keys += [ + key + for key in ( + f"{base}.{expert_idx}.{suffix}" for suffix in spec.optional_suffixes + ) + if checkpoint_key(key) in weight_map + ] + files = {weight_map[checkpoint_key(key)] for key in keys} + job = (moe, base, local_slot_id, expert_idx) + if len(files) == 1: + per_file.setdefault(files.pop(), []).append(job) + else: + split_file_jobs.append((job, files)) + + def drop_file_pages(file_name: str): + path = os.path.join(ckpt_dir, file_name) + try: + fd = os.open(path, os.O_RDONLY) + try: + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(fd) + except OSError: + pass + + def load_expert_file(file_name: str, jobs: list): + if device.type == "cuda": + torch.cuda.set_device(device) + path = os.path.join(ckpt_dir, file_name) + with safe_open(path, framework="pt", device="cpu") as fh: + # A job lands here only when every key it needs -- optional + # ones included, see the grouping above -- is in this file, + # so presence in this file is the whole question. + file_keys = set(fh.keys()) + for moe, base, local_slot_id, expert_idx in jobs: + load_expert( + moe, + base, + local_slot_id, + expert_idx, + lambda key: fh.get_tensor(checkpoint_key(key)), + lambda key: ( + fh.get_tensor(checkpoint_key(key)) + if checkpoint_key(key) in file_keys + else None + ), + ) + # Handle closed -> pages unmapped -> the drop takes effect. + drop_file_pages(file_name) + + def load_split_file_expert(job, files): + if device.type == "cuda": + torch.cuda.set_device(device) + with ExitStack() as stack: + handles = { + file_name: stack.enter_context( + safe_open( + os.path.join(ckpt_dir, file_name), framework="pt", device="cpu" + ) + ) + for file_name in files + } + + def get_tensor(key): + source_key = checkpoint_key(key) + return handles[weight_map[source_key]].get_tensor(source_key) + + def get_optional_tensor(key): + source_key = checkpoint_key(key) + # Absent from the index means the checkpoint does not + # store it; present means the grouping above put its + # file in ``handles``. + if source_key not in weight_map: + return None + return handles[weight_map[source_key]].get_tensor(source_key) + + load_expert(*job, get_tensor, get_optional_tensor) + for file_name in files: + drop_file_pages(file_name) + + run_concurrently(load_expert_file, sorted(per_file.items()), num_workers=4) + run_concurrently(load_split_file_expert, split_file_jobs, num_workers=4) + else: + # Falling back is a silent loss of the whole point of the block + # above: the shared lazy dict keeps every shard mapped, which is + # the OOM this streaming path exists to avoid. Say so. + if expert_jobs: + logger.warning( + f"Kimi K3: no safetensors index at '{index_path}', so routed " + "experts are loaded from the shared lazy weight dict instead " + "of being streamed per shard file. Every shard stays mapped " + "for the whole load, which OOM-kills the step at DEP8 scale." + ) + run_concurrently(load_experts_from_weights, expert_jobs, num_workers=4) + + for _, moe, _ in expert_jobs: + spec = moe.expert_ckpt_spec + backend = moe.routed_experts.backend + loaded_slots = getattr(backend, spec.loaded_slots_attr, set()) + expected_slots = set(range(backend.expert_size_per_partition)) + if loaded_slots != expected_slots: + missing_slots = sorted(expected_slots - loaded_slots) + raise RuntimeError( + "Kimi K3 streaming expert loading did not fill all backend " + f"slots; missing {missing_slots[:10]}." + ) + if ( + spec.needs_layer_finalize + and expected_slots + and id(backend) not in finalized_backends + ): + # Unreachable via load_expert (the last slot finalizes), so + # reaching it means the two bookkeeping paths disagree. Guarded + # on expected_slots because a layer that owns no local slots is + # never prepared and so is legitimately never finalized. + raise RuntimeError( + "Kimi K3 streaming expert loading filled every slot but " + "never finalized the layer." + ) + backend._weights_transformed = False + + def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: + """Finalize checkpoint-loaded attention and optional BF16 MLP conversion.""" + for layer in self.model.layers: + if not _has_weights(layer): + continue + if getattr(layer, "is_kda", False): + attention = layer.linear_attn + attention.finalize_decode_weights() + attention._build_mtp_conv_weights() + else: + attention = layer.self_attn.mixer + for name, linear in attention.named_children(): + if _is_fp8_block_projection(linear): + _swap_linear_to_fp8_weight_read(attention, name, (type(linear),)) + fused = getattr(attention, "qkvg_proj", None) + if isinstance(fused, _Fp8BlockScaleWeightReadLinear): + offset = 0 + for name in ("q_proj", "k_proj", "v_proj", "g_proj"): + projection = getattr(attention, name) + projection._fused_projection = (weakref.ref(fused), offset) + offset += projection.out_features + # Repoint to the fused rows now rather than in post_load: + # deferring kept every layer's standalone codes alongside + # its fused copy, ~22.6 GiB at peak across K3's 69 KDA + # layers, which overflows a DEP8 context rank. + projection.transform_weights() + logger.info( + f"Kimi K3: loaded {num_params} parameters and expert slices of {num_moe_layers} layers" + ) + if self._fp8_weight_read_moe_mlp and is_sm_100f(): + converted = _convert_moe_mlps_to_fp8_weight_read( + self.model, + include_fused_gate_up=self.model_config.mapping.enable_attention_dp, + ) + logger.info(f"Kimi K3: quantized {converted} BF16 shared/latent MLP projections to FP8") + + +class KimiK3DefaultWeights: + """Checkpoint loading of the vision-language model: the vision tower, then the text model. + + The first base of ``ModelingV2KimiK3Nvfp4Sm107Default``. + """ + + def load_draft_weights(self, *args, **kwargs): + return self.llm.load_draft_weights(*args, **kwargs) + + def load_weights(self, weights) -> None: + """Load vision + projector + LLM weights from checkpoint.""" + # Create vision encoder if it was skipped during meta init. + # Use the saved VLM config since model_config.pretrained_config + # has been overwritten with text_config by __init__. + if self.mm_encoder is not None and _has_meta_tensors(self.mm_encoder): + logger.info("Recreating deferred vision encoder after MetaInitMode") + self.mm_encoder = None + + if self.mm_encoder is None and not DISAGG: + vision_model_config = copy.copy(self.model_config) + vision_model_config._frozen = False + vision_model_config.pretrained_config = self._vlm_pretrained_config + vision_model_config._frozen = True + self.mm_encoder = self._VISION_MODEL_CLS(vision_model_config) + if self.mm_encoder is not None: + self.mm_encoder.load_weights(weights) + + if any(k.startswith(self._LANG_PREFIX) for k in weights): + lm_weights = filter_weights("language_model", weights) + lm_weights = ConsumableWeightsDict(lm_weights) + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = self._LANG_PREFIX + else: + lm_weights = weights + self.llm.load_weights(lm_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py index c4370e5c1e07..0fc0ec2f7903 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py @@ -8,17 +8,22 @@ uses the built-in Kimi K3 implementation; in ``require`` it raises that error. The checkpoint declares the vision-language wrapper, and the language model's -shape lives in its ``text_config``, so that is what this tree reads. The -targets are text-only: they load no vision tower and refuse image input. +shape lives in its ``text_config``, so that is what this tree reads. The MXFP4 +targets on sm100 are text-only: they load no vision tower and refuse image input. + +The NVFP4 checkpoint on sm107 (Rubin) has one target, ``kimi_k3_nvfp4__sm_107__default``, +which derives its sharding from the mapping and serves the disaggregated +deployments listed in ``_sm107_parallel``. """ from __future__ import annotations from typing import TYPE_CHECKING, Any, Optional, Tuple +from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget, routes +from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig @@ -93,6 +98,102 @@ def _parallel(m) -> Optional[str]: return None +# --------------------------------------------------------------------------- +# Kimi K3 NVFP4 on Rubin (sm107). +# --------------------------------------------------------------------------- + +_SM107 = (10, 7) + +# (layers, hidden, routed experts, experts per token, attention heads, q_lora_rank, +# kv_lora_rank, vocab), read from ``text_config``. +_SM107_SHAPE = (93, 7168, 896, 16, 96, 1536, 512, 163840) +_SM107_SHAPE_FIELDS = ( + "num_hidden_layers", + "hidden_size", + "num_experts", + "num_experts_per_token", + "num_attention_heads", + "q_lora_rank", + "kv_lora_rank", + "vocab_size", +) + + +def _sm107_parallel(mapping) -> str: + """The deployment this mapping is, or a description of it when it is none of them. + + The sm107 target serves one-rank DWDP context servers (``tp1``), generation over + Helix CP8 (``helix8``), over Helix CP16 with the experts split MoE TP4 x EP4 + (``helix16_moetp4ep4``), and attention DP over 16 ranks (``dep16``). + """ + helix = mapping.cp_size > 1 and mapping.has_cp_helix() + if mapping.pp_size != 1: + return f"pp{mapping.pp_size}" + if mapping.world_size == 1: + return "tp1" + if helix and not mapping.enable_attention_dp and mapping.tp_size == 1: + if mapping.cp_size == 8: + return "helix8" + if mapping.cp_size == 16 and (mapping.moe_tp_size, mapping.moe_ep_size) == (4, 4): + return "helix16_moetp4ep4" + if ( + mapping.enable_attention_dp + and mapping.world_size == 16 + and mapping.tp_size == 16 + and mapping.cp_size == 1 + ): + return "dep16" + return ( + f"world_size={mapping.world_size} tp={mapping.tp_size} cp={mapping.cp_size} " + f"helix={helix} attention_dp={mapping.enable_attention_dp} " + f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size}" + ) + + +def _check_sm107(model_config: ModelConfig, llm_args: TorchLlmArgs) -> str: + """Return the sm107 deployment name, or raise ``NoTarget`` saying why the target does not fit.""" + pretrained_config = model_config.pretrained_config + text = getattr(pretrained_config, "text_config", None) + if text is None: + raise NoTarget("the sm107 Kimi K3 target needs the composite checkpoint's text_config") + shape = tuple(getattr(text, name, None) for name in _SM107_SHAPE_FIELDS) + if shape != _SM107_SHAPE: + raise NoTarget(f"text_config shape {shape} is not Kimi K3's {_SM107_SHAPE}") + quant = getattr(pretrained_config, "quantization_config", None) or {} + experts = [ + value + for name, value in (quant.get("quantized_layers", {}) or {}).items() + if name.endswith(".experts") + ] + if not experts or any(value.get("quant_algo") != "NVFP4" for value in experts): + raise NoTarget("the sm107 Kimi K3 target needs the NVFP4 checkpoint (NVFP4 routed experts)") + spec_config = model_config.spec_config + if spec_config is not None and not spec_config.spec_dec_mode.is_dspark(): + raise NoTarget( + f"the sm107 Kimi K3 target runs DSpark or no speculative decoding, got " + f"{spec_config.spec_dec_mode.name}" + ) + kv_algo = engine_kv_cache_quant_algo(model_config, llm_args) + if kv_algo not in (QuantAlgo.NVFP4, QuantAlgo.FP8): + raise NoTarget(f"the sm107 Kimi K3 target runs an NVFP4 or FP8 KV cache, got {kv_algo}") + mapping = model_config.mapping + parallel = _sm107_parallel(mapping) + if parallel not in ("tp1", "helix8", "helix16_moetp4ep4", "dep16"): + raise NoTarget( + f"{parallel}: the sm107 Kimi K3 target serves one rank per process (DWDP context " + "servers), Helix over 8 ranks, Helix over 16 ranks with MoE TP4 x EP4, or attention " + "DP over 16 ranks" + ) + logger.info( + f"modeling_v2 kimi_k3_vl: parallel={parallel} world_size={mapping.world_size} " + f"rank={mapping.rank} tp={mapping.tp_size} cp={mapping.cp_size} " + f"attention_dp={mapping.enable_attention_dp} moe_tp={mapping.moe_tp_size} " + f"moe_ep={mapping.moe_ep_size} dwdp_size={getattr(mapping, 'dwdp_size', None)} " + f"kv={kv_algo} spec={spec_config.spec_dec_mode.name if spec_config else None}" + ) + return parallel + + @routes("KimiK3ForConditionalGeneration") def route( model_config: ModelConfig, llm_args: TorchLlmArgs, compute_capability: Tuple[int, int] @@ -102,6 +203,16 @@ def route( Each target is imported where it is returned, so declining costs no import of modeling code. """ + if compute_capability == _SM107: + parallel = _check_sm107(model_config, llm_args) + from .kimi_k3_nvfp4__sm_107__default import features + + # The target's modules read the deployment's feature profile when imported. + features.select(parallel) + from .kimi_k3_nvfp4__sm_107__default.modeling import ModelingV2KimiK3Nvfp4Sm107Default + + return ModelingV2KimiK3Nvfp4Sm107Default + pretrained_config = model_config.pretrained_config mapping = model_config.mapping text_config = getattr(pretrained_config, "text_config", None) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index f2a3cf01748d..389b1270ff92 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -14,7 +14,12 @@ from tensorrt_llm._torch.kimi_k3_cache_policy import get_kimi_k3_fp8_kv_layer_ids from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager -from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned +from tensorrt_llm._utils import ( + TensorWrapper, + binding_to_torch_dtype, + convert_to_torch_tensor, + prefer_pinned, +) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( @@ -112,6 +117,8 @@ class Fp4MlaV2CacheLayoutPolicy: # Set by configure_manager: a disaggregated context worker keeps neither # V scales nor packed V (see is_fp4_mla_context_only_cache). fp4_mla_context_only_cache: bool = False + # Manager layers of a standalone drafter appended to the lifecycle manager. + _draft_manager_layer_ids: list = [] def __init__(self, *args, **kwargs) -> None: kv_cache_config = args[0] if args else kwargs.get("kv_cache_config") @@ -311,12 +318,13 @@ def _fp8_mla_local_layer_indices(self) -> list[int]: ] def _fp4_mla_local_layer_indices(self) -> list[int]: - fallback_layers = set(self._fp8_mla_local_layer_indices()) + # Draft layers appended by a standalone drafter keep their own storage. + excluded = set(self._fp8_mla_local_layer_indices()) | set(self._appended_draft_layer_specs) is_linear_layer = getattr(self, "_is_local_mamba_layer", None) return [ local_layer for local_layer in range(self.num_local_layers) - if local_layer not in fallback_layers + if local_layer not in excluded and (not callable(is_linear_layer) or not is_linear_layer(local_layer)) ] @@ -383,13 +391,18 @@ def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: if layer_idx is None: return self.dtype local_layer = self.layer_offsets[layer_idx] + if local_layer in self._appended_draft_layer_specs: + return self._appended_draft_layer_specs[local_layer].dtype if local_layer in self._fp8_mla_local_layer_indices(): return DataType.FP8 return self.dtype def get_kv_cache_num_blocks(self, layer_idx: int) -> int: local_layer = self.layer_offsets[layer_idx] - if local_layer in self._fp8_mla_local_layer_indices(): + if ( + local_layer in self._appended_draft_layer_specs + or local_layer in self._fp8_mla_local_layer_indices() + ): manager_layer = LayerId(local_layer) else: manager_layer = self._cache_manager_layer_ids[ @@ -497,6 +510,26 @@ def _build_cache_config(self, config: KVCacheManagerConfig) -> KVCacheManagerCon sliding_window_size=fp8_lifecycle_window, num_sink_tokens=None, ) + # Draft layers take one more distinct non-evicting lifecycle: the FP4 + # and FP8 tables address pages with their own pools' index scales. + self._draft_manager_layer_ids = [ + LayerId(local_layer) for local_layer in sorted(self._appended_draft_layer_specs) + ] + for layer_id in self._draft_manager_layer_ids: + draft_lifecycle_window = ( + self.max_seq_len + + self._generation_kv_capacity_headroom + + self._kv_reserve_draft_tokens + + 2 + ) + draft_layer = cache_layers[int(layer_id)] + cache_layers[int(layer_id)] = AttentionLayerConfig( + layer_id=layer_id, + buffers=list(draft_layer.buffers), + sliding_window_size=draft_lifecycle_window, + num_sink_tokens=draft_layer.num_sink_tokens, + reuse_policy=draft_layer.reuse_policy, + ) self._fp4_mla_local_to_compact = { local_layer: compact_layer for compact_layer, local_layer in enumerate(fp4_local_layers) } @@ -534,6 +567,9 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: fp8_pool_id = int(self.impl.get_layer_group_id(fp8_layer)) if fp8_pool_id in (cache_pool_id, hp_pool_id): raise RuntimeError("FP4 MLA, FP8 MLA, and HP state must use distinct lifecycles.") + draft_pool_id = None + if self._draft_manager_layer_ids: + draft_pool_id = int(self.impl.get_layer_group_id(self._draft_manager_layer_ids[0])) num_pools = len(self.impl.layer_grouping) pool_pointers = [[[0, 0], [0, 0]] for _ in range(num_pools)] @@ -548,16 +584,18 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: [int(self.impl.get_mem_pool_base_address(hp_layer, Role.MLA_HP_TAIL)), 0], [0, 0], ] - if fp8_pool_id is not None: - pool_pointers[fp8_pool_id] = [ - [ - int( - self.impl.get_mem_pool_base_address( - self._fp8_mla_manager_layer_ids[0], Role.KEY - ) - ), - 0, - ], + # Single-buffer pools: the FP8 fallback latents and the drafter's. + single_buffer_pools = [ + (pool_id, layer_ids) + for pool_id, layer_ids in ( + (fp8_pool_id, self._fp8_mla_manager_layer_ids), + (draft_pool_id, self._draft_manager_layer_ids), + ) + if pool_id is not None + ] + for pool_id, layer_ids in single_buffer_pools: + pool_pointers[pool_id] = [ + [int(self.impl.get_mem_pool_base_address(layer_ids[0], Role.KEY)), 0], [0, 0], ] self.kv_cache_pool_pointers = torch.tensor( @@ -572,13 +610,10 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: converter = self.impl.get_page_index_converter(layer_id, Role.KEY) local_layer = self._fp4_mla_compact_to_local[compact_layer] mapping[local_layer] = [cache_pool_id, int(converter.layer_offset)] - if fp8_pool_id is not None: - for layer_id in self._fp8_mla_manager_layer_ids: + for pool_id, layer_ids in single_buffer_pools: + for layer_id in layer_ids: converter = self.impl.get_page_index_converter(layer_id, Role.KEY) - mapping[int(layer_id)] = [ - fp8_pool_id, - int(converter.layer_offset), - ] + mapping[int(layer_id)] = [pool_id, int(converter.layer_offset)] self.kv_cache_pool_mapping = torch.tensor( mapping, dtype=torch.int32, @@ -595,11 +630,9 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: self.index_scales[hp_pool_id] = int( self.impl.get_page_index_converter(hp_layer, Role.MLA_HP_TAIL).scale ) - if fp8_pool_id is not None: - self.index_scales[fp8_pool_id] = int( - self.impl.get_page_index_converter( - self._fp8_mla_manager_layer_ids[0], Role.KEY - ).scale + for pool_id, layer_ids in single_buffer_pools: + self.index_scales[pool_id] = int( + self.impl.get_page_index_converter(layer_ids[0], Role.KEY).scale ) self.kv_offset = torch.zeros( num_pools, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() @@ -646,6 +679,17 @@ def _validate_fp4_mla_layer_groups(self) -> None: "FP8 MLA fallback layers require one lifecycle distinct from " f"FP4 and HP; got fp8={sorted(fp8_groups)}." ) + draft_groups = { + int(self.impl.get_layer_group_id(layer_id)) + for layer_id in self._draft_manager_layer_ids + } + if self._draft_manager_layer_ids and ( + len(draft_groups) != 1 or draft_groups & (cache_groups | hp_groups | fp8_groups) + ): + raise RuntimeError( + "Draft layers require one lifecycle distinct from FP4, FP8 and HP; " + f"got draft={sorted(draft_groups)}." + ) cache_roles = [Role.KEY, Role.KEY_BLOCK_SCALE] if not self.fp4_mla_context_only_cache: cache_roles.append(Role.MLA_V_SCALE) @@ -847,6 +891,21 @@ def _iter_fp4_mla_physical_pool_views(self) -> Iterable[torch.Tensor]: self.head_dim_per_layer[fp8_local_layer], ), ) + if self._draft_manager_layer_ids: + # Every appended draft layer lives in this one single-buffer pool. + draft_layer = self._draft_manager_layer_ids[0] + draft_spec = self._appended_draft_layer_specs[int(draft_layer)] + yield self._role_pool_base_view( + draft_layer, + Role.KEY, + binding_to_torch_dtype(draft_spec.dtype), + ( + draft_spec.kv_factor, + self.tokens_per_block, + draft_spec.num_kv_heads, + draft_spec.head_dim, + ), + ) if not self.fp4_mla_context_only_cache: yield self.get_mla_v_scale_pool_base().view(torch.float8_e4m3fn) if self.mla_v_head_dim is not None: @@ -1051,9 +1110,12 @@ def _fp4_mla_disagg_global_layer_id(self, internal_layer: int) -> Optional[int]: A model layer's latent cache (FP4 or FP8 fallback) maps to ``2 * L`` and its HP-ring virtual layer to ``2 * L + 1``, where ``L`` is the - global model layer. Returns ``None`` for any other internal layer. + global model layer; an appended draft layer ``D`` maps to ``2 * D``. + Returns ``None`` for any other internal layer. """ internal_layer = int(internal_layer) + if internal_layer in self._appended_draft_layer_specs: + return 2 * int(self.pp_layers[internal_layer]) for compact_layer, layer_id in enumerate(self._cache_manager_layer_ids): if int(layer_id) == internal_layer: local_layer = self._fp4_mla_compact_to_local[compact_layer] @@ -1093,6 +1155,9 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[list[int], list[Opti for local_layer in fp4_local_layers ) windows.extend([self._fp4_mla_hp_pool_size] * len(fp4_local_layers)) + draft_layers = self._appended_draft_layer_specs.values() + sizes.extend(layer.bytes_per_token() for layer in draft_layers) + windows.extend([None] * len(draft_layers)) return sizes, windows def _get_generation_request_capacity(self) -> int: diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 972e8a7979a8..d32f4c7d49b5 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -71,6 +71,12 @@ _SKIP_CORRECTION_SUPPORTED_SMS = frozenset((100, 103)) +def _has_fp8_context_mla_fmha(sm_version: int) -> bool: + """Whether MLA context attention with an FP8 or NVFP4 KV cache runs the op's + FP8 context MLA FMHA on this SM (``mFP8ContextMLA`` in thop/attentionOp.cpp).""" + return sm_version == 90 or is_sm_100f(sm_version) or is_sm_120f(sm_version) + + def _resolve_skip_correction_threshold(threshold: float, sm_version: int, *, @@ -2064,6 +2070,13 @@ def uses_fp4_mla_attention(self) -> bool: return (self.is_mla_enable and self.has_fp4_kv_cache and self.sparse_params is None) + @property + def has_fp8_context_mla(self) -> bool: + """Whether MLA context attention with an FP8 KV cache runs the op's FP8 + context MLA FMHA, which reads FP8 Q, K and V.""" + return (self.is_mla_enable and self.has_fp8_kv_cache + and _has_fp8_context_mla_fmha(get_sm_version())) + def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.quant_config = new_quant_config or QuantConfig() self.quant_mode = self.quant_config.layer_quant_mode @@ -2171,12 +2184,10 @@ def runtime_workspace_bytes_per_token(cls, model_config: "ModelConfig", if fp8_fallback_layers and fp8_fallback_layers.intersection( mapping.pp_layers(config.num_hidden_layers)): nvfp4_uses_fp8_context = True - sm_version = get_sm_version() fp8_context_mla = (quant_mode is not None and (quant_mode.has_fp8_kv_cache() or nvfp4_uses_fp8_context) - and (sm_version == 90 or is_sm_100f(sm_version) - or is_sm_120f(sm_version))) + and _has_fp8_context_mla_fmha(get_sm_version())) if not fp8_context_mla: return 0 # Attention-DP runs the full head set per rank; otherwise heads shard across TP (mirror diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 7f35307c75b1..20f7388b1525 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -175,6 +175,37 @@ def validate_bounded_replay_compatible(self, peer_ri: RankInfo) -> None: f"peer={peer_ri.bounded_replay_on_generation})." ) + def validate_draft_layers_compatible(self, peer_ri: RankInfo) -> None: + """Raise ValueError when the peers' shared draft KV layouts disagree. + + Both sides run it: the sender when it registers a receiver, and the + receiver before it registers with a sender, so a mismatch fails the + receiver's requests instead of leaving them waiting for KV. + """ + if self._ri.num_draft_layers != peer_ri.num_draft_layers: + # The layer-count tolerance in _check_peer_compatible would transfer + # only the target layers and leave the drafter's prompt KV unset on + # the receiver. + raise ValueError( + "PeerRegistrar: unified draft KV layer count differs " + f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " + "both peers must use the same speculative decoding KV cache layout; " + "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." + ) + # Draft layers carry the target's KV heads per rank, so they move with + # the target's head mapping only when both sides shard heads alike. A + # K-only (MLA) cache keeps the whole latent on every rank. + if ( + self._ri.num_draft_layers + and not (self._ri.attention.is_mla and peer_ri.attention.is_mla) + and self._ri.tp_size_per_dp_group != peer_ri.tp_size_per_dp_group + ): + raise ValueError( + "PeerRegistrar: a shared draft KV cache needs the same attention TP size " + f"(local={self._ri.tp_size_per_dp_group}, peer={peer_ri.tp_size_per_dp_group}); " + "unset TRTLLM_DISAGG_SHARED_DRAFT_KV or shard attention alike on both workers." + ) + def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: # Wire-format gate first: a peer on another protocol version cannot be # trusted on any other field. Raises ValueError naming both versions. @@ -196,29 +227,7 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) - if self._ri.num_draft_layers != peer_ri.num_draft_layers: - # The layer-count tolerance below would transfer only the target - # layers and leave the drafter's prompt KV unset on the receiver. - raise ValueError( - "PeerRegistrar: unified draft KV layer count differs " - f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " - "both peers must use the same speculative decoding KV cache layout; " - "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." - ) - if self._ri.num_draft_layers: - if self._ri.tp_size != peer_ri.tp_size: - raise ValueError( - "PeerRegistrar: a shared draft KV cache needs equal tp_size " - f"(local={self._ri.tp_size}, peer={peer_ri.tp_size}); " - "unset TRTLLM_DISAGG_SHARED_DRAFT_KV or use the same TP on both workers." - ) - if self._ri.attention.enable_attention_dp != peer_ri.attention.enable_attention_dp: - raise ValueError( - "PeerRegistrar: a shared draft KV cache needs the same attention DP " - f"setting (local={self._ri.attention.enable_attention_dp}, " - f"peer={peer_ri.attention.enable_attention_dp}); unset " - "TRTLLM_DISAGG_SHARED_DRAFT_KV or match it on both workers." - ) + self.validate_draft_layers_compatible(peer_ri) self.validate_view_head_geometry(peer_ri) self_layers = sum(self._ri.layer_num_per_pp) diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index eb432bd63d29..5ba03f3087a4 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -123,6 +123,10 @@ def from_kv_cache_manager( if bytes_for_eight_elements % 8 == 0 else bytes_for_eight_elements / 8 ) + # Explicit draft layers (layer specs) ride the target's pages and + # requests; folded MTP/Eagle layers travel as ordinary layers. + get_draft_view = getattr(kvm, "get_draft_kv_cache_view", None) + shares_draft_layers = get_draft_view is not None and get_draft_view() is kvm return cls( instance_name=instance_name, instance_rank=m.rank, @@ -150,9 +154,7 @@ def from_kv_cache_manager( aux_meta=aux_buffer_meta, page_table=build_page_table_from_manager(kvm), bounded_replay_on_generation=getattr(kvm, "bounded_replay_on_generation", False), - num_draft_layers=( - len(kvm.draft_layer_ids) if getattr(kvm, "layer_specs", None) is not None else 0 - ), + num_draft_layers=len(kvm.draft_layer_ids) if shares_draft_layers else 0, ) @classmethod diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index c4c6ea0eef6a..ca3dba4810b6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -3390,6 +3390,7 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: self._registrar.self_extractor.page_table, sender_info.page_table, ) + self._registrar.validate_draft_layers_compatible(sender_info) self._registrar.validate_view_head_geometry(sender_info) except ValueError as e: msg = ( diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 1d94771554b1..d605ea7b13cc 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -606,6 +606,26 @@ class DFlashForCausalLM(nn.Module): # tests), so the config cap stands. _runtime_position_ceiling = None + def reads_managed_ctx_pages(self, tokens_per_block: int) -> bool: + """Whether the block decode reads the context KV from a KV cache + manager's pages of ``tokens_per_block`` tokens. + + The TRTLLM block decode takes only the page sizes its trtllm-gen + kernels support; at any other size the drafter keeps its own pool. + """ + from ..attention.backends.fmha.flashinfer_trtllm_gen import FlashInferTrtllmGenFmha + + return ( + self.dflash_attention_backend == "TRTLLM" + and tokens_per_block in FlashInferTrtllmGenFmha.SUPPORTED_TOKENS_PER_BLOCK + ) + + @property + def reads_helix_sharded_ctx_pages(self) -> bool: + """Whether the block decode reads a context whose pages are sharded over + helix CP, merging the ranks' partial attention.""" + return False + @classmethod def _resolve_auto_attention_backend(cls) -> str: """Turn ``AUTO`` into a concrete backend for this drafter family.""" @@ -711,6 +731,9 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): self.block_size = dflash_config.get( "block_size", getattr(pretrained_config, "block_size", None) ) + # The block's draft tokens sit in slots 0..K-1 (the DSpark shift_label + # convention, set by DSpark heads) instead of 1..K. + self.shift_label = False self.logit_multiplier, self.logit_softcap = resolve_draft_logit_transform(dflash_config) if dflash_attention_backend == "AUTO": dflash_attention_backend = self._resolve_auto_attention_backend() diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 701f3bfc6f76..165ed34cd213 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2904,7 +2904,7 @@ def _init_dspark_heads(self, cfg) -> None: # runs one slot past the block and reads the next request's anchor. # An explicit false still selects the legacy layout. shift_label = resolve_dspark_head_config(cfg, "shift_label") - self._dspark_shift_label = True if shift_label is None else bool(shift_label) + self.shift_label = True if shift_label is None else bool(shift_label) self._dspark_markov_rank = int(resolve_dspark_head_config(cfg, "markov_rank") or 0) self._dspark_markov_head_type = str( resolve_dspark_head_config(cfg, "markov_head_type") or "vanilla" diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index 98662f5382df..053f9ac9add8 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -932,6 +932,8 @@ class KimiK3HybridCacheManager(MambaHybridCacheManagerV2): """KDA/MLA cache layout and K3-only replay lifecycle handling.""" kda_context_only = False + # The DSpark MLA drafter keeps its context KV in this manager. + accepts_appended_draft_layers = True @override def __init__( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b35ed3fc40e5..f97720cd744b 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1745,7 +1745,8 @@ def _get_token_num_for_estimation(self) -> int: if self._is_kv_cache_manager_v2: unified_layer_specs = self._build_unified_dflash_layer_specs( self._kv_cache_config) - if unified_layer_specs: + if unified_layer_specs and not issubclass( + self._kv_cache_manager_cls, MambaHybridCacheManagerV2): # Target and draft layers can differ in page size and life # cycle, so the group count comes from the combined layout. num_pool_groups = _get_num_storage_pool_groups( @@ -1757,6 +1758,10 @@ def _get_token_num_for_estimation(self) -> int: self._model_engine.max_seq_len, self._kv_cache_config, ) + if unified_layer_specs: + # A hybrid manager's appended draft layers. + num_pool_groups += _get_num_storage_pool_groups( + unified_layer_specs, self._tokens_per_block) num_cache_blocks *= num_pool_groups # Dummy context requests use the configured maximum beam width. Scale @@ -2323,10 +2328,54 @@ def _unified_drafter_attention_backend(self) -> Optional[str]: backend = getattr(draft_model, "dflash_attention_backend", None) return backend + def _loaded_standalone_dspark_drafter(self): + """The loaded standalone DSpark drafter, whose declarations decide how + it reads its context, or None. + + DFlash keeps its configured backend. An embedded DSpark draft is not a + ``DFlashForCausalLM`` and declares nothing: its context lives in a + worker-owned rolling window. + """ + spec_config = self._speculative_config + spec_dec_mode = spec_config.spec_dec_mode + if (spec_dec_mode.is_dflash() or not spec_dec_mode.is_dspark() + or spec_config.draft_is_embedded_in_target): + return None + return getattr(self._model_engine.model, "draft_model", None) + + def _unified_drafter_reads_paged_kv(self, backend: Optional[str]) -> bool: + """Whether the drafter's block attention reads its context KV from the + manager's pages at this run's page size. + + A loaded DSpark drafter declares it; otherwise only TRTLLM pages of 16, + 32 or 64 tokens. + """ + drafter = self._loaded_standalone_dspark_drafter() + if drafter is None: + return (backend == "TRTLLM" + and self._tokens_per_block in (16, 32, 64)) + return drafter.reads_managed_ctx_pages(self._tokens_per_block) + + def _unified_drafter_reads_helix_pages(self) -> bool: + """Whether the loaded drafter reads a context sharded over helix CP.""" + drafter = self._loaded_standalone_dspark_drafter() + return drafter is not None and drafter.reads_helix_sharded_ctx_pages + + def _unified_target_manager_cls_supported(self) -> bool: + """Whether the target's manager class can hold the drafter's layers. + + Hybrid managers describe their own target layers and append the + drafter's after them when they opt in. + """ + cls = self._kv_cache_manager_cls + return cls is KVCacheManagerV2 or ( + issubclass(cls, MambaHybridCacheManagerV2) + and cls.accepts_appended_draft_layers) + def _unified_drafter_block_width(self) -> int: """Use the loaded checkpoint's slot convention for sizing and allocation.""" draft_model = getattr(self._model_engine.model, "draft_model", None) - shift_label = getattr(draft_model, "_dspark_shift_label", False) + shift_label = draft_model is not None and draft_model.shift_label return get_draft_block_width(self._speculative_config, shift_label=shift_label) @@ -2366,10 +2415,12 @@ def _dflash_unify_blocker(self) -> Optional[str]: if spec_config is None: return "speculative decoding is not enabled" backend = self._unified_drafter_attention_backend() - if backend != "TRTLLM": - return f"the drafter attention_backend is {backend}, not TRTLLM" - if self._tokens_per_block not in (16, 32, 64): - return f"tokens_per_block is {self._tokens_per_block}, not 16, 32, or 64" + if backend is None: + return "the drafter has no block attention over paged draft KV" + if not self._unified_drafter_reads_paged_kv(backend): + return (f"the drafter's {backend} block attention does not read " + "paged draft KV from a KV cache manager at " + f"tokens_per_block={self._tokens_per_block}") if not spec_config._allow_separate_draft_kv_cache: return "a separate draft KV cache is not allowed for this configuration" if spec_config.draft_len_schedule is not None: @@ -2378,8 +2429,11 @@ def _dflash_unify_blocker(self) -> Optional[str]: return "there is no draft model config" if self._mapping.pp_size != 1: return f"pipeline_parallel_size is {self._mapping.pp_size}, not 1" - if self._mapping.cp_size != 1: - return f"context_parallel_size is {self._mapping.cp_size}, not 1" + if self._mapping.cp_size != 1 and not ( + self._mapping.has_cp_helix() + and self._unified_drafter_reads_helix_pages()): + return (f"context_parallel_size is {self._mapping.cp_size} and the " + "drafter does not read a helix-sharded context") if self._llm_args.attn_backend != "TRTLLM": return f"the target attn_backend is {self._llm_args.attn_backend}, not TRTLLM" if _has_variable_kv_cache_geometry( @@ -2392,8 +2446,9 @@ def _dflash_unify_blocker(self) -> Optional[str]: return "KV cache compression is enabled" # Specialized managers add cache roles and lifecycle protocols that an # external drafter cannot yet describe. Keep their existing fallback. - if self._kv_cache_manager_cls is not KVCacheManagerV2: - return "the target does not use KVCacheManagerV2" + if not self._unified_target_manager_cls_supported(): + return (f"the target's {self._kv_cache_manager_cls.__name__} " + "cannot hold draft layers") target_layout = self._dflash_kv_layout( self._model_engine.model.model_config, self._model_engine.model.config.torch_dtype) @@ -2448,15 +2503,16 @@ def _dflash_disagg_layouts_compatible( per-layer bytes over the rank's KV-head count, so dtype and head_dim may differ. The KV-head count and K-only flag come from the rank's single AttentionInfo and must match. NVFP4 adds block-scale views, so - an NVFP4 side is only paired with an identical layout. + an NVFP4 side is only paired with an identical layout, unless both are + K-only: every rank holds the whole latent, so no view is head-sliced. """ target_heads, _, target_dtype, target_k_only = target_layout draft_heads, _, draft_dtype, draft_k_only = draft_layout if (target_heads, target_k_only) != (draft_heads, draft_k_only): return False nvfp4 = tensorrt_llm.bindings.DataType.NVFP4 - return target_dtype == draft_dtype or nvfp4 not in (target_dtype, - draft_dtype) + return (target_k_only or target_dtype == draft_dtype + or nvfp4 not in (target_dtype, draft_dtype)) def _dflash_kv_layout( self, model_config: ModelConfig, @@ -2515,9 +2571,16 @@ def _layer_spec_roles(dtype, *, k_only: bool) -> tuple: def _build_unified_dflash_layer_specs( self, kv_cache_config: KvCacheConfig) -> tuple[KVCacheLayerSpec, ...]: - """Build the rank-local target+DFlash storage layout used by cost and allocation.""" + """Build the rank-local target+DFlash storage layout used by cost and allocation. + + A hybrid target manager describes its own layers, so it gets only the + drafter's. + """ if not self._should_unify_dflash_kv_cache(): return () + draft_specs = self._build_unified_draft_layer_specs(kv_cache_config) + if issubclass(self._kv_cache_manager_cls, MambaHybridCacheManagerV2): + return draft_specs target_model_config = self._model_engine.model.model_config target_config = target_model_config.pretrained_config @@ -2540,20 +2603,6 @@ def _build_unified_dflash_layer_specs( target_local_ids, ) - draft_model_config = self._get_draft_kv_model_config() - draft_config = draft_model_config.pretrained_config - draft_layers = int(draft_config.num_hidden_layers) - draft_local_ids = list(self._mapping.pp_layers(draft_layers)) - draft_heads, draft_head_dim, draft_dtype, draft_k_only = ( - self._dflash_kv_layout(draft_model_config, - draft_config.torch_dtype)) - draft_windows = _resolve_v2_max_attention_window_vec( - self._get_draft_max_attention_window(self._max_seq_len, - kv_cache_config), - self._max_seq_len, - draft_local_ids, - ) - specs = [ KVCacheLayerSpec( global_layer_id=layer_id, @@ -2567,7 +2616,28 @@ def _build_unified_dflash_layer_specs( reuse_policy=ReusePolicy.REQUIRED, ) for layer_id, window in zip(target_local_ids, target_windows) ] - specs.extend( + return (*specs, *draft_specs) + + def _build_unified_draft_layer_specs( + self, + kv_cache_config: KvCacheConfig) -> tuple[KVCacheLayerSpec, ...]: + """Rank-local storage of the drafter's layers, after the target's.""" + target_layers = int(self._model_engine.model.model_config. + pretrained_config.num_hidden_layers) + draft_model_config = self._get_draft_kv_model_config() + draft_config = draft_model_config.pretrained_config + draft_layers = int(draft_config.num_hidden_layers) + draft_local_ids = list(self._mapping.pp_layers(draft_layers)) + draft_heads, draft_head_dim, draft_dtype, draft_k_only = ( + self._dflash_kv_layout(draft_model_config, + draft_config.torch_dtype)) + draft_windows = _resolve_v2_max_attention_window_vec( + self._get_draft_max_attention_window(self._max_seq_len, + kv_cache_config), + self._max_seq_len, + draft_local_ids, + ) + return tuple( KVCacheLayerSpec( global_layer_id=target_layers + layer_id, owner=KVCacheLayerOwner.DRAFT, @@ -2579,7 +2649,6 @@ def _build_unified_dflash_layer_specs( sliding_window_size=window, reuse_policy=ReusePolicy.REQUIRED, ) for layer_id, window in zip(draft_local_ids, draft_windows)) - return tuple(specs) def _joint_reuse_supported(self) -> bool: """Span half of the pairing decision; ``build_managers`` ANDs it with @@ -5296,7 +5365,10 @@ def _adjust_torch_mem_fraction(): return mem_reserved = torch.cuda.memory_reserved() mem_free, mem_total = torch.cuda.mem_get_info() - safety_margin = 32 * 1024**2 + # Keep room outside the PyTorch cap for driver allocations made after setup + # (streams, events, copy launches in the KV cache manager and transfer agents), + # up to the free memory, so the cap never drops below what PyTorch holds. + safety_margin = min(1024**3, mem_free) mem_torch_max = mem_free + mem_reserved - safety_margin mem_torch_fraction = mem_torch_max / mem_total logger.info( diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index f2af478d1718..3de4d257cd69 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1856,6 +1856,10 @@ class KVCacheManagerV2(BaseResourceManager): inforward_snapshot_rows: int = 0 # Block width of a drafter whose layers share this manager, else 0. _shared_draft_block_width: int = 0 + # Local layer id -> storage of the draft-owned layers a subclass appends + # after the layers it describes itself (MambaHybridCacheManagerV2). + # Empty when ``layer_specs`` describes every layer. + _appended_draft_layer_specs: dict[int, KVCacheLayerSpec] = {} @property def _kv_alloc_extra_tokens(self) -> int: @@ -2019,6 +2023,21 @@ def __init__( layer.owner is not KVCacheLayerOwner.DRAFT for layer in self.layer_specs ): raise ValueError("a draft mirror manager cannot own target cache layers") + for local_layer_id, layer in self._appended_draft_layer_specs.items(): + if ( + layer.owner is not KVCacheLayerOwner.DRAFT + or self.pp_layers[local_layer_id] != layer.global_layer_id + ): + raise ValueError( + f"appended draft layer {layer.global_layer_id} must be a draft-owned " + f"layer at local id {local_layer_id}" + ) + # Explicit storage of the layers a drafter owns in this manager. + draft_layer_specs = ( + self.layer_specs + if self.layer_specs is not None + else tuple(self._appended_draft_layer_specs.values()) or None + ) # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) # can distinguish the throwaway estimation-phase managers from the @@ -2031,10 +2050,7 @@ def __init__( self.enable_swa_scratch_reuse = ( kv_cache_config.enable_swa_scratch_reuse and not self.is_draft - and not ( - self.layer_specs is not None - and any(layer.owner is KVCacheLayerOwner.DRAFT for layer in self.layer_specs) - ) + and not any(layer.owner is KVCacheLayerOwner.DRAFT for layer in draft_layer_specs or ()) ) block_reuse_config = kv_cache_config.block_reuse_config self.block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) @@ -2099,12 +2115,14 @@ def __init__( "all CP ranks, so publishing it requires a group-wide " "protocol. Set kv_cache_config.enable_block_reuse=False." ) + # K-only draft latents shard by page over CP like the target's. if is_draft or any( - layer.owner is KVCacheLayerOwner.DRAFT for layer in self.layer_specs or () + layer.owner is KVCacheLayerOwner.DRAFT and layer.kv_factor != 1 + for layer in draft_layer_specs or () ): raise ValueError( - "KVCacheManagerV2 does not support a draft cache " - "manager or draft-owned layers with helix context parallelism." + "KVCacheManagerV2 does not support a draft cache manager or " + "draft-owned K/V layers with helix context parallelism." ) if mapping.enable_attention_dp and mapping.pp_size > 1: # Only the pp=1 executor loops call @@ -2153,7 +2171,7 @@ def __init__( self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) self._shared_draft_block_width = _get_shared_draft_block_width( - spec_config, self.layer_specs, shared_draft_block_width + spec_config, draft_layer_specs, shared_draft_block_width ) # The draft request view (LlmRequest.use_draft_model) reads the context # position and chunk a separate draft engine advances for its own @@ -2200,7 +2218,7 @@ def __init__( _get_generation_kv_capacity( spec_config, is_draft=self.is_draft, - layer_specs=self.layer_specs, + layer_specs=draft_layer_specs, shared_draft_block_width=self._shared_draft_block_width, ) ) @@ -2232,6 +2250,11 @@ def __init__( else layer.sliding_window_size for layer in self.layer_specs ] + for local_layer_id, layer in self._appended_draft_layer_specs.items(): + window = layer.sliding_window_size + self.max_attention_window_vec[local_layer_id] = ( + None if window is None or window >= max_seq_len else window + ) # A D-token match backoff needs D more physical SWA history at the # committed endpoint so commit_min_snapshot still covers the rewind. @@ -2361,6 +2384,11 @@ def append_to_kv_heads_per_layer( f"Per-layer head_dim: {len(self.head_dim_per_layer)} layers, " f"unique values={set(self.head_dim_per_layer)}" ) + # Appended draft specs are rank-local already. + for local_layer_id, layer in self._appended_draft_layer_specs.items(): + self.num_kv_heads_per_layer[local_layer_id] = layer.num_kv_heads + self.total_num_kv_heads_per_layer[layer.global_layer_id] = layer.num_kv_heads + self.head_dim_per_layer[local_layer_id] = layer.head_dim self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) @@ -3086,26 +3114,36 @@ def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: role_b = None if self.kv_cache_type == CacheTypeCpp.SELFKONLY else Role.VALUE return Role.KEY, role_b - def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: + def _explicit_layer_spec(self, local_layer_id: int) -> Optional[KVCacheLayerSpec]: + """Explicitly supplied storage of a local layer, or None.""" if self.layer_specs is not None: - return self.layer_specs[local_layer_id].kv_factor + return self.layer_specs[local_layer_id] + return self._appended_draft_layer_specs.get(local_layer_id) + + def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: + layer = self._explicit_layer_spec(local_layer_id) + if layer is not None: + return layer.kv_factor return self.kv_factor def _get_layer_page_table_roles( self, local_layer_id: int, pool_id: int ) -> Tuple[DataRole, Optional[DataRole]]: - if self.layer_specs is None: + layer = self._explicit_layer_spec(local_layer_id) + if layer is None: return self._get_pool_roles(pool_id) - return Role.KEY, (Role.VALUE if self.layer_specs[local_layer_id].kv_factor == 2 else None) + return Role.KEY, (Role.VALUE if layer.kv_factor == 2 else None) def _get_layer_dtype_by_local_id(self, local_layer_id: int) -> DataType: - if self.layer_specs is not None: - return self.layer_specs[local_layer_id].dtype + layer = self._explicit_layer_spec(local_layer_id) + if layer is not None: + return layer.dtype return self.dtype def _get_layer_roles_by_local_id(self, local_layer_id: int) -> tuple[DataRole, ...]: - if self.layer_specs is not None: - return self.layer_specs[local_layer_id].buffer_roles + layer = self._explicit_layer_spec(local_layer_id) + if layer is not None: + return layer.buffer_roles return tuple(self._get_buffer_roles_for_layer(local_layer_id)) def _get_block_scale_role(self, role_a: DataRole, layer_id: int) -> Optional[DataRole]: @@ -3286,6 +3324,12 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: if self.layer_specs is not None: storage_layouts = {(layer.dtype, layer.kv_factor) for layer in self.layer_specs} self._use_per_layer_page_tables |= len(storage_layouts) > 1 + elif self._appended_draft_layer_specs: + storage_layouts = {(self.dtype, self.kv_factor)} | { + (layer.dtype, layer.kv_factor) + for layer in self._appended_draft_layer_specs.values() + } + self._use_per_layer_page_tables |= len(storage_layouts) > 1 if self._use_per_layer_page_tables: self.num_attention_op_pools = self.num_local_layers else: @@ -3817,8 +3861,9 @@ def _copy_batch_block_offsets_per_layer( def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: """Return the primary cache roles allocated for one local layer.""" - if self.layer_specs is not None: - return list(self.layer_specs[local_layer_idx].buffer_roles) + layer = self._explicit_layer_spec(local_layer_idx) + if layer is not None: + return list(layer.buffer_roles) roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: roles.append(Role.VALUE) @@ -3833,6 +3878,31 @@ def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: roles.append(Role.VALUE_BLOCK_SCALE) return roles + def _explicit_layer_config( + self, local_layer_id: int, layer: KVCacheLayerSpec, tokens_per_block: int + ) -> AttentionLayerConfig: + """Storage config of a layer described by an explicit layer spec.""" + scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + return AttentionLayerConfig( + layer_id=LayerId(local_layer_id), + buffers=[ + BufferConfig( + role=role, + size=layer.bytes_per_token(role) * tokens_per_block, + # Keep NVFP4 data (1), its scales (2), and other data (0) + # separate while preserving matching K/V page order. + coalescing_group=( + 1 + int(role in scale_roles) if layer.dtype == DataType.NVFP4 else 0 + ), + ) + for role in (Role.KEY, Role.VALUE, *scale_roles) + if role in layer.buffer_roles + ], + sliding_window_size=self.max_attention_window_vec[local_layer_id], + num_sink_tokens=layer.num_sink_tokens, + reuse_policy=layer.reuse_policy, + ) + def _build_base_config( self, kv_cache_config: KvCacheConfig, @@ -3913,27 +3983,8 @@ def _build_base_config( ) if self.layer_specs is not None: - scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) layer_configs = [ - AttentionLayerConfig( - layer_id=LayerId(local_layer_id), - buffers=[ - BufferConfig( - role=role, - size=layer.bytes_per_token(role) * tokens_per_block, - # Keep NVFP4 data (1), its scales (2), and other data (0) - # separate while preserving matching K/V page order. - coalescing_group=( - 1 + int(role in scale_roles) if layer.dtype == DataType.NVFP4 else 0 - ), - ) - for role in (Role.KEY, Role.VALUE, *scale_roles) - if role in layer.buffer_roles - ], - sliding_window_size=self.max_attention_window_vec[local_layer_id], - num_sink_tokens=layer.num_sink_tokens, - reuse_policy=layer.reuse_policy, - ) + self._explicit_layer_config(local_layer_id, layer, tokens_per_block) for local_layer_id, layer in enumerate(self.layer_specs) ] else: @@ -3950,6 +4001,12 @@ def _build_base_config( layer_configs: List[AttentionLayerConfig] = [] for layer_id in typed_range(LayerId(self.num_local_layers)): + appended_layer = self._appended_draft_layer_specs.get(int(layer_id)) + if appended_layer is not None: + layer_configs.append( + self._explicit_layer_config(int(layer_id), appended_layer, tokens_per_block) + ) + continue buffer_type = self._get_buffer_roles_for_layer(layer_id) buffers = [ BufferConfig( @@ -4464,20 +4521,26 @@ def get_disagg_pool_view_partition( def get_layer_cache_spec(self, global_layer_id: int) -> KVCacheLayerSpec: """Return an explicitly supplied rank-local layer specification.""" - if self.layer_specs is None: + if self.layer_specs is None and not self._appended_draft_layer_specs: raise RuntimeError("this manager was not constructed with explicit layer specs") try: local_layer_id = self.layer_offsets[global_layer_id] except KeyError as error: raise KeyError(f"global layer {global_layer_id} is not local to this rank") from error - return self.layer_specs[local_layer_id] + layer = self._explicit_layer_spec(local_layer_id) + if layer is None: + raise RuntimeError(f"global layer {global_layer_id} has no explicit layer spec") + return layer def get_layer_owner(self, global_layer_id: int) -> KVCacheLayerOwner: if self.layer_specs is not None: return self.get_layer_cache_spec(global_layer_id).owner + local_layer_id = self.layer_offsets.get(global_layer_id) + if local_layer_id in self._appended_draft_layer_specs: + return KVCacheLayerOwner.DRAFT if global_layer_id in self._legacy_draft_layer_ids: return KVCacheLayerOwner.DRAFT - if global_layer_id not in self.layer_offsets: + if local_layer_id is None: raise KeyError(f"global layer {global_layer_id} is not local to this rank") return KVCacheLayerOwner.TARGET @@ -4516,6 +4579,8 @@ def draft_layer_ids(self) -> tuple[int, ...]: def get_draft_kv_cache_view(self) -> Optional["KVCacheManagerV2"]: """Return this manager when explicit draft layers share its request timeline.""" + if self._appended_draft_layer_specs: + return self if self.layer_specs is None or not self.draft_layer_ids: return None return self @@ -7649,15 +7714,17 @@ def get_cache_bytes_per_token(self) -> int: if self.kv_cache_type != CacheTypeCpp.SELFKONLY: data_roles.append(Role.VALUE_BLOCK_SCALE) - return sum( + appended = self._appended_draft_layer_specs + return sum(layer.bytes_per_token() for layer in appended.values()) + sum( self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=data_role) for local_layer_idx in range(self.num_local_layers) + if local_layer_idx not in appended for data_role in data_roles ) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): - if self.layer_specs is not None: - layer = self.layer_specs[local_layer_idx] + layer = self._explicit_layer_spec(local_layer_idx) + if layer is not None: if data_role == Role.ALL: return layer.bytes_per_token() return layer.bytes_per_token(data_role) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 518bd08ca37b..9af459fe63a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -45,6 +45,8 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _RESERVED_REQUEST_IDS, BlockReusePolicy, + KVCacheLayerOwner, + KVCacheLayerSpec, KVCacheManagerV2, Role, ) @@ -63,6 +65,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttnLifeCycle, BatchDesc, BufferConfig, DataRole, @@ -110,6 +113,39 @@ _INFORWARD_SNAPSHOT_STAGING_BYTES = 512 << 20 +def _insert_draft_pool_ratio( + target_ratio: Sequence[float], + life_cycle_block_bytes: Sequence[int], + full_history: Sequence[bool], + draft_life_cycle: int, +) -> List[float]: + """Pool ratio per lifecycle with the drafter's share inserted. + + ``target_ratio`` lists the other lifecycles' shares in registration order. + The drafter keeps one context entry per target token, so its share is + solved for as many blocks as the scarcest full-history target lifecycle + holds: ``d / B_draft == (1 - d) * min_i(r_i / B_i)``. The target shares + keep their split of ``1 - d``. + """ + target_ids = [i for i in range(len(life_cycle_block_bytes)) if i != draft_life_cycle] + if len(target_ratio) != len(target_ids): + raise ValueError( + f"pool_ratio has {len(target_ratio)} entries for {len(target_ids)} target layer groups" + ) + blocks_per_byte = [ + ratio / life_cycle_block_bytes[i] + for ratio, i in zip(target_ratio, target_ids) + if full_history[i] + ] + if not blocks_per_byte: + raise ValueError("the drafter's pool share needs a full-history target layer group") + draft_blocks = life_cycle_block_bytes[draft_life_cycle] * min(blocks_per_byte) + draft_share = draft_blocks / (1 + draft_blocks) + ratio = [share * (1 - draft_share) for share in target_ratio] + ratio.insert(draft_life_cycle, draft_share) + return ratio + + class InforwardSnapshotEntry(NamedTuple): """One request's in-forward snapshot in this step's context chunk.""" @@ -159,6 +195,9 @@ class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): _inforward_stage_conv: Optional[torch.Tensor] = None # Recurrent snapshots use their own commit/history endpoints without draft backoff. _supports_reuse_match_backoff = False + # Whether the KV cache creator may append a standalone drafter's layers + # (``layer_specs``) to this manager; set by the subclasses validated with one. + accepts_appended_draft_layers = False # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the # cache manager __init__ coerces a positive spec recompute tail to a full @@ -214,6 +253,9 @@ def __init__( Args: gdn_fused_verify: Whether the GDN layers verify drafts with the fused kernels (see gdn_replay_partitioned). + layer_specs: Draft-owned attention layers a standalone drafter + keeps in this manager, appended after the target layers that + the mask arguments describe (see ``_append_draft_layers``). """ if conv_state_layout not in ("x_b_c", "q_k_v"): raise ValueError(f"Unsupported convolution state layout: {conv_state_layout!r}") @@ -226,6 +268,15 @@ def __init__( f"(got {kwargs['model_type']!r}); pass " "conv_state_layout='x_b_c' or 'q_k_v' instead" ) + draft_layer_specs = tuple(kwargs.pop("layer_specs", None) or ()) + mamba_layer_mask, layer_mask, num_kv_heads = self._append_draft_layers( + draft_layer_specs, + mamba_layer_mask, + layer_mask, + num_kv_heads, + k_only=kv_cache_type == CacheTypeCpp.SELFKONLY, + ) + num_layers += len(draft_layer_specs) total_layers = len(mamba_layer_mask) if layer_mask is None: full_attention_layer_mask = [False] * total_layers @@ -278,6 +329,12 @@ def __init__( layer_idx for layer_idx in self.pp_layers if mamba_layer_mask[layer_idx] ] self.local_num_mamba_layers = len(self.mamba_pp_layers) + local_ids = {layer_idx: offset for offset, layer_idx in enumerate(self.pp_layers)} + self._appended_draft_layer_specs = { + local_ids[layer.global_layer_id]: layer + for layer in draft_layer_specs + if layer.global_layer_id in local_ids + } if self.local_num_mamba_layers > 0: tp_size = _mamba_effective_tp_size(mapping) @@ -413,6 +470,53 @@ def __init__( self._setup_state_pool() + @staticmethod + def _append_draft_layers( + draft_layer_specs: Sequence[KVCacheLayerSpec], + mamba_layer_mask: List[bool], + layer_mask: Optional[List[bool]], + num_kv_heads: Union[int, List[Optional[int]]], + *, + k_only: bool, + ) -> Tuple[List[bool], Optional[List[bool]], Union[int, List[Optional[int]]]]: + """Masks and KV heads with a standalone drafter's layers appended. + + The masks describe this manager's recurrent and target attention + layers; the drafter's layers follow them in global layer id order and + keep the storage of their specs. + """ + if not draft_layer_specs: + return mamba_layer_mask, layer_mask, num_kv_heads + num_target_layers = len(mamba_layer_mask) + expected_ids = list(range(num_target_layers, num_target_layers + len(draft_layer_specs))) + if [layer.global_layer_id for layer in draft_layer_specs] != expected_ids: + raise ValueError( + "MambaHybridCacheManagerV2 layer_specs must be the drafter's layers " + f"{expected_ids} after the target layers" + ) + for layer in draft_layer_specs: + if layer.owner is not KVCacheLayerOwner.DRAFT: + raise ValueError( + "MambaHybridCacheManagerV2 describes its target layers itself; " + f"layer_specs may add only draft-owned layers (layer {layer.global_layer_id})" + ) + if layer.kv_factor != (1 if k_only else 2): + raise ValueError( + f"draft layer {layer.global_layer_id} has kv_factor {layer.kv_factor}, " + "unlike the target cache type" + ) + if layer_mask is None: + layer_mask = [False] * num_target_layers + if isinstance(num_kv_heads, int): + num_kv_heads = [num_kv_heads] * num_target_layers + num_draft = len(draft_layer_specs) + return ( + list(mamba_layer_mask) + [False] * num_draft, + list(layer_mask) + [True] * num_draft, + # Placeholders: the base manager takes the rank-local heads from the specs. + list(num_kv_heads) + [layer.num_kv_heads for layer in draft_layer_specs], + ) + @override def _disagg_transfer_overwrites_whole_cached_prefix(self) -> bool: # A transferred recurrent state summarizes and overwrites the whole prefix. @@ -444,6 +548,10 @@ def get_cache_size_per_token( **kwargs, ): spec_config = kwargs.get("spec_config") + # Appended draft layers grow with every token like the target's pages. + draft_layer_specs = kwargs.pop("layer_specs", None) or () + kwargs.pop("shared_draft_block_width", None) + draft_slope = sum(layer.bytes_per_token() for layer in draft_layer_specs) num_reserved_dummy_slots = _get_num_cuda_graph_padding_dummy_slots( spec_config, max_batch_size ) + int(mapping.enable_attention_dp) @@ -505,7 +613,14 @@ def get_cache_size_per_token( attention_page_bytes_override=attention_page_bytes_override, **kwargs, ) - return slope, intercept + attention_intercept + return slope + draft_slope, intercept + attention_intercept + + @classmethod + @override + def _get_per_request_layer_tokens(cls, *args, layer_specs=None, **kwargs): + # ``layer_specs`` holds only the appended draft layers here; the + # recurrent and target layers keep the base estimate (None for hybrids). + return super()._get_per_request_layer_tokens(*args, layer_specs=None, **kwargs) @classmethod def _extra_state_bytes_per_rank(cls, model_config, mapping: Mapping, **kwargs) -> int: @@ -889,7 +1004,61 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC # Adds the FP4 MLA HP-ring virtual layers and the fallback-layer # lifecycle; SSM layers are left unchanged. result = layout_policy._build_cache_config(self, result) - return result + return self._with_draft_pool_ratio(result) + + def _with_draft_pool_ratio(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + """Give appended draft layers of their own lifecycle a pool share. + + A configured ``pool_ratio`` lists the target's lifecycles; the draft + lifecycle's share is inserted for token parity (see + ``_insert_draft_pool_ratio``). Native lifecycles are registered in + layer and buffer order. + """ + ratio = config.initial_pool_ratio + if ratio is None or not self._appended_draft_layer_specs: + return config + keys: List[tuple] = [] + block_bytes: List[int] = [] + draft_keys = set() + target_keys = set() + for layer in config.layers: + # Empty buffers are dropped before the native manager sees them. + for buffer in (buffer for buffer in layer.buffers if buffer.size): + policy = ( + buffer.reuse_policy if buffer.reuse_policy is not None else layer.reuse_policy + ) + if isinstance(layer, SsmLayerConfig): + key = ("ssm", policy) + else: + life_cycle = AttnLifeCycle.make( + layer.sliding_window_size, + layer.num_sink_tokens, + config.tokens_per_block, + policy, + ) + key = (life_cycle.window_size, life_cycle.num_sink_blocks, policy) + if key not in keys: + keys.append(key) + block_bytes.append(0) + block_bytes[keys.index(key)] += buffer.size + is_draft = int(layer.layer_id) in self._appended_draft_layer_specs + (draft_keys if is_draft else target_keys).add(key) + if len(ratio) == len(keys): + return config + if len(draft_keys) != 1 or draft_keys & target_keys or len(ratio) != len(keys) - 1: + raise ValueError( + f"kv_cache_config.pool_ratio has {len(ratio)} entries; give one per target " + f"layer group ({len(keys) - 1}) or one per layer group ({len(keys)})" + ) + full_history = [ + key[0] != "ssm" and (key[0] is None or key[0] >= self.max_seq_len) for key in keys + ] + return replace( + config, + initial_pool_ratio=_insert_draft_pool_ratio( + ratio, block_bytes, full_history, keys.index(draft_keys.pop()) + ), + ) def _setup_state_pool(self) -> None: """Resolve logical slots, bind base views, then initialize model state.""" diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 36960dd14ecb..f78988965b4c 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -1313,9 +1313,7 @@ def _draft_block_width(self, draft_model) -> int: checkpoint at max_draft_len=7, which is how both published DSpark drafters are meant to run. """ - return self.block_width_for( - self.spec_config, shift_label=getattr(draft_model, "_dspark_shift_label", False) - ) + return self.block_width_for(self.spec_config, shift_label=draft_model.shift_label) def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int @@ -1326,9 +1324,8 @@ def _draft_slot_ids( draft model rather than assumed: a DSpark drafter trained with the legacy DFlash slot layout keeps the base class' slots 1..K. """ - shift_label = getattr(draft_model, "_dspark_shift_label", False) return dflash_draft_slot_ids( - num_gens, block_size, num_draft_tokens, shift_label, device="cuda" + num_gens, block_size, num_draft_tokens, draft_model.shift_label, device="cuda" ) def _vocab_parallel_draft_tokens(self, draft_model) -> Optional[VocabParallelDraftTokens]: diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index d24c45b8bcde..eb9fece05e8e 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -63,7 +63,8 @@ validate_internal_request_auth_key) from tensorrt_llm.llmapi.llm_args import (SERVE_CLI_CONFIG_FIELDS, KvCacheConnectorConfig, SleepConfig, - TorchLlmArgs, build_cli_config) + TorchLlmArgs, build_cli_config, + drafter_reads_helix_sharded_context) from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -2820,46 +2821,84 @@ def set_cuda_device(): def _shared_draft_kv_blocker(server_configs) -> Optional[str]: - """What keeps the context and generation servers from sharing draft KV, or None.""" - with_spec = { - cfg.type - for cfg in server_configs if cfg.other_args.get("speculative_config") - } - if with_spec != {"ctx", "gen"}: + """What keeps the context and generation servers from sharing draft KV, or None. + + A config may list only one side: a DWDP context MPI world launches its + context servers from a config without generation servers, which run as a + separate launch. Such a config is checked for its own side; the peer + handshake rejects a pair whose sides disagree. + """ + if not server_configs or not all( + cfg.other_args.get("speculative_config") for cfg in server_configs): return ("speculative_config in both context_servers and " "generation_servers") - if len({(cfg.other_args.get("tensor_parallel_size", 1), - bool(cfg.other_args.get("enable_attention_dp", False))) - for cfg in server_configs}) != 1: - return ("the same tensor_parallel_size and enable_attention_dp in " - "context_servers and generation_servers") + # Ranks hold the same KV heads when their attention TP matches: attention + # DP keeps every head on each rank. + if len({ + 1 if cfg.other_args.get("enable_attention_dp", False) else + cfg.other_args.get("tensor_parallel_size", 1) + for cfg in server_configs + }) != 1: + return ("the same attention tensor parallelism (tensor_parallel_size, " + "or 1 with enable_attention_dp) in context_servers and " + "generation_servers") + if any( + cfg.other_args.get("context_parallel_size", 1) != 1 + for cfg in server_configs + ) and not all( + drafter_reads_helix_sharded_context( + cfg.other_args["speculative_config"].get("decoding_type"), + cfg.other_args["speculative_config"].get("attention_backend")) + for cfg in server_configs): + return ("context_parallel_size 1, or DSpark with attention_backend " + "CUTEDSL, in both sections") for cfg in server_configs: spec = cfg.other_args["speculative_config"] backend = spec.get("attention_backend") if not ((spec.get("decoding_type") == "DFlash" and backend == "TRTLLM") or (spec.get("decoding_type") == "DSpark" - and backend in (None, "AUTO", "TRTLLM"))): + and backend in (None, "AUTO", "TRTLLM", "CUTEDSL"))): return ("speculative_config to be DFlash with attention_backend " - "TRTLLM, or DSpark with attention_backend AUTO or TRTLLM, " - "in both sections") + "TRTLLM, or DSpark with attention_backend AUTO, TRTLLM " + "or CUTEDSL, in both sections") if (spec.get("draft_len_schedule") is not None or spec.get("max_concurrency") is not None): return "draft_len_schedule and max_concurrency to be unset in both sections" if cfg.other_args.get("attn_backend", "TRTLLM") != "TRTLLM": return "attn_backend TRTLLM in both sections" - if (cfg.other_args.get("pipeline_parallel_size", 1) != 1 - or cfg.other_args.get("context_parallel_size", 1) != 1): - return "pipeline_parallel_size and context_parallel_size to be 1 in both sections" + if cfg.other_args.get("pipeline_parallel_size", 1) != 1: + return "pipeline_parallel_size 1 in both sections" if cfg.other_args.get("kv_cache_compression_config") is not None: return "kv_cache_compression_config to be unset in both sections" kv_cache = cfg.other_args.get("kv_cache_config") or {} if str(kv_cache.get("dtype", "auto")).lower().startswith("mixed_"): return "a uniform KV cache dtype in both sections" - if kv_cache.get("tokens_per_block", 32) not in (16, 32, 64): + tokens_per_block = kv_cache.get("tokens_per_block", 32) + if spec.get("decoding_type") == "DSpark" and backend == "CUTEDSL": + # The CuTe DSL MLA block decode pages by any divisor of 128 above 1. + if tokens_per_block <= 1 or 128 % tokens_per_block: + return ( + "kv_cache_config.tokens_per_block to divide 128 for " + "DSpark with attention_backend CUTEDSL in both sections") + elif tokens_per_block not in (16, 32, 64): return "kv_cache_config.tokens_per_block to be 16, 32, or 64 in both sections" return None +def _shared_draft_kv_default(server_configs, blocker: Optional[str]) -> str: + """Default TRTLLM_DISAGG_SHARED_DRAFT_KV value set by the launcher. + + The launcher sets it when the variable is unset or "auto". Only a config + that lists both sides and passes the check turns sharing on. + A one-sided config cannot see how the other side is launched, and a + generation server started with the variable unset keeps a separate draft + KV cache, so it gets "0"; sharing across separate launches needs the + variable set to "1" on both. + """ + sides = {cfg.type for cfg in server_configs} + return "auto" if blocker is None and sides == {"ctx", "gen"} else "0" + + @click.command("disaggregated_mpi_worker") @stability_option( "-c", @@ -2915,8 +2954,8 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, raise click.UsageError( f"TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires {blocker}.") if os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") in (None, "auto"): - os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("0" - if blocker else "auto") + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = _shared_draft_kv_default( + disagg_cfg.server_configs, blocker) # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f94c0fa58038..e1b2f0d6999b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4027,6 +4027,15 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.DSPARK +def drafter_reads_helix_sharded_context( + decoding_type: Optional[str], attention_backend: Optional[str]) -> bool: + """Whether a speculative_config's drafter reads a helix-sharded context. + + Only the standalone MLA DSpark drafter's CUTEDSL block decode does. + """ + return decoding_type == "DSpark" and attention_backend == "CUTEDSL" + + class AutoDecodingConfig(DecodingBaseConfig): """Configuration for auto speculative decoding. diff --git a/tests/integration/defs/accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/__init__.py b/tests/integration/defs/accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/integration/defs/accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py b/tests/integration/defs/accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py new file mode 100644 index 000000000000..ce9fbec89377 --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Full-GSM8K gates for the NVFP4 Kimi K3 target kimi_k3_nvfp4__sm_107__default on SM107. + +The target runs in disaggregated serving: one-GPU DWDP context servers and a +generation server laid out as attention DP16 with expert parallel EP16, Helix +CP8, or Helix CP16 with MoE TP4 x EP4. Each case runs one multi-node deployment +through the perf-sanity disaggregated harness (``jenkins/scripts/perf/submit.py``; +``jenkins/scripts/perf/local/submit.py`` outside CI) from a +``tests/scripts/perf-sanity/disaggregated/vr200_*`` config, and replaces the +benchmark client with GSM8K against the disaggregated frontend. Every role of the +harness runs the same test id; only the client role evaluates and asserts. + +The configs require modeling_v2 on every worker, so an unsupported deployment +cannot silently test the built-in model. +""" + +import contextlib +import os +from concurrent.futures import Future +from typing import Iterator, Optional + +import openai +import pytest +import requests + +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.llmapi import ( + CompletionOutput, + DSparkDecodingConfig, + RequestOutput, + SamplingParams, +) +from tensorrt_llm.llmapi.llm_args import LlmArgs +from tensorrt_llm.llmapi.tokenizer import load_hf_tokenizer + +from .....conftest import llm_models_root +from ....accuracy_core import GSM8K, LlmapiAccuracyTestHarness +from ....test_disaggregated_serving import DuckLLM, MyThreadPoolExecutor, Result + +_MODEL_NAME = "moonshotai/Kimi-K3" +# The perf-sanity configs name the same directories (model_dir_name and +# speculative_model, both under the models root). +_NVFP4_MODEL_PATH = f"{llm_models_root()}/Kimi-K3-NVFP4" +_DSPARK_MODEL_PATH = f"{llm_models_root()}/Inferact-Kimi-K3-DSpark" + +# Deployment config -> DSpark draft length used by its generation server. +_CASES = { + # Context on 16 one-GPU DWDP servers; generation on Helix CP8, batch 20. + "disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs20_dspark5": 5, + # Generation on Helix CP8 with experts over EP8, batch 8. + "disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs8_dspark7": 7, + # Low-latency generation on Helix CP16 with experts over MoE TP4 x EP4. + "disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx12_dwdp12_gen1_helix16_moetp4ep4_bs4_dspark7": 7, + # Generation with attention DP over 16 ranks and experts over EP16. + "disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_dep16_bs16_dspark5": 5, +} + + +@contextlib.contextmanager +def _frontend_llm(base_url: str, max_draft_len: int, max_workers: int) -> Iterator[DuckLLM]: + """An LLM stand-in that sends GSM8K's requests to the disaggregated frontend.""" + response = requests.get(f"{base_url}/models", timeout=60) + response.raise_for_status() + served_model = response.json()["data"][0]["id"] + # Seconds per request attempt; a stalled request ends well inside the test timeout. + client = openai.OpenAI(api_key="unused", base_url=base_url, timeout=3600) + # The fields that pick the GSM8K reference row: the NVFP4 export declares + # MIXED_PRECISION, the workers run an NVFP4 KV cache and DSpark. + args = LlmArgs(model=_NVFP4_MODEL_PATH) + args.quant_config.quant_algo = "MIXED_PRECISION" + args.quant_config.kv_cache_quant_algo = "NVFP4" + args.speculative_config = DSparkDecodingConfig( + max_draft_len=max_draft_len, speculative_model=_DSPARK_MODEL_PATH + ) + tokenizer = load_hf_tokenizer(_NVFP4_MODEL_PATH, trust_remote_code=True) + + with MyThreadPoolExecutor(max_workers=max_workers) as thread_pool: + + def send_request( + prompt: str, sampling_params: Optional[SamplingParams], streaming: bool + ) -> RequestOutput: + if sampling_params is None: + sampling_params = SamplingParams() + response = client.completions.create( + model=served_model, + prompt=prompt, + stream=streaming, + max_tokens=sampling_params.max_tokens, + n=sampling_params.n, + temperature=( + sampling_params.temperature if sampling_params.top_p is not None else 0 + ), + top_p=sampling_params.top_p, + stop=sampling_params.stop, + seed=sampling_params.seed, + ) + result = Result( + id=0, + sampling_params=sampling_params, + outputs=[ + CompletionOutput(text=choice.text, index=index) + for index, choice in enumerate(response.choices) + ], + ) + output = RequestOutput._from_generation_result(result, prompt=prompt) + setattr(output, "result", result.result) + return output + + def generate_async( + prompt: str, sampling_params: Optional[SamplingParams] = None, streaming: bool = False + ) -> Future: + future = thread_pool.submit(send_request, prompt, sampling_params, streaming) + thread_pool.futures.append(future) + return future + + yield DuckLLM(args, tokenizer, generate_async) + + +def _run_disaggregated_gsm8k(output_dir: str, test_case: str, max_draft_len: int) -> None: + """Run ``test_case``'s deployment in the role this process was launched as. + + The client role runs GSM8K against the frontend and asserts the score; + the worker and frontend roles serve until the client is done. + """ + from .....perf.test_perf_sanity import AccuracyConfig, PerfSanityTestConfig + + assert "DISAGG_SERVING_TYPE" in os.environ, ( + "launch through the multi-node disaggregated harness (jenkins/scripts/perf/submit.py)" + ) + + class Gsm8kGate(AccuracyConfig): + def __init__(self, max_workers: int) -> None: + super().__init__({"enable_accuracy_test": True}) + self.max_workers = max_workers + self.evaluated = False + + def run( + self, + model_name: str, + server_hostname: str, + server_port: int, + output_dir: str, + server_idx: int, + ) -> None: + base_url = f"http://{server_hostname}:{server_port}/v1" + with _frontend_llm(base_url, max_draft_len, self.max_workers) as llm: + GSM8K(_MODEL_NAME).evaluate(llm) + self.evaluated = True + + config = PerfSanityTestConfig(test_case, output_dir) + config.parse_config_file() + commands = config.get_commands() + gates = [] + for clients in commands.client_configs.values(): + for client in clients: + gate = Gsm8kGate(max_workers=client.concurrency) + client.only_run_accuracy = True + client.accuracy_config = gate + gates.append(gate) + assert gates, f"{test_case} defines no client to run GSM8K from" + config.run_ex(commands) + if commands.disagg_serving_type == "BENCHMARK": + assert all(gate.evaluated for gate in gates), "GSM8K did not run against every frontend" + + +@pytest.mark.skipif( + get_sm_version() != 107, reason="This modeling_v2 Kimi K3 target requires SM107 (Rubin)." +) +@pytest.mark.timeout(7200) +class TestModelingV2KimiK3Nvfp4Sm107Default(LlmapiAccuracyTestHarness): + """GSM8K against each supported disaggregated deployment of the target.""" + + MODEL_NAME = _MODEL_NAME + MODEL_PATH = _NVFP4_MODEL_PATH + + @pytest.mark.parametrize("test_case", list(_CASES)) + def test_gsm8k(self, output_dir: str, test_case: str) -> None: + _run_disaggregated_gsm8k(output_dir, test_case, max_draft_len=_CASES[test_case]) diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index f2821dbfb4b6..0bc7a427f306 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -379,6 +379,13 @@ moonshotai/Kimi-K3: - quant_algo: MIXED_PRECISION kv_cache_quant_algo: NVFP4 accuracy: 95.4 + # The same NVFP4 KV with DSpark, which verifies against the target: the + # disaggregated SM107 modeling_v2 target + # (accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py). + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: NVFP4 + spec_dec_algo: DSpark + accuracy: 95.4 nvidia/Nemotron-MOE: - accuracy: 88.249 - quant_algo: FP8 diff --git a/tests/integration/defs/perf/test_perf_sanity.py b/tests/integration/defs/perf/test_perf_sanity.py index f874cd8570d0..c51c1e1dcf1e 100644 --- a/tests/integration/defs/perf/test_perf_sanity.py +++ b/tests/integration/defs/perf/test_perf_sanity.py @@ -3050,11 +3050,13 @@ def _hostnames_dir(self, server_idx: int) -> str: Reused across runs; slurm_launch_draft.sh clears it before each attempt. The step id is excluded so all roles (separate srun steps) agree on it. - A DWDP ctx task sees SLURM_JOB_ID only under DWDP_SAVED_ENV_PREFIX. + Every role reads the job id from its own environment only: slurm_run.sh + clears the Slurm variables of the BENCHMARK and DISAGG_SERVER tasks and + trtllm-llmapi-launch those of the worker pytest, so a DWDP ctx task must + not take the copy dwdp_ctx_rank.sh saved, or it would publish its address + where no other role looks. """ - run_id = os.environ.get( - "SLURM_JOB_ID", os.environ.get(f"{DWDP_SAVED_ENV_PREFIX}SLURM_JOB_ID", "local") - ) + run_id = os.environ.get("SLURM_JOB_ID", "local") return os.path.join(self.test_output_dir, f"hostnames-{run_id}-{server_idx}") def _hostname_file(self, server_idx: int, worker_name: str = "") -> str: @@ -3639,7 +3641,8 @@ def _write_dwdp_mpi_worker_config( urls = [] for rank in range(world_size): path = os.path.join(reserved_dir, f"{rank}.txt") - urls.append(wait_for_reported_addr(path, max(1, int(deadline - time.time())))) + host, port = wait_for_reported_addr(path, max(1, int(deadline - time.time()))) + urls.append(f"{host}:{port}") # Same LLM args a plain ctx server gets from its --config file, with # the speculative model / load balancer paths already resolved. diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 02b7e2a083a4..58453564cce7 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -17,6 +17,12 @@ accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[sa] # ModelingV2 Kimi K3: 16 GB200 GPUs in one NVLink domain (SM100). accuracy/modeling_v2/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/test_accuracy.py::TestModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4::test_gsm8k accuracy/modeling_v2/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/test_accuracy.py::TestModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1::test_gsm8k +# ModelingV2 Kimi K3 on VR200 (SM107), disaggregated: launched through the +# perf-sanity disaggregated harness (jenkins/scripts/perf/submit.py), 6-8 nodes. +accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py::TestModelingV2KimiK3Nvfp4Sm107Default::test_gsm8k[disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs20_dspark5] +accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py::TestModelingV2KimiK3Nvfp4Sm107Default::test_gsm8k[disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs8_dspark7] +accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py::TestModelingV2KimiK3Nvfp4Sm107Default::test_gsm8k[disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx12_dwdp12_gen1_helix16_moetp4ep4_bs4_dspark7] +accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py::TestModelingV2KimiK3Nvfp4Sm107Default::test_gsm8k[disagg-e2e-vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_dep16_bs16_dspark5] # ModelingV2 GLM-5.2 on VR200 (SM107), TP4/TP8/TP16. accuracy/modeling_v2/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/test_accuracy.py::TestModelingV2Glm52Nvfp4Sm107LowLatency::test_mtp7_gsm8k_and_acceptance accuracy/modeling_v2/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/test_accuracy.py::TestModelingV2Glm52Nvfp4Sm107LowLatency::test_tp4_mtp5_gsm8k_and_acceptance diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ff8f27ec4469..7165411785c7 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -144,6 +144,15 @@ l0_cpu: - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/test_tp16_moetp16ep1_uncertified_generic_calls.py - unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/test_modeling_v2_gpt_oss_default_shards.py + # The Kimi K3 Rubin target's cpu_only cases: CuTe DSL hint tables, the context + # MLA split-K dispatch gates, the FP4 MLA decode plan geometry and the + # drafter's binding to the unified draft KV layers. + # Its GPU tests are in l0_gr100. + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3r_dsl_hints.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_features.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_fp4_mla_decode.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_ctx_split_k.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_unified_draft_binding.py # MiniMax-M3 sm_107 tp2 target: its device-free (cpu_only) tests. - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 4b680ce39474..c038fb43013d 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -33,6 +33,9 @@ l0_gb300_multi_gpus: - accuracy/modeling_v2/deepseek_v3/r1_0528_nvfp4__sm_10x__default/test_accuracy.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[dep4] - accuracy/modeling_v2/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_accuracy.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k - accuracy/modeling_v2/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_accuracy.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/comm/test_k3_lamport_ar.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_allreduce.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_push_cells.py # qwen4_exp target, post-merge: three drafts over quantized KV -- NVFP4 KV at tp4 with the # breakable prefill graph, FP8 KV at dep4 (MTP under attention DP), FP8 KV at tep4 on # FlashInfer's ring replay, NVFP4 KV at tep4 without block reuse -- and dep2 context and diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index 7f2c3b87b10a..4b148524b19c 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/tests/integration/test_lists/test-db/l0_gr100.yml @@ -87,3 +87,24 @@ l0_gr100: - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec - unittest/_torch/visual_gen/kernels/parallel/test_fmha_dense_fp8.py - unittest/_torch/visual_gen/kernels/parallel/test_fused_dit_qk_norm_rope_fp8.py + # The Kimi K3 Rubin target's unit tests, one entry per op directory plus its + # top-level files; core/ and ops/gemm/ are listed per file because their + # all-reduce and GEMM-push tests run several ranks. The multi-rank tests + # (comm/, core all-reduce, ops/gemm push) are in l0_gb300_multi_gpus until + # there is a multi-GPU GR100 stage. Its cpu_only cases are in l0_cpu. + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_linear.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_mxfp8_gemm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_norm_mlp_embedding.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_cells.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_engine_probes.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_glue.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_ull_shared_fork.py diff --git a/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx12_dwdp12_gen1_helix16_moetp4ep4_bs4_dspark7.yaml b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx12_dwdp12_gen1_helix16_moetp4ep4_bs4_dspark7.yaml new file mode 100644 index 000000000000..30fb4be4b76b --- /dev/null +++ b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx12_dwdp12_gen1_helix16_moetp4ep4_bs4_dspark7.yaml @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# GSM8K deployment of the kimi_k3_nvfp4__sm_107__default target +# (accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py): +# one generation server with Helix CP16 and routed experts over MoE TP4 x EP4, +# batch 4, DSpark 7, in its low-latency mode, fed by 12 one-GPU DWDP context +# servers. 7 Rubin nodes. The test replaces the client with GSM8K against the +# disaggregated frontend. +metadata: + model_name: kimi_k3_nvfp4 + precision: fp4 + model_dir_name: Kimi-K3-NVFP4 + supported_gpus: [VR200] +slurm: + timeout: 3600 + numa_bind: false +benchmark: + mode: e2e + concurrency_list: '32' +hardware: + gpus_per_node: 4 + num_ctx_servers: 12 + num_gen_servers: 1 +environment: + # Both roles store the BF16 shared and latent MLP projections as FP8 + # (KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1; without it the context servers run out of + # memory) and keep the KV cache of the last two MLA layers in FP8. + # UCX_TLS keeps the KV cache transfer on the CUDA IPC, shared-memory and TCP + # transports: the servers reach each other over NVLink. With the RDMA transports + # also enabled, the KV cache pool sized during the startup memory profiling stays + # registered with UCX after it is released, and its memory does not come back + # for the final KV cache allocation. + worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_MODELING_V2_VALIDATE=1 TRTLLM_DISAGG_SHARED_DRAFT_KV=1 + TRTLLM_NIXL_NUM_THREADS=1 TRTLLM_KV_TRANSFER_NUM_THREADS=8 TLLM_KIMI_K3_FP8_KV_LAYERS=91,92 + KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + NCCL_GRAPH_MIXING_SUPPORT=0 TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 + TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES=2097152 TLLM_NUMA_AWARE_WORKER_AFFINITY=pin_executor + UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp + ctx_worker_env_var: TLLM_FMHA_LIBS=-flashinfer_trtllm_gen TLLM_MOE_ASYNC_INPUT_QUANT=0 + TLLM_ROUTING_DISABLE_COOP_LAUNCH=1 +worker_config: + ctx: + modeling_v2: require + attention_dp_config: + kv_cache_routing_conversation_affinity: true + kv_cache_routing_new_conv_placement: least_queued + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + cuda_graph_config: null + disable_overlap_scheduler: true + dwdp_config: + contention_opt: false + copy_mode: gather + dwdp_size: 12 + gather_kernel: tma + num_experts_per_worker: 82 + num_groups: 1 + num_prefetch_experts: 74 + enable_attention_dp: true + enable_chunked_prefill: true + kv_cache_config: + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: true + free_gpu_memory_fraction: 0.7 + host_cache_size: 171798691840 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + # Layer-group order: ssm, fp4_mla, fp8_mla, mla_hp_tail. + pool_ratio: [0.14, 0.635, 0.205, 0.02] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 16 + max_num_tokens: 16384 + max_seq_len: 996672 + moe_config: + backend: CUTEDSL + max_num_tokens: 16384 + use_low_precision_moe_combine: true + moe_expert_parallel_size: 1 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + stream_interval: 10 + tensor_parallel_size: 1 + trust_remote_code: true + speculative_config: + block_size: 7 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 7 + speculative_model: Inferact-Kimi-K3-DSpark + advanced_sampling_mode: no_topk_no_topp + gen: + enable_min_latency: true + modeling_v2: require + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + context_parallel_size: 16 + cp_config: + cp_type: HELIX + fifo_version: 2 + tokens_per_block: 128 + use_nccl_for_alltoall: false + cuda_graph_config: + batch_sizes: [1, 2, 3, 4] + enable_padding: true + enable_attention_dp: false + kv_cache_config: + avg_seq_len: 149442 + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: false + free_gpu_memory_fraction: 0.8 + host_cache_size: 0 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 4 + max_num_tokens: 32 + max_seq_len: 997696 + moe_config: + backend: CUTEDSL_FC12 + use_low_precision_moe_combine: true + moe_expert_parallel_size: 4 + moe_tensor_parallel_size: 4 + num_postprocess_workers: 4 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + speculative_config: + block_size: 7 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 7 + speculative_model: Inferact-Kimi-K3-DSpark + advanced_sampling_mode: no_topk_no_topp + stream_interval: 100 + tensor_parallel_size: 1 + trust_remote_code: true diff --git a/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_dep16_bs16_dspark5.yaml b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_dep16_bs16_dspark5.yaml new file mode 100644 index 000000000000..b366705f41a8 --- /dev/null +++ b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_dep16_bs16_dspark5.yaml @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# GSM8K deployment of the kimi_k3_nvfp4__sm_107__default target +# (accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py): +# one generation server with attention DP over 16 ranks and routed experts over +# EP16, batch 16, DSpark 5, fed by 16 one-GPU DWDP context servers. 8 Rubin +# nodes. The test replaces the client with GSM8K against the disaggregated frontend. +metadata: + model_name: kimi_k3_nvfp4 + precision: fp4 + model_dir_name: Kimi-K3-NVFP4 + supported_gpus: [VR200] +slurm: + timeout: 3600 + numa_bind: false +benchmark: + mode: e2e + concurrency_list: '128' +hardware: + gpus_per_node: 4 + num_ctx_servers: 16 + num_gen_servers: 1 +environment: + # Both roles store the BF16 shared and latent MLP projections as FP8 + # (KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1; without it the context servers run out of + # memory) and keep the KV cache of the last two MLA layers in FP8. + # UCX_TLS keeps the KV cache transfer on the CUDA IPC, shared-memory and TCP + # transports: the servers reach each other over NVLink. With the RDMA transports + # also enabled, the KV cache pool sized during the startup memory profiling stays + # registered with UCX after it is released, and its memory does not come back + # for the final KV cache allocation. + worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_MODELING_V2_VALIDATE=1 TRTLLM_DISAGG_SHARED_DRAFT_KV=1 + TRTLLM_NIXL_NUM_THREADS=1 TRTLLM_KV_TRANSFER_NUM_THREADS=8 TLLM_KIMI_K3_FP8_KV_LAYERS=91,92 + KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + NCCL_GRAPH_MIXING_SUPPORT=0 TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 + UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp + ctx_worker_env_var: TLLM_FMHA_LIBS=-flashinfer_trtllm_gen TLLM_MOE_ASYNC_INPUT_QUANT=0 + TLLM_ROUTING_DISABLE_COOP_LAUNCH=1 TLLM_V2_CTX_DEMOTE_MAX_ITERS=8 + TLLM_V2_CTX_DEMOTE_UNCACHED_TOKENS=16384 +worker_config: + ctx: + modeling_v2: require + attention_dp_config: + kv_cache_routing_conversation_affinity: true + kv_cache_routing_new_conv_placement: least_queued + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + cuda_graph_config: null + disable_overlap_scheduler: true + dwdp_config: + contention_opt: false + copy_mode: gather + dwdp_size: 16 + gather_kernel: tma + num_experts_per_worker: 56 + num_groups: 1 + num_prefetch_experts: 56 + enable_attention_dp: true + enable_chunked_prefill: true + kv_cache_config: + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: true + free_gpu_memory_fraction: 0.7 + host_cache_size: 171798691840 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + enable_branch_snapshot: true + # Layer-group order: ssm, fp4_mla, fp8_mla, mla_hp_tail. + pool_ratio: [0.455, 0.458333, 0.066667, 0.02] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 16 + max_num_tokens: 16384 + max_seq_len: 996672 + moe_config: + backend: CUTEDSL + max_num_tokens: 16384 + use_low_precision_moe_combine: true + moe_expert_parallel_size: 1 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + stream_interval: 10 + tensor_parallel_size: 1 + trust_remote_code: true + speculative_config: + block_size: 5 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 5 + speculative_model: Inferact-Kimi-K3-DSpark + gen: + attention_dp_config: + enable_kv_balanced_admission: true + modeling_v2: require + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + cuda_graph_config: + batch_sizes: [1, 2, 4, 8, 16] + enable_padding: true + enable_attention_dp: true + kv_cache_config: + avg_seq_len: 149442 + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: false + free_gpu_memory_fraction: 0.8 + host_cache_size: 0 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 16 + max_num_tokens: 96 + max_seq_len: 997696 + moe_config: + backend: MEGAMOE_CUTEDSL + use_low_precision_moe_combine: true + moe_expert_parallel_size: 16 + num_postprocess_workers: 4 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + speculative_config: + block_size: 5 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 5 + speculative_model: Inferact-Kimi-K3-DSpark + stream_interval: 100 + tensor_parallel_size: 16 + trust_remote_code: true diff --git a/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs20_dspark5.yaml b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs20_dspark5.yaml new file mode 100644 index 000000000000..d538a0dc14cd --- /dev/null +++ b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs20_dspark5.yaml @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# GSM8K deployment of the kimi_k3_nvfp4__sm_107__default target +# (accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py): +# 16 one-GPU DWDP context servers feeding one Helix CP8 generation server +# (batch 20, DSpark 5). 6 Rubin nodes. The test replaces the client with GSM8K +# against the disaggregated frontend. +metadata: + model_name: kimi_k3_nvfp4 + precision: fp4 + model_dir_name: Kimi-K3-NVFP4 + supported_gpus: [VR200] +slurm: + timeout: 3600 + numa_bind: false +benchmark: + mode: e2e + concurrency_list: '64' +hardware: + gpus_per_node: 4 + num_ctx_servers: 16 + num_gen_servers: 1 +environment: + # Both roles store the BF16 shared and latent MLP projections as FP8 + # (KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1; without it the context servers run out of + # memory) and keep the KV cache of the last two MLA layers in FP8. + # UCX_TLS keeps the KV cache transfer on the CUDA IPC, shared-memory and TCP + # transports: the servers reach each other over NVLink. With the RDMA transports + # also enabled, the KV cache pool sized during the startup memory profiling stays + # registered with UCX after it is released, and its memory does not come back + # for the final KV cache allocation. + worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_MODELING_V2_VALIDATE=1 TRTLLM_DISAGG_SHARED_DRAFT_KV=1 + TRTLLM_NIXL_NUM_THREADS=1 TRTLLM_KV_TRANSFER_NUM_THREADS=8 TLLM_KIMI_K3_FP8_KV_LAYERS=91,92 + KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + NCCL_GRAPH_MIXING_SUPPORT=0 TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 + TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES=2097152 + UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp + ctx_worker_env_var: TLLM_FMHA_LIBS=-flashinfer_trtllm_gen TLLM_MOE_ASYNC_INPUT_QUANT=0 + TLLM_ROUTING_DISABLE_COOP_LAUNCH=1 +worker_config: + ctx: + modeling_v2: require + attention_dp_config: + kv_cache_routing_conversation_affinity: true + kv_cache_routing_new_conv_placement: least_queued + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + cuda_graph_config: null + disable_overlap_scheduler: true + dwdp_config: + contention_opt: false + copy_mode: gather + dwdp_size: 16 + gather_kernel: tma + num_experts_per_worker: 56 + num_groups: 1 + num_prefetch_experts: 56 + enable_attention_dp: true + enable_chunked_prefill: true + kv_cache_config: + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: true + free_gpu_memory_fraction: 0.7 + host_cache_size: 171798691840 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + # Layer-group order: ssm, fp4_mla, fp8_mla, mla_hp_tail. + pool_ratio: [0.14, 0.635, 0.205, 0.02] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 16 + max_num_tokens: 16384 + max_seq_len: 996672 + moe_config: + backend: CUTEDSL + max_num_tokens: 16384 + use_low_precision_moe_combine: true + moe_expert_parallel_size: 1 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + stream_interval: 10 + tensor_parallel_size: 1 + trust_remote_code: true + speculative_config: + block_size: 5 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 5 + speculative_model: Inferact-Kimi-K3-DSpark + advanced_sampling_mode: no_topk_no_topp + gen: + modeling_v2: require + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + context_parallel_size: 8 + cp_config: + cp_type: HELIX + fifo_version: 2 + tokens_per_block: 128 + use_nccl_for_alltoall: false + cuda_graph_config: + batch_sizes: [1, 2, 4, 8, 16, 20] + enable_padding: true + enable_attention_dp: false + kv_cache_config: + avg_seq_len: 149442 + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: false + free_gpu_memory_fraction: 0.8 + host_cache_size: 0 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 20 + max_num_tokens: 120 + max_seq_len: 997696 + moe_config: + backend: CUTEDSL_FC12 + use_low_precision_moe_combine: true + moe_expert_parallel_size: 4 + moe_tensor_parallel_size: 2 + num_postprocess_workers: 4 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + speculative_config: + block_size: 5 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 5 + speculative_model: Inferact-Kimi-K3-DSpark + advanced_sampling_mode: no_topk_no_topp + stream_interval: 100 + tensor_parallel_size: 1 + trust_remote_code: true diff --git a/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs8_dspark7.yaml b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs8_dspark7.yaml new file mode 100644 index 000000000000..44f46e5d28c5 --- /dev/null +++ b/tests/scripts/perf-sanity/disaggregated/vr200_kimi-k3-nvfp4-dspark_gsm8k_ctx16_dwdp16_gen1_helix8_bs8_dspark7.yaml @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# GSM8K deployment of the kimi_k3_nvfp4__sm_107__default target +# (accuracy/modeling_v2/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_accuracy.py): +# one generation server with Helix CP8 and routed experts over EP8, batch 8, +# DSpark 7, fed by 16 one-GPU DWDP context servers. 6 Rubin nodes. The test +# replaces the client with GSM8K against the disaggregated frontend. +metadata: + model_name: kimi_k3_nvfp4 + precision: fp4 + model_dir_name: Kimi-K3-NVFP4 + supported_gpus: [VR200] +slurm: + timeout: 3600 + numa_bind: false +benchmark: + mode: e2e + concurrency_list: '64' +hardware: + gpus_per_node: 4 + num_ctx_servers: 16 + num_gen_servers: 1 +environment: + # Both roles store the BF16 shared and latent MLP projections as FP8 + # (KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1; without it the context servers run out of + # memory) and keep the KV cache of the last two MLA layers in FP8. + # UCX_TLS keeps the KV cache transfer on the CUDA IPC, shared-memory and TCP + # transports: the servers reach each other over NVLink. With the RDMA transports + # also enabled, the KV cache pool sized during the startup memory profiling stays + # registered with UCX after it is released, and its memory does not come back + # for the final KV cache allocation. + worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_MODELING_V2_VALIDATE=1 TRTLLM_DISAGG_SHARED_DRAFT_KV=1 + TRTLLM_NIXL_NUM_THREADS=1 TRTLLM_KV_TRANSFER_NUM_THREADS=8 TLLM_KIMI_K3_FP8_KV_LAYERS=91,92 + KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + NCCL_GRAPH_MIXING_SUPPORT=0 TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 + TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES=2097152 + UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp + ctx_worker_env_var: TLLM_FMHA_LIBS=-flashinfer_trtllm_gen TLLM_MOE_ASYNC_INPUT_QUANT=0 + TLLM_ROUTING_DISABLE_COOP_LAUNCH=1 +worker_config: + ctx: + modeling_v2: require + attention_dp_config: + kv_cache_routing_conversation_affinity: true + kv_cache_routing_new_conv_placement: least_queued + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + cuda_graph_config: null + disable_overlap_scheduler: true + dwdp_config: + contention_opt: false + copy_mode: gather + dwdp_size: 16 + gather_kernel: tma + num_experts_per_worker: 56 + num_groups: 1 + num_prefetch_experts: 56 + enable_attention_dp: true + enable_chunked_prefill: true + kv_cache_config: + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: true + free_gpu_memory_fraction: 0.7 + host_cache_size: 257698037760 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + enable_branch_snapshot: true + # Layer-group order: ssm, fp4_mla, fp8_mla, mla_hp_tail. + pool_ratio: [0.14, 0.635, 0.205, 0.02] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 16 + max_num_tokens: 16384 + max_seq_len: 996672 + moe_config: + backend: CUTEDSL + max_num_tokens: 16384 + use_low_precision_moe_combine: true + moe_expert_parallel_size: 1 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + stream_interval: 10 + tensor_parallel_size: 1 + trust_remote_code: true + speculative_config: + block_size: 7 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 7 + speculative_model: Inferact-Kimi-K3-DSpark + advanced_sampling_mode: no_topk_no_topp + gen: + modeling_v2: require + cache_transceiver_config: + backend: NIXL + kv_cache_bounce_size_mb: 0 + kv_transfer_timeout_ms: 600000 + transceiver_runtime: PYTHON + context_parallel_size: 8 + cp_config: + cp_type: HELIX + fifo_version: 2 + tokens_per_block: 128 + use_nccl_for_alltoall: false + cuda_graph_config: + batch_sizes: [1, 2, 4, 8] + enable_padding: true + enable_attention_dp: false + kv_cache_config: + avg_seq_len: 149442 + block_reuse_config: + max_num_turns: 5 + policy: per_conversation + dtype: nvfp4 + enable_block_reuse: false + free_gpu_memory_fraction: 0.8 + host_cache_size: 0 + mamba_ssm_cache_dtype: bfloat16 + mamba_state_config: + additional_snapshot_offsets_from_end: [0] + tokens_per_block: 128 + use_kv_cache_manager_v2: true + max_batch_size: 8 + max_num_tokens: 64 + max_seq_len: 997696 + moe_config: + backend: CUTEDSL_FC12 + use_low_precision_moe_combine: true + moe_expert_parallel_size: 8 + num_postprocess_workers: 4 + num_serve_frontends: 8 + pipeline_parallel_size: 1 + print_iter_log: true + speculative_config: + block_size: 7 + decoding_type: DSpark + attention_backend: CUTEDSL + max_draft_len: 7 + speculative_model: Inferact-Kimi-K3-DSpark + advanced_sampling_mode: no_topk_no_topp + stream_interval: 100 + tensor_parallel_size: 1 + trust_remote_code: true diff --git a/tests/test_common/error_utils.py b/tests/test_common/error_utils.py index 65c9b905dc4f..fedcc43f3bc0 100644 --- a/tests/test_common/error_utils.py +++ b/tests/test_common/error_utils.py @@ -42,6 +42,15 @@ r"\s*WARNING: Retrying \(Retry\([^)]*\)\) after connection broken by " r"'(?:Read|Connect)TimeoutError\(.*\)': \S+\s*" ) +# A multi-frontend server's secondary frontends forward the disaggregated +# server's clock-offset POST to the launcher frontend. Before the launcher +# listens, the forward is refused and the POST gets a 503; the disaggregated +# server then leaves that server's perf-metric timestamps unadjusted. +LAUNCHER_CLOCK_OFFSET_REFUSED = re.compile( + r".*\[W\] \[serve\] Forwarding POST /steady_clock_offset to the launcher frontend failed: " + r"\w*ClientConnectorError\(.*ConnectionRefusedError\(111, 'Connection refused'\)\)\s*" +) +LAUNCHER_CLOCK_OFFSET_KEYWORDS = ("ClientConnectorError", "ConnectionRefusedError") def is_benign_line(line: str) -> bool: @@ -52,6 +61,12 @@ def is_benign_line(line: str) -> bool: # warning lines containing another error keyword must still fail the scan. if PIP_TIMEOUT_RETRY.fullmatch(line): return not any(keyword in line for keyword in ERROR_KEYWORDS if keyword != "TimeoutError") + if LAUNCHER_CLOCK_OFFSET_REFUSED.fullmatch(line): + return not any( + keyword in line + for keyword in ERROR_KEYWORDS + if keyword not in LAUNCHER_CLOCK_OFFSET_KEYWORDS + ) return AUTOTUNER_MARKER in line and any(text in line for text in AUTOTUNER_BENIGN_TEXTS) diff --git a/tests/unittest/_torch/attention/test_mla_fp8_context_inputs.py b/tests/unittest/_torch/attention/test_mla_fp8_context_inputs.py new file mode 100644 index 000000000000..010d76a73863 --- /dev/null +++ b/tests/unittest/_torch/attention/test_mla_fp8_context_inputs.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MLA cached-KV context attention with caller-quantized FP8 Q / K / V. + +With an FP8 KV cache, the attention op's FP8 context MLA FMHA reads FP8 Q, K and V. Given BF16 +K / V it quantizes them itself (unit scale, round to nearest, saturating). A caller may hand it +those FP8 tensors directly (``key`` / ``value`` in E4M3, and ``quant_q_buffer`` with +``quant_scale_qkv``); the output must be bit-identical to handing it the BF16 inputs. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + RopeParams, +) +from tensorrt_llm._torch.attention.backends.utils import get_attention_backend +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +NUM_HEADS = 32 +Q_LORA_RANK = 1536 +KV_LORA_RANK = 512 +QK_NOPE_HEAD_DIM = 128 +QK_ROPE_HEAD_DIM = 64 +V_HEAD_DIM = 128 +QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM +MAX_POSITIONS = 4096 +FP8_MAX = 448.0 + + +def _make_layer(): + rope_config = SimpleNamespace( + hidden_size=7168, + num_attention_heads=NUM_HEADS, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 40.0, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + max_position_embeddings=MAX_POSITIONS, + rope_theta=10000.0, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + model_type="deepseek_v3", + ) + return get_attention_backend("TRTLLM")( + layer_idx=0, + num_heads=NUM_HEADS, + head_dim=QK_HEAD_DIM, + num_kv_heads=NUM_HEADS, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value), + q_scaling=1.0, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.yarn, + rope=RopeParams.from_config(rope_config), + is_neox=False, + ), + mla_params=MLAParams( + q_lora_rank=Q_LORA_RANK, + kv_lora_rank=KV_LORA_RANK, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + qk_nope_head_dim=QK_NOPE_HEAD_DIM, + v_head_dim=V_HEAD_DIM, + predicted_tokens_per_seq=1, + ), + ) + + +def _to_e4m3(x: torch.Tensor) -> torch.Tensor: + """What the op's unit-scale copy writes: round to nearest, saturating.""" + return x.float().clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + + +@pytest.mark.parametrize( + "cached,new", [([256], [128]), ([256, 200], [128, 184])], ids=["one_request", "two_requests"] +) +def test_mla_cached_kv_context_fp8_inputs_match_bf16_inputs(cached, new) -> None: + if not torch.cuda.is_available(): + pytest.skip("needs a GPU") + layer = _make_layer() + if not layer.has_fp8_context_mla: + pytest.skip("this SM has no FP8 context MLA FMHA") + scale = layer.kv_scale_orig_quant + assert isinstance(scale, torch.Tensor) and bool(torch.all(scale == 1.0)) + + torch.manual_seed(7) + batch = len(cached) + total = [c + n for c, n in zip(cached, new)] + mapping = Mapping(world_size=1, tp_size=1, rank=0) + tokens_per_block = 32 + manager = KVCacheManagerV2( + KvCacheConfig( + max_tokens=sum(-(-t // tokens_per_block) * tokens_per_block for t in total), + enable_block_reuse=False, + ), + CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=KV_LORA_RANK + QK_ROPE_HEAD_DIM, + tokens_per_block=tokens_per_block, + max_seq_len=max(total), + max_batch_size=batch, + mapping=mapping, + dtype=str_dtype_to_binding(torch_dtype_to_str(torch.float8_e4m3fn)), + ) + metadata_cls = get_attention_backend("TRTLLM").Metadata + + def context_metadata(seq_lens, cached_lens): + metadata = metadata_cls( + seq_lens=torch.tensor(seq_lens, dtype=torch.int), + request_ids=list(range(batch)), + max_num_requests=batch, + num_contexts=batch, + prompt_lens=list(seq_lens), + max_num_tokens=sum(seq_lens), + kv_cache_manager=manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached_lens), + mapping=mapping, + enable_context_mla_with_cached_kv=True, + ) + metadata.prepare() + return metadata + + def random_inputs(num_tokens): + q = torch.randn(num_tokens, NUM_HEADS * QK_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + latent = torch.randn( + num_tokens, KV_LORA_RANK + QK_ROPE_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + return q, latent + + try: + stream = torch.cuda.current_stream().cuda_stream + caches = [] + for req_id, length in enumerate(cached): + kv_cache = manager._create_kv_cache(req_id, None, None) + assert kv_cache.resume(stream) + kv_cache.capacity = length + caches.append(kv_cache) + # First chunk: write the cached tokens' latent KV into the FP8 cache. + q, latent = random_inputs(sum(cached)) + layer.mla_rope_append_paged_kv_assign_q(q, latent, context_metadata(cached, [0] * batch)) + + for kv_cache, length in zip(caches, total): + kv_cache.capacity = length + metadata = context_metadata(new, cached) + q, latent = random_inputs(sum(new)) + layer.mla_rope_append_paged_kv_assign_q(q, latent, metadata) + compressed_kv, k_pe = layer.load_paged_kv_cache_for_mla(metadata, torch.bfloat16) + assert compressed_kv.shape[0] == sum(total) + + kv_b = ( + torch.randn( + KV_LORA_RANK, + NUM_HEADS * (QK_NOPE_HEAD_DIM + V_HEAD_DIM), + dtype=torch.bfloat16, + device="cuda", + ) + * KV_LORA_RANK**-0.5 + ) + k_nope, v = (compressed_kv @ kv_b).split( + [NUM_HEADS * QK_NOPE_HEAD_DIM, NUM_HEADS * V_HEAD_DIM], -1 + ) + k = torch.cat( + ( + k_nope.reshape(-1, NUM_HEADS, QK_NOPE_HEAD_DIM), + k_pe.view(-1, 1, QK_ROPE_HEAD_DIM).expand(-1, NUM_HEADS, -1), + ), + dim=-1, + ).reshape(-1, NUM_HEADS * QK_HEAD_DIM) + v = v.contiguous() + + def run(key, value, **fp8_q_args): + return layer.forward( + q.clone(), + key, + value, + metadata, + forward_args=AttentionForwardArgs( + attention_input_type=AttentionInputType.context_only, + latent_cache=None, + **fp8_q_args, + ), + ) + + bf16_out = run(k, v) + q_fp8 = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(q, scale)[0] + fp8_out = run(_to_e4m3(k), _to_e4m3(v), quant_q_buffer=q_fp8, quant_scale_qkv=scale) + fp8_kv_only_out = run(_to_e4m3(k), _to_e4m3(v)) + torch.cuda.synchronize() + + assert bf16_out.abs().sum() > 0 + assert torch.equal(fp8_out.view(torch.int16), bf16_out.view(torch.int16)) + # Without quant_q_buffer the op quantizes Q itself; K / V still come from the caller. + assert torch.equal(fp8_kv_only_out.view(torch.int16), bf16_out.view(torch.int16)) + + with pytest.raises(RuntimeError, match="same dtype"): + run(_to_e4m3(k), v) + with pytest.raises(RuntimeError, match="same dtype"): + run(k, _to_e4m3(v)) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/test_kimi_k3_draft_block_rows.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/test_kimi_k3_draft_block_rows.py index ebc4da336738..f0741cce1e13 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/test_kimi_k3_draft_block_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/test_kimi_k3_draft_block_rows.py @@ -89,7 +89,7 @@ def _is_view(rows, out): def test_rows_equal_the_stock_gather(num_gens, k, block_extra, shift_label): block = k + block_extra worker = _worker() - drafter = SimpleNamespace(_dspark_shift_label=shift_label) + drafter = SimpleNamespace(shift_label=shift_label) out = _block(num_gens, block) rows = worker._draft_block_hidden_states(drafter, out, num_gens, block, k) @@ -107,7 +107,7 @@ def test_rows_equal_the_stock_gather(num_gens, k, block_extra, shift_label): def test_dspark_decode_shape_is_a_view(): """DSpark's shift_label drafter at a block of K = 7: every batch of the decode path reads its rows in place.""" worker = _worker() - drafter = SimpleNamespace(_dspark_shift_label=True) + drafter = SimpleNamespace(shift_label=True) for num_gens in range(1, 9): out = _block(num_gens, 7) rows = worker._draft_block_hidden_states(drafter, out, num_gens, 7, 7) @@ -126,7 +126,7 @@ def test_dflash_slots_of_several_requests_are_gathered(): def test_slot_ids_are_built_once_per_shape(monkeypatch): worker = _worker() - drafter = SimpleNamespace(_dspark_shift_label=True) + drafter = SimpleNamespace(shift_label=True) built = [] real = DSparkWorker._draft_slot_ids monkeypatch.setattr( @@ -144,7 +144,7 @@ def test_capture_gathers_a_shape_it_has_not_seen(monkeypatch): """Under capture an unseen shape is gathered and not kept (building the run needs a host read); a shape seen before capture keeps its view.""" worker = _worker() - drafter = SimpleNamespace(_dspark_shift_label=True) + drafter = SimpleNamespace(shift_label=True) seen, unseen = _block(2, 7), _block(3, 7) assert _is_view(worker._draft_block_hidden_states(drafter, seen, 2, 7, 7), seen) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/comm/test_k3_lamport_ar.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/comm/test_k3_lamport_ar.py new file mode 100644 index 000000000000..18e982bc1125 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/comm/test_k3_lamport_ar.py @@ -0,0 +1,570 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pushed MNNVL all-reduce of the Kimi K3 package (``kernels/k3_lamport_ar.py``, the low-M GEMM +push epilogue and the ``layers.py`` API on main's ``AllReduce``): bitwise equal to the C++ +MNNVL all-reduce, with pushed and C++ calls interleaved on one workspace. + +``test_k3_lamport_entry_cells`` drives the CELLS of the collective entries on the shared MNNVL +workspace of the group, each call under ``validating`` and against the entry's reference, which +takes every rank's input in rank order: + +* ``k3_lamport_push``: the slots of this rank's buffer (one-shot: every token's ``W`` slots; + two-shot: the tokens this rank owns) hold every rank's sanitized rows, bit for bit; +* ``k3_lamport_finish``: the rank-order sum in the C++ kernel's order, bit for bit; +* ``k3_attn_res_lamport``: the updated prefix bit for bit, the normed mix within the family + attn_res tests' cosine / relative-L2 band, the MXFP8 output bit for bit the low-M GEMM's + quantization of the stored output; +* ``k3_latent_ar_rmsnorm``: the normed latent within the same band, its MXFP8 output (row-major + or the mid-M GEMM's atoms) bit for bit the quantization of the stored output. +""" + +import os +import pickle +import sys +import traceback + +import cloudpickle +import pytest +import torch +from mpi4py import MPI +from utils.util import skip_pre_blackwell + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__( + cloudpickle.dumps, + cloudpickle.loads, + pickle.HIGHEST_PROTOCOL, +) + +pytestmark = pytest.mark.threadleak(enabled=False) + + +def _special(m, h, gen): + """BF16 ``[m, h]`` with -0, +0, denormals, large values, inf and NaN lanes.""" + x = torch.randn(m, h, device="cuda", generator=gen) * torch.exp2( + torch.randint(-20, 20, (m, 1), device="cuda", generator=gen).float() + ) + x = x.to(torch.bfloat16) + bits = x.view(torch.int16) + r = torch.rand(m, h, device="cuda", generator=gen) + bits[r < 0.05] = -32768 + bits[(r >= 0.05) & (r < 0.08)] = 0 + den = (r >= 0.08) & (r < 0.11) + bits[den] = torch.randint(1, 127, (int(den.sum()),), device="cuda", generator=gen).to( + torch.int16 + ) + bits[(r >= 0.11) & (r < 0.115)] = 0x7F80 + bits[(r >= 0.115) & (r < 0.117)] = 0x7FC1 + return x + + +def _input(m, h, seed, rank, kind): + gen = torch.Generator(device="cuda").manual_seed(seed * 1000 + rank) + if kind == 0: + return (torch.randn(m, h, device="cuda", generator=gen) * 0.25).to(torch.bfloat16) + if kind == 1: + return _special(m, h, gen) + # Every rank sends denormals only: the sums stay denormal (no flush to zero anywhere). + bits = torch.randint(-127, 127, (m, h), device="cuda", generator=gen).to(torch.int16) + return bits.view(torch.bfloat16) + + +def _block_weights(n, k, seed): + """E4M3 ``[n, k]`` weight and its power-of-two FP32 128x128 block scales.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, device="cuda", generator=gen) * 0.5).to(torch.float8_e4m3fn) + scale = torch.exp2( + torch.randint(-8, -2, (n // 128, k // 128), device="cuda", generator=gen).float() + ) + return w, scale + + +_K3 = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) +# tests/unittest/_torch/modeling_v2, where _validating lives (the MPI workers import it by path). +_MV2_TESTS = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) +_ENTRIES = ( + "ops.comm.k3_lamport_push", + "ops.comm.k3_lamport_finish", + "ops.comm.k3_attn_res_lamport", + "ops.comm.k3_latent_ar_rmsnorm", +) + + +def _weights(n, k, seed): + import importlib + + lowm = importlib.import_module(f"{_K3}.kernels.k3_fp8w_lowm_gemm") + + return lowm.prepare_k3_fp8w_lowm_weight(*_block_weights(n, k, seed)) + + +def _run_rank(world_size): + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + import importlib + + from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy + + # Imported by name: an attribute named ``ops`` next to the global ``torch`` would make + # cloudpickle ship the ``torch.ops`` object with this function. + from tensorrt_llm._torch.distributed.ops import ( + _MNNVL_ONE_SHOT_THRESHOLD_BYTES, + get_or_scale_allreduce_mnnvl_workspace, + ) + from tensorrt_llm.mapping import Mapping + + lowm = importlib.import_module(f"{_K3}.kernels.k3_fp8w_lowm_gemm") + lamport = importlib.import_module(f"{_K3}.kernels.k3_lamport_ar") + push = importlib.import_module(f"{_K3}.ops.comm.k3_lamport_push").k3_lamport_push + finish = importlib.import_module(f"{_K3}.ops.comm.k3_lamport_finish").k3_lamport_finish + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + torch.cuda.set_device(rank) + try: + mapping = Mapping(world_size=world_size, tp_size=world_size, rank=rank) + allreduce = AllReduce( + mapping=mapping, strategy=AllReduceStrategy.MNNVL, dtype=torch.bfloat16 + ) + assert allreduce.mnnvl_allreduce is not None + workspace = get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16) + ws = lamport.LamportWorkspace.from_mnnvl(workspace, world_size, rank) + threshold = _MNNVL_ONE_SHOT_THRESHOLD_BYTES + bad = 0 + + def differ(ref, out): + return int((ref.view(torch.int16) != out.view(torch.int16)).sum()) + + # Push kernel + finish kernel against the C++ all-reduce, both transports reproducing + # the C++ kernel the size selects, with C++ calls in between. + for m, h in ((1, 7168), (14, 7168), (64, 7168), (300, 7168), (14, 3584), (5, 64)): + cpp_one_shot = lamport.one_shot_for(m, h, world_size, threshold) + for kind in (0, 1, 2): + x = _input(m, h, 3 * m + kind, rank, kind) + ref = allreduce(x) + for one_shot in (False, True): + if lamport.required_bytes(m, h, world_size, one_shot) > ws.bytes_per_buffer: + continue + out = torch.empty_like(x) + push(ws, x, one_shot) + finish(ws, out, one_shot, cpp_one_shot=cpp_one_shot) + bad += differ(ref, out) + bad += differ(ref, allreduce(x)) + + # GEMM push epilogue + finish against the GEMM followed by the C++ all-reduce. + for n, k in ((7168, 768), (7168, 384)): + w, w_scale = _weights(n, k, 11 + rank) + for m in (1, 7, 14, 32, 64): + cpp_one_shot = lamport.one_shot_for(m, n, world_size, threshold) + x = _input(m, k, 5 * m, rank, 0) + y = torch.empty(m, n, device="cuda", dtype=torch.bfloat16) + lowm.run_k3_fp8w_lowm_gemm(x, w, w_scale, y) + ref = allreduce(y) + for one_shot in (False, True): + lowm.run_k3_fp8w_lowm_gemm_ar_push(x, w, w_scale, ws, one_shot) + out = finish(ws, torch.empty_like(y), one_shot, cpp_one_shot=cpp_one_shot) + bad += differ(ref, out) + + # One CUDA graph with pushed and C++ calls on the workspace, replayed with new inputs. + n, k, m = 7168, 768, 14 + cpp_one_shot = lamport.one_shot_for(m, n, world_size, threshold) + w, w_scale = _weights(n, k, 99 + rank) + xs = [_input(m, k, 700 + i, rank, 0) for i in range(4)] + ys = [torch.empty(m, n, device="cuda", dtype=torch.bfloat16) for _ in range(4)] + outs = [torch.empty_like(y) for y in ys] + + def calls(): + res = [] + for i in range(4): + if i % 2 == 0: + lowm.run_k3_fp8w_lowm_gemm_ar_push(xs[i], w, w_scale, ws, cpp_one_shot) + res.append(finish(ws, outs[i], cpp_one_shot, cpp_one_shot=cpp_one_shot)) + else: + lowm.run_k3_fp8w_lowm_gemm(xs[i], w, w_scale, ys[i]) + res.append(allreduce(ys[i])) + return res + + calls() + torch.cuda.synchronize() + comm.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + results = calls() + for rep in range(3): + for i in range(4): + xs[i].copy_(_input(m, k, 900 + 10 * rep + i, rank, 0)) + torch.cuda.synchronize() + comm.Barrier() + graph.replay() + torch.cuda.synchronize() + for i in range(4): + y = torch.empty(m, n, device="cuda", dtype=torch.bfloat16) + lowm.run_k3_fp8w_lowm_gemm(xs[i], w, w_scale, y) + bad += differ(allreduce(y), results[i]) + del graph + torch.cuda.synchronize() + assert comm.allreduce(bad) == 0, f"rank {rank}: {bad} values differ from the C++ path" + except Exception: + traceback.print_exc() + raise + return True + + +def _run_rank_api(world_size): + """``layers`` on main's ``AllReduce``: a pushed tensor finished as a tensor and deferred + (``LamportPending``) against ``AllReduce.forward``, eagerly and in a CUDA graph.""" + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + import importlib + + from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + api = importlib.import_module(f"{_K3}.layers") + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + torch.cuda.set_device(rank) + + def differ(ref, out): + return int((ref.view(torch.int16) != out.view(torch.int16)).sum()) + + try: + mapping = Mapping(world_size=world_size, tp_size=world_size, rank=rank) + allreduce = AllReduce( + mapping=mapping, strategy=AllReduceStrategy.MNNVL, dtype=torch.bfloat16 + ) + assert allreduce.mnnvl_allreduce is not None + bad = 0 + kinds = {"plain": 0, "deferred": 0} + for m, h in ((1, 7168), (14, 7168), (64, 7168), (300, 7168), (14, 3584)): + for kind in (0, 1): + x = _input(m, h, 7 * m + kind, rank, kind) + ref = allreduce(x) + + target = api.lamport_push_target(allreduce, m, h) + assert target is not None + try: + api.finish_lamport_push(allreduce, target) + raise AssertionError("finish_lamport_push accepted a target nobody pushed") + except RuntimeError: + pass + api.lamport_push_tensor(allreduce, target, x) + out = api.finish_lamport_push(allreduce, target) + bad += differ(ref, out) + kinds["plain"] += 1 + + target = api.lamport_push_target(allreduce, m, h, deferred=True) + assert target.deferred + api.lamport_push_tensor(allreduce, target, x) + pending = api.finish_lamport_push(allreduce, target, deferred=True) + assert isinstance(pending, api.LamportPending) + bad += differ(ref, pending.materialize()) + kinds["deferred"] += 1 + + bad += differ(ref, allreduce(x)) + + # MNNVLAllReduce-level target with a forced transport (the drafter's embedding site). + for one_shot in (True, False): + x = _input(9, 7168, 50 + int(one_shot), rank, 0) + ref = allreduce(x) + target = api.mnnvl_lamport_push_target( + allreduce.mnnvl_allreduce, 9, 7168, one_shot=one_shot + ) + assert target.one_shot == one_shot + api.lamport_push_tensor(allreduce, target, x) + bad += differ(ref, api.finish_lamport_push(allreduce, target)) + + # CUDA graph: plain push, C++ call, deferred push on one workspace, replayed. + m, h = 14, 7168 + xs = [torch.zeros(m, h, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + + def calls(): + t0 = api.lamport_push_target(allreduce, m, h) + api.lamport_push_tensor(allreduce, t0, xs[0]) + a = api.finish_lamport_push(allreduce, t0) + b = allreduce(xs[1]) + t2 = api.lamport_push_target(allreduce, m, h, deferred=True) + api.lamport_push_tensor(allreduce, t2, xs[2]) + c = api.finish_lamport_push(allreduce, t2, deferred=True).materialize() + return a, b, c + + calls() + torch.cuda.synchronize() + comm.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = calls() + for rep in range(3): + for i, x in enumerate(xs): + x.copy_(_input(m, h, 600 + 10 * rep + i, rank, rep % 2)) + torch.cuda.synchronize() + comm.Barrier() + graph.replay() + torch.cuda.synchronize() + for x, out in zip(xs, outs): + bad += differ(allreduce(x), out) + del graph + torch.cuda.synchronize() + assert comm.allreduce(bad) == 0, f"rank {rank}: {bad} values differ from the C++ path" + assert all(count == 10 for count in kinds.values()), kinds + except Exception: + traceback.print_exc() + raise + return True + + +def _similar(actual, expected): + """The family attn_res tests' metric: cosine > 0.9999 and relative L2 < 5e-3.""" + a, e = actual.float().flatten(), expected.float().flatten() + cosine = torch.nn.functional.cosine_similarity(a, e, dim=0).item() + rel_l2 = ((a - e).norm() / (e.norm() + 1e-12)).item() + return cosine > 0.9999 and rel_l2 < 5e-3, f"cosine {cosine:.6f}, relative L2 {rel_l2:.3e}" + + +def _run_rank_cells(world_size): + """Every cell of the four collective entries on the group's shared MNNVL workspace; returns + this rank's failures.""" + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + import importlib + import zlib + + from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch + from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy + from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace + from tensorrt_llm.mapping import Mapping + + sys.path.insert(0, _MV2_TESTS) + from _validating import validating + + lamport = importlib.import_module(f"{_K3}.kernels.k3_lamport_ar") + norm_ops = importlib.import_module(f"{_K3}.ops.norm.k3_attn_res_norm") + push, finish, attn_res, latent = ( + getattr(importlib.import_module(f"{_K3}.{name}"), name.rpartition(".")[2]) + for name in _ENTRIES + ) + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + torch.cuda.set_device(rank) + W = world_size + arch = {(10, 0): Arch.SM_100, (10, 3): Arch.SM_103, (10, 7): Arch.SM_107}.get( + torch.cuda.get_device_capability() + ) + failures = [] + + def check(entry, cell, ok, what): + if not ok: + failures.append(f"rank {rank} {type(entry).__name__} [{cell.why}]: {what}") + + def bits(t): + return t.contiguous().view(torch.int16) + + def same(a, b): + return torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)) + + def seed(cell, salt=0): + return (zlib.crc32(cell.why.encode()) + salt) % 100003 + + def sync(): + torch.cuda.synchronize() + comm.Barrier() + + def rnd(gen, *shape, scale=1.0, shift=0.0): + x = torch.randn(*shape, generator=gen, device="cuda") * scale + shift + return x.to(torch.bfloat16) + + try: + mapping = Mapping(world_size=W, tp_size=W, rank=rank) + AllReduce(mapping=mapping, strategy=AllReduceStrategy.MNNVL, dtype=torch.bfloat16) + # The largest call of the cells: 96 one-shot rows of 7168. + need = lamport.required_bytes(96, 7168, W, True) + workspace = get_or_scale_allreduce_mnnvl_workspace( + mapping, torch.bfloat16, buffer_size_bytes=need + ) + ws = lamport.LamportWorkspace.from_mnnvl(workspace, W, rank) + slots = workspace["uc_buffer"].view(torch.int16) + + if arch in push.ARCHS: + for cell in push.CELLS: + m, h, one_shot = cell.spec["rows"], cell.spec["hidden"], cell.spec["one_shot"] + xs = [_input(m, h, seed(cell), r, 0) for r in range(W)] + sync() + flags = ws.flags.view(torch.int32) + base = int(flags[0].item()) * int(flags[2].item()) // 2 + with validating(push): + push(ws, xs[rank], one_shot) + sync() + want = torch.stack([bits(push.reference(ws, x, one_shot)) for x in xs]) + if one_shot: + got = slots[base : base + m * W * h].view(m, W, h).transpose(0, 1) + else: + owned = torch.arange(rank, m, W, device="cuda") + rows = -(-m // W) + got = slots[base : base + rows * W * h].view(rows, W, h)[owned // W] + got, want = got.transpose(0, 1), want[:, owned] + check(push, cell, torch.equal(got, want), "slot contents") + out = finish(ws, torch.empty_like(xs[rank]), one_shot) + check(push, cell, same(out, finish.reference(xs, out, one_shot)), "sum") + + if arch in finish.ARCHS: + for cell in finish.CELLS: + spec = cell.spec + m, h, one_shot = spec["rows"], spec["hidden"], spec["one_shot"] + xs = [_input(m, h, seed(cell), r, 0) for r in range(W)] + sync() + push(ws, xs[rank], one_shot) + out = torch.empty_like(xs[rank]) + with validating(finish): + got = finish(ws, out, one_shot, cpp_one_shot=spec["cpp_one_shot"]) + want = finish.reference(xs, out, one_shot, cpp_one_shot=spec["cpp_one_shot"]) + check(finish, cell, got.data_ptr() == out.data_ptr(), "returned another tensor") + check(finish, cell, same(got, want), "sum") + + if arch in attn_res.ARCHS: + for cell in attn_res.CELLS: + spec = cell.spec + m, n, h = spec["rows"], spec["candidates"], spec["hidden"] + gen = torch.Generator(device="cuda").manual_seed(seed(cell, 1)) + prefix = rnd(gen, m, h, scale=0.5) + block = rnd(gen, n - 1, m, h, scale=0.5) + weights = ( + rnd(gen, h, scale=0.02), + rnd(gen, h, scale=0.02, shift=1.0), + rnd(gen, h, scale=0.02, shift=1.0), + ) + xs = [_input(m, h, seed(cell), r, 0) for r in range(W)] + two_shot = spec["two_shot"] + chunk = lamport.reduce_chunk(W, spec["cpp_one_shot"]) + tail = (W, chunk, two_shot, ws.mc_base, rank, spec["add_rn"]) + tail += (spec["early_poll"], spec["entry_trigger"], False) + eps = (spec["rms_eps"], spec["out_eps"]) + mxfp8 = spec["mxfp8"] + updated = torch.empty_like(prefix) if spec["into"] else None + sync() + push(ws, xs[rank], not two_shot) + ws_args = (ws.uc_local, ws.flags.data_ptr()) + with validating(attn_res): + got = attn_res( + prefix, *ws_args, block, *weights, *eps, *tail, mxfp8=mxfp8, updated=updated + ) + ref = attn_res.reference(prefix, xs, block, *weights, *eps, *tail, mxfp8=mxfp8) + check(attn_res, cell, len(got) == len(ref), "number of outputs") + check(attn_res, cell, same(got[0], ref[0]), "updated prefix") + if updated is not None: + check(attn_res, cell, got[0].data_ptr() == updated.data_ptr(), "not in updated") + ok, detail = _similar(got[1], ref[1]) + check(attn_res, cell, ok, f"output: {detail}") + if mxfp8: + want_q, want_s = norm_ops.mxfp8_rows_reference(got[1]) + check(attn_res, cell, same(got[2], want_q), "E4M3 codes") + check(attn_res, cell, same(got[3], want_s), "UE8M0 scales") + + if arch in latent.ARCHS: + for cell in latent.CELLS: + spec = cell.spec + m, h = spec["rows"], spec["hidden"] + gen = torch.Generator(device="cuda").manual_seed(seed(cell, 1)) + gamma = rnd(gen, h, scale=0.02, shift=1.0) + xs = [_input(m, h, seed(cell), r, 0) for r in range(W)] + out = torch.empty_like(xs[rank]) + mx, sf_blocks = None, 0 + if spec["mxfp8"] is not None: + sf_blocks = 1 if spec["mxfp8"] == "atoms" else 0 + scale_shape = ((h // 128) * 512,) if sf_blocks else (m, h // 32) + mx = ( + torch.empty(m, h, dtype=torch.float8_e4m3fn, device="cuda"), + torch.empty(scale_shape, dtype=torch.uint8, device="cuda"), + ) + kwargs = dict( + uc_addr=ws.uc_local, + mc_addr=ws.mc_base, + flags=ws.flags, + rank=rank, + world_size=W, + mxfp8_out=mx, + sf_blocks=sf_blocks, + early_state=spec["early_state"], + row=spec["row"], + ) + sync() + with validating(latent): + got = latent(xs[rank], gamma, spec["eps"], out, **kwargs) + ref = latent.reference(xs, gamma, spec["eps"], out, **kwargs) + got_out, ref_out = (got, ref) if mx is None else (got[0], ref[0]) + check(latent, cell, got_out.data_ptr() == out.data_ptr(), "not in out") + ok, detail = _similar(got_out, ref_out) + check(latent, cell, ok, f"output: {detail}") + if mx is not None: + want_q, want_s = norm_ops.mxfp8_rows_reference(got_out) + check(latent, cell, same(mx[0], want_q), "E4M3 codes") + if sf_blocks: + offsets = norm_ops.mx_atom_offsets(m, h, out.device) + check(latent, cell, same(mx[1][offsets], want_s.reshape(-1)), "atoms") + else: + check(latent, cell, same(mx[1], want_s), "UE8M0 scales") + sync() + except Exception: + traceback.print_exc() + raise + return failures + + +@skip_pre_blackwell +@pytest.mark.parametrize("mpi_pool_executor", [2, 4], indirect=True) +def test_k3_lamport_entry_cells(mpi_pool_executor): + """The CELLS of k3_lamport_push, k3_lamport_finish, k3_attn_res_lamport and + k3_latent_ar_rmsnorm (see the module docstring).""" + import importlib + + from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch + + world_size = mpi_pool_executor.num_workers + if torch.cuda.device_count() < world_size: + pytest.skip(f"needs {world_size} GPUs") + arch = {(10, 0): Arch.SM_100, (10, 3): Arch.SM_103, (10, 7): Arch.SM_107}.get( + torch.cuda.get_device_capability() + ) + ops = [ + getattr(importlib.import_module(f"{_K3}.{name}"), name.rpartition(".")[2]) + for name in _ENTRIES + ] + if not any(arch in op.ARCHS for op in ops): + pytest.skip(f"{arch} is in none of the entries' ARCHS") + results = mpi_pool_executor.map(_run_rank_cells, *zip(*[(world_size,)] * world_size)) + failures = [f for r in results for f in r] + assert not failures, "\n".join(failures[:40]) + + +@skip_pre_blackwell +@pytest.mark.parametrize("mpi_pool_executor", [2, 4], indirect=True) +def test_k3_lamport_ar(mpi_pool_executor): + world_size = mpi_pool_executor.num_workers + if torch.cuda.device_count() < world_size: + pytest.skip(f"needs {world_size} GPUs") + results = mpi_pool_executor.map(_run_rank, *zip(*[(world_size,)] * world_size)) + for r in results: + assert r is True + + +@skip_pre_blackwell +@pytest.mark.parametrize("mpi_pool_executor", [2, 4], indirect=True) +def test_k3_lamport_push_api(mpi_pool_executor): + world_size = mpi_pool_executor.num_workers + if torch.cuda.device_count() < world_size: + pytest.skip(f"needs {world_size} GPUs") + results = mpi_pool_executor.map(_run_rank_api, *zip(*[(world_size,)] * world_size)) + for r in results: + assert r is True diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/conftest.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/conftest.py new file mode 100644 index 000000000000..9b9442866249 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/conftest.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import dataclasses + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + features, +) + + +@pytest.fixture +def k3_features(monkeypatch): + """``k3_features(=, ...)`` sets K3Features fields for one test. + + Only fields read when called take effect this way; a field read at import is a module + constant of its reader, which the test patches instead. + """ + + def set_fields(**fields): + monkeypatch.setattr(features, "_ACTIVE", dataclasses.replace(features.active(), **fields)) + return features.active() + + return set_fields diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_allreduce.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_allreduce.py new file mode 100644 index 000000000000..5abcd70032aa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_allreduce.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``layers`` all-reduce, gated MLP (dense and MoE shared expert), row-parallel BF16 +projection and vocab-parallel embedding over 2 / 4 ranks against main's ``AllReduce`` / +``GatedMLP`` / ``Linear`` / ``Embedding``: bitwise, on every rank. + +``TLLM_TEST_MNNVL=1`` makes both sides take MNNVL on one node, as AUTO does across nodes; the +pushed all-reduce of ``layers.py`` runs on the primitive's handle against main's result.""" + +import os +import pickle +import sys +import traceback + +import cloudpickle +import pytest +import torch +from mpi4py import MPI +from utils.util import skip_pre_blackwell + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL) + +pytestmark = pytest.mark.threadleak(enabled=False) + +_K3 = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) + + +def _same(a, b) -> int: + """Number of differing elements (bit patterns), 0 when equal.""" + if isinstance(a, (tuple, list)): + return sum(_same(x, y) for x, y in zip(a, b)) + abs(len(a) - len(b)) + if a.shape != b.shape or a.dtype != b.dtype: + return max(a.numel(), b.numel()) + if a.dtype.is_floating_point and a.element_size() == 2: + return int((a.view(torch.int16) != b.view(torch.int16)).sum()) + return int((a != b).sum()) + + +def _run_rank(world_size, mnnvl): + if mnnvl: + os.environ["TLLM_TEST_MNNVL"] = "1" + else: + os.environ.pop("TLLM_TEST_MNNVL", None) + import importlib + + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.modules.embedding import Embedding + from tensorrt_llm._torch.modules.gated_mlp import GatedMLP + from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode + from tensorrt_llm._torch.modules.situ import SituAndMul + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + core_ar = importlib.import_module(f"{_K3}.layers") + core_mlp = importlib.import_module(f"{_K3}.layers") + core_emb = importlib.import_module(f"{_K3}.layers") + core_linear = importlib.import_module(f"{_K3}.layers") + rank = MPI.COMM_WORLD.Get_rank() + torch.cuda.set_device(rank) + bad = {} + try: + mapping = Mapping(world_size=world_size, tp_size=world_size, rank=rank) + strategies = ( + [AllReduceStrategy.AUTO] + if mnnvl + else [AllReduceStrategy.NCCL, AllReduceStrategy.ONESHOT] + ) + for strategy in strategies: + ref = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + prim = core_ar.TpAllReduce(mapping, torch.bfloat16, strategy) + assert (ref.mnnvl_allreduce is not None) == (prim.transport == "mnnvl") == mnnvl + gen = torch.Generator(device="cuda").manual_seed(rank) + # One-shot and two-shot MNNVL sizes (threshold 1 MiB of tokens * hidden * ranks * 2); + # 3584: the routed MoE latent (RMS_NORM fusion). + shapes = ((1, 7168), (8, 7168), (32, 7168), (64, 7168), (300, 7168), (5, 512)) + for m, h in (*shapes, (16, 3584), (200, 3584)): + x = (torch.randn(m, h, device="cuda", generator=gen) * 0.5).to(torch.bfloat16) + key = f"{strategy.name}/none/{m}x{h}" + bad[key] = _same(prim(x), ref(x)) + if not mnnvl: + continue + w = (torch.randn(h, device="cuda", generator=gen)).to(torch.bfloat16) + r = (torch.randn(m, h, device="cuda", generator=gen)).to(torch.bfloat16) + for op in (AllReduceFusionOp.RESIDUAL_RMS_NORM, AllReduceFusionOp.RMS_NORM): + params = AllReduceParams( + fusion_op=op, + residual=r if op == AllReduceFusionOp.RESIDUAL_RMS_NORM else None, + norm_weight=w, + eps=1e-6, + ) + bad[f"{key}/{op.name}"] = _same( + prim(x, all_reduce_params=params), ref(x, all_reduce_params=params) + ) + + if mnnvl: + # FP32, as the helix vocab-parallel pair exchange runs main's MNNVLAllReduce directly. + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + ref = MNNVLAllReduce(mapping, torch.float32) + prim = core_ar.TpAllReduce(mapping, torch.float32) + assert prim.transport == "mnnvl" + for m, h in ((4, 64), (8, 10240), (100, 7168)): + x = torch.randn(m, h, device="cuda") + want = ref(x, all_reduce_params=AllReduceParams()) + bad[f"fp32/none/{m}x{h}"] = _same(prim(x), want) + + # The pushed all-reduce (layers.py) on the primitive's handle. + lamport = importlib.import_module(f"{_K3}.layers") + ref = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + prim = core_ar.TpAllReduce(mapping, torch.bfloat16) + for m in (1, 8, 100): + x = torch.randn(m, 7168, device="cuda").to(torch.bfloat16) + want = ref(x) + target = lamport.lamport_push_target(prim, m, 7168) + assert target is not None + lamport.lamport_push_tensor(prim, target, x) + bad[f"lamport/{m}"] = _same(lamport.finish_lamport_push(prim, target), want) + + # The SiTU MLP: the MoE shared expert (the model's TP group, is_shared_expert) and the dense + # MLP over the full group and, for 4 ranks, its TP-2 subgroups; reduced, and as the + # partial (enable_allreduce=False). Without MNNVL both sides take NCCL (main autotunes + # AUTO's strategy there, the primitive does not). + strategy = AllReduceStrategy.AUTO if mnnvl else AllReduceStrategy.NCCL + no_reduce = AllReduceParams(enable_allreduce=False) + torch.manual_seed(0) + hidden, inter = 2048, 3072 + gate = (torch.randn(inter, hidden, device="cuda") * 0.02).to(torch.bfloat16) + up = (torch.randn(inter, hidden, device="cuda") * 0.02).to(torch.bfloat16) + down = (torch.randn(hidden, inter, device="cuda") * 0.02).to(torch.bfloat16) + for sub in (None, *sorted({world_size, 2})): + with torch.device("cuda"): + ref = GatedMLP( + hidden_size=hidden, + intermediate_size=inter, + bias=False, + activation=SituAndMul(beta=1.0, linear_beta=8.0, use_fused_activation=True), + dtype=torch.bfloat16, + config=ModelConfig(mapping=mapping, allreduce_strategy=strategy), + overridden_tp_size=sub, + reduce_output=True, + is_shared_expert=sub is None, + ) + ref.gate_up_proj.load_weights([{"weight": gate}, {"weight": up}]) + ref.down_proj.load_weights([{"weight": down}]) + prim = core_mlp.GatedMlp( + hidden, + inter, + activation="situ", + beta=1.0, + linear_beta=8.0, + mapping=mapping if sub is None else core_mlp.subgroup_mapping(mapping, sub), + allreduce_strategy=strategy, + device="cuda", + ) + prim.load(gate, up, down) + tag = "shared" if sub is None else f"tp{sub}" + down_ar, ref_ar = prim.down_proj.all_reduce, ref.down_proj.all_reduce + assert prim.down_proj.reduce_output == ref.down_proj.reduce_output + assert (down_ar.mnnvl_allreduce is None) == (ref_ar.mnnvl_allreduce is None) + assert down_ar.mapping.tp_group == ref_ar.mapping.tp_group + for m in (1, 16, 200): + x = torch.randn(m, hidden, device="cuda").to(torch.bfloat16) + bad[f"mlp/{tag}/{m}"] = _same(prim(x), ref(x)) + bad[f"mlp/{tag}/partial/{m}"] = _same( + prim(x, final_all_reduce_params=no_reduce), + ref(x, final_all_reduce_params=no_reduce), + ) + + # A BF16 row-parallel projection with its all-reduce (the drafter's o_proj): plain, the + # partial, and the residual RMSNorm fusion. + hidden, width = 7168, 1024 * world_size + weight = (torch.randn(hidden, width, device="cuda") * 0.02).to(torch.bfloat16) + with torch.device("cuda"): + ref = Linear( + width, + hidden, + bias=False, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.ROW, + reduce_output=True, + allreduce_strategy=strategy, + ) + ref.load_weights([{"weight": weight}]) + prim = core_linear.Bf16Linear( + width, + hidden, + tp_size=world_size, + tp_rank=rank, + tp_mode=TensorParallelMode.ROW, + all_reduce=core_ar.TpAllReduce(mapping, torch.bfloat16, strategy), + device="cuda", + ) + prim.load(weight) + norm_w = torch.randn(hidden, device="cuda").to(torch.bfloat16) + for m in (1, 8, 64): + x = torch.randn(m, prim.in_features, device="cuda").to(torch.bfloat16) + res = torch.randn(m, hidden, device="cuda").to(torch.bfloat16) + fused = AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=res, + norm_weight=norm_w, + eps=1e-6, + ) + for name, params in (("plain", None), ("partial", no_reduce), ("res_norm", fused)): + bad[f"row/{name}/{m}"] = _same( + prim(x, all_reduce_params=params), ref(x, all_reduce_params=params) + ) + + # Vocab-parallel embedding with its all-reduce. + vocab = 1003 + table = torch.randn(vocab, 512, device="cuda").to(torch.bfloat16) + with torch.device("cuda"): + ref = Embedding( + vocab, + 512, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + allreduce_strategy=strategy, + ) + ref.load_weights([{"weight": table}]) + prim = core_emb.EmbeddingTable( + vocab, + 512, + dtype=torch.bfloat16, + mapping=mapping, + vocab_parallel=True, + allreduce_strategy=strategy, + device="cuda", + ) + prim.load(table) + ids = torch.tensor([0, 5, vocab - 1, -1, 400, 251, 252, 753], device="cuda") + bad["embedding"] = _same(prim(ids), ref(ids)) + except Exception: + traceback.print_exc() + raise + return bad + + +@skip_pre_blackwell +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs") +@pytest.mark.parametrize("mnnvl", [True, False], ids=["mnnvl", "op"]) +@pytest.mark.parametrize("mpi_pool_executor", [2, 4], indirect=True) +def test_core_collectives_match_main(mpi_pool_executor, mnnvl): + world_size = mpi_pool_executor.num_workers + if torch.cuda.device_count() < world_size: + pytest.skip(f"needs {world_size} GPUs") + results = list(mpi_pool_executor.map(_run_rank, *zip(*[(world_size, mnnvl)] * world_size))) + for rank, bad in enumerate(results): + failed = {key: n for key, n in bad.items() if n} + assert not failed, f"rank {rank}: {failed}" + assert bad, f"rank {rank} compared nothing" diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_linear.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_linear.py new file mode 100644 index 000000000000..155b85af6ae8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_linear.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``layers.py`` (linear) against the module chains K3 runs today: bitwise. + +* BF16: main's ``Linear`` (unquantized, F.linear) on the same weight. +* FP8 block scales: main's ``Linear(FP8_BLOCK_SCALES)`` loaded from the checkpoint codes and grid, + swapped into K3's FP8 weight-read projection (``_Fp8BlockScaleWeightReadLinear``, as K3 swaps its + attention projections) and post-loaded; the primitive loads the same full tensors itself. +* Shards: main's ``Linear.load_shard`` for COLUMN / ROW over 1..8 ranks. +""" + +import pytest +import torch +from utils.util import getSMVersion + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as linear, +) +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + +FP8 = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) +COL, ROW = TensorParallelMode.COLUMN, TensorParallelMode.ROW +# Row counts across the low-M / MXFP8 (DeepGEMM) cut-overs and the K3 tuning buckets. +ROWS = (1, 3, 8, 16, 31, 64, 100, 128, 129, 192, 256, 300, 512, 1000) + +requires_block_fp8 = pytest.mark.skipif( + getSMVersion() not in (100, 103, 107) + or (getSMVersion() == 107 and not IS_CUTLASS_DSL_RUBIN_AVAILABLE), + reason="needs a supported block-scaled FP8 GPU and toolchain", +) + + +def _checkpoint_fp8(n, k, seed): + """E4M3 codes and FP32 (non power-of-two) 128x128 block scales of a random ``[n, k]`` weight.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + weight = torch.randn(n, k, device="cuda", generator=gen) * 0.02 + blocks = torch.nn.functional.pad(weight, (0, -k % 128, 0, -n % 128)).view( + -(-n // 128), 128, -(-k // 128), 128 + ) + scale = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-8) / 448.0 + expanded = scale.repeat_interleave(128, 0)[:n].repeat_interleave(128, 1)[:, :k] + return (weight / expanded).to(torch.float8_e4m3fn), scale + + +def _main_linear(n, k, quant_config=None, mapping=None, mode=None, dtype=torch.bfloat16): + with torch.device("cuda"): + return Linear( + k, + n, + bias=False, + dtype=dtype, + mapping=mapping, + tensor_parallel_mode=mode, + quant_config=quant_config, + reduce_output=False, + ) + + +@pytest.mark.parametrize("tp", [1, 2, 3, 4, 8]) +@pytest.mark.parametrize("mode", [COL, ROW]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_shard_matches_main_load_shard(tp, mode, fp8): + n, k = 1536, 2048 + weight = torch.randn(n, k, device="cuda") + scale = torch.randn(n // 128, k // 128, device="cuda") + for rank in range(tp): + mapping = Mapping(world_size=tp, tp_size=tp, rank=rank) + ref = _main_linear(n, k, FP8 if fp8 else None, mapping, mode) + cls = linear.Fp8BlockScaleLinear if fp8 else linear.Bf16Linear + prim = cls(k, n, tp_size=tp, tp_rank=rank, tp_mode=mode, device="cuda") + assert (prim.out_features, prim.in_features) == (ref.out_features, ref.in_features) + assert torch.equal(prim.shard(weight), ref.load_shard(weight, device=weight.device)) + if fp8: + assert torch.equal( + prim.shard(scale, scale_span=128), + ref.load_shard(scale, scale_span=128, device=scale.device), + ) + + +# MLA projection widths that K3 splits over TP ranks: o (ROW over 64 heads x 128), q_b / kv_b +# (COLUMN over 64 heads x 192 / 256). +MLA_SPLITS = [(7168, 8192, ROW), (12288, 1536, COL), (16384, 512, COL)] + + +@pytest.mark.parametrize("tp", [2, 4, 8, 16]) +@pytest.mark.parametrize("n,k,mode", MLA_SPLITS) +@pytest.mark.parametrize("fp8", [False, True]) +def test_shard_matches_main_deferred_weights(n, k, mode, tp, fp8): + """Main's MLA builds its projections with ``skip_create_weights_in_init`` (an even per-tensor + split instead of the block-aligned one); on K3's widths both give the same slices.""" + weight = torch.randn(n, k) + scale = torch.randn(n // 128, k // 128) + for rank in range(tp): + mapping = Mapping(world_size=tp, tp_size=tp, rank=rank) + ref = Linear( + k, + n, + bias=False, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=mode, + quant_config=FP8 if fp8 else None, + reduce_output=False, + skip_create_weights_in_init=True, + ) + assert ref.tp_sharding is None + cls = linear.Fp8BlockScaleLinear if fp8 else linear.Bf16Linear + prim = cls(k, n, tp_size=tp, tp_rank=rank, tp_mode=mode, device="cpu") + assert (prim.out_features, prim.in_features) == (ref.out_features, ref.in_features) + assert torch.equal(prim.shard(weight), ref.load_shard(weight)) + if fp8: + assert torch.equal(prim.shard(scale, scale_span=128), ref.load_shard(scale)) + + +@pytest.mark.parametrize("n,k", [(1536, 7168), (7168, 2048), (24, 7168)]) +@torch.inference_mode() +def test_bf16_linear_matches_main(n, k): + weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.02 + ref = _main_linear(n, k) + ref.load_weights([{"weight": weight}]) + ref.post_load_weights() + prim = linear.Bf16Linear(k, n, device="cuda") + prim.load(weight) + assert torch.equal(prim.weight, ref.weight) + for m in ROWS: + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + assert torch.equal(prim(x), ref(x)), m + + +def _weight_read(codes, scale): + """K3's decode chain: the checkpoint loaded into an unprepared ``Fp8BlockScaleLinear``, swapped to + the weight-read projection, which prepares the codes and scales itself.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.fp8_weight_read import ( # noqa: E501 + _Fp8BlockScaleWeightReadLinear, + ) + + n, k = codes.shape + loader = linear.Fp8BlockScaleLinear(k, n, device="cuda") + loader.load(codes, scale) + proj = _Fp8BlockScaleWeightReadLinear.from_linear(loader) + proj.post_load_weights() + return proj + + +# K3's checkpoint FP8 projection shapes (q_a + kv_a fused, q_b / kv_b / o / g heads, KDA q/k/v/g/o). +FP8_SHAPES = [(2112, 7168), (1152, 1536), (2304, 512), (7168, 1024), (512, 7168), (4096, 7168)] + + +@requires_block_fp8 +@pytest.mark.parametrize("n,k", FP8_SHAPES) +@torch.inference_mode() +def test_fp8_block_linear_matches_k3_weight_read(n, k): + codes, scale = _checkpoint_fp8(n, k, seed=n + k) + ref = _weight_read(codes.clone(), scale.clone()) + prim = linear.Fp8BlockScaleLinear(k, n, device="cuda") + prim.load(codes, scale) + prim.prepare() + assert torch.equal(prim.weight.view(torch.uint8), ref.weight.view(torch.uint8)) + assert torch.equal(prim.weight_scale, ref.weight_scale) + assert prim.low_m_max == ref._low_m_max + if prim.low_m_scale is not None: + assert torch.equal(prim.low_m_scale, ref._low_m_weight_scale) + for m in ROWS: + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + assert torch.equal(prim(x), ref(x)), m + + +@requires_block_fp8 +@torch.inference_mode() +def test_fp8_block_linear_registers_main_parameter_names(): + n, k = 1152, 1536 + ref = _main_linear(n, k, FP8, mode=COL) + prim = linear.Fp8BlockScaleLinear(k, n, device="cuda") + ref_params = {name: (p.shape, p.dtype) for name, p in ref.named_parameters()} + prim_params = {name: (p.shape, p.dtype) for name, p in prim.named_parameters()} + assert prim_params == ref_params diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_mxfp8_gemm.py new file mode 100644 index 000000000000..45a1ce8edfef --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_mxfp8_gemm.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The K3 MXFP8 GEMM (``ops/gemm/k3_mxfp8_gemm.py`` on ``ops/gemm/_mxfp8_gemm.py``) on SM107. + +* Its runner's candidate tactics are main's ``CuteDSLKimiK3MXFP8RubinLinear`` list with K3's rules + applied (checked here rule by rule, independently of ``keeps_tactic``), in main's order, over the + K3 FP8 projection shapes and every K3 tuning bucket. +* With one tactic, its launch is bitwise main's runner with that tactic. +* The op asks the autotuner under its own name with that runner, and untuned it is bitwise main's + ``cute_dsl_mxfp8_gemm_rubin(fine_grained_m=True)``. +* ``k3_mxfp8_gemm.CELLS`` under ``validating`` against the entry's FP64 ``reference``. +""" + +import pytest +import torch +from utils.util import getSMVersion + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 + +pytestmark = pytest.mark.skipif( + getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="the SM107 CuTe DSL MXFP8 GEMM", +) + +K3 = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) +# Per-rank [N, K] of the K3 FP8 block-scale projections that reach the MXFP8 GEMM. +SHAPES = ((2112, 7168), (1152, 1536), (2304, 512), (7168, 1024), (512, 7168), (4096, 7168)) +# Shapes and rows whose tactics run against main's runner (each new tactic compiles a kernel). +RUN_SHAPES = ((2112, 7168), (7168, 1024)) +RUN_ROWS = (1, 48, 256, 1000) +RUN_TACTICS = 2 + + +def _mods(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as linear, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.gemm import ( # noqa: E501 + _mxfp8_gemm as mxfp8_gemm, + ) + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as cdo + + return linear, mxfp8_gemm, cdo + + +def _entry(): + import importlib + + return importlib.import_module(f"{K3}.ops.gemm.k3_mxfp8_gemm").k3_mxfp8_gemm + + +def _tactic_inputs(m, n, k): + """Shape-only operands for ``get_valid_tactics``.""" + a = torch.empty(m, k, dtype=torch.float8_e4m3fn, device="cuda") + b = torch.empty(n, k, dtype=torch.float8_e4m3fn, device="cuda") + return [a, b, None, None, None] + + +def _k3_rules_hold(tactic, m): + if tactic[0] == "base": + _, tiler, inst, _cluster, swap_ab, use_prefetch, _sched, _raster, split_k = tactic + else: + _, tiler, inst, _pref, _fallback, swap_ab, use_prefetch, *_ = tactic + split_k = 1 + mma_m = (tiler[0], inst[0]) + return ( + mma_m != (256, 128) + and (not use_prefetch or (not swap_ab and mma_m == (512, 256))) + and (split_k == 1 or m <= 256) + ) + + +def test_tactics_are_main_list_under_k3_rules(): + _, mxfp8_gemm, cdo = _mods() + main_runner = cdo.CuteDSLKimiK3MXFP8RubinLinear(output_dtype=torch.bfloat16) + k3_runner = mxfp8_gemm.K3MXFP8RubinRunner(output_dtype=torch.bfloat16) + buckets = cdo._get_kimi_k3_mxfp8_tuning_buckets(8192) + dropped = 0 + for n, k in SHAPES: + for m in buckets: + inputs = _tactic_inputs(m, n, k) + main = main_runner.get_valid_tactics(inputs, None) + got = k3_runner.get_valid_tactics(inputs, None) + want = [t for t in main if _k3_rules_hold(t, m)] + assert got == want, (n, k, m) + assert got, (n, k, m) + assert all(len(t) == 9 for t in got if t[0] == "base"), (n, k, m) + dropped += len(main) - len(got) + assert dropped + + +def _operands(m, n, k, seed=0): + linear, _, _ = _mods() + gen = torch.Generator(device="cuda").manual_seed(seed) + w = torch.randn(n, k, device="cuda", generator=gen) * 0.02 + blocks = torch.nn.functional.pad(w, (0, -k % 128, 0, -n % 128)).view( + -(-n // 128), 128, -(-k // 128), 128 + ) + scale = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-8) / 448.0 + expanded = scale.repeat_interleave(128, 0)[:n].repeat_interleave(128, 1)[:, :k] + weight, weight_scale = linear.prepare_fp8_block_weight( + (w / expanded).to(torch.float8_e4m3fn), scale + ) + x = torch.randn(m, k, device="cuda", generator=gen).to(torch.bfloat16) + act, act_scale = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(x) + return act, weight, act_scale, weight_scale + + +@torch.inference_mode() +def test_each_tactic_matches_main_runner(): + _, mxfp8_gemm, cdo = _mods() + main_runner = cdo.CuteDSLKimiK3MXFP8RubinLinear(output_dtype=torch.bfloat16, use_tvm_ffi=True) + for n, k in RUN_SHAPES: + for m in RUN_ROWS: + act, weight, act_scale, weight_scale = _operands(m, n, k) + alpha = cdo._get_mxfp8_gemm_alpha(act.device) + inputs = [act, weight, act_scale, weight_scale, alpha] + bucket = cdo._kimi_k3_mxfp8_tuning_bucket(m) + tactics = mxfp8_gemm.K3MXFP8RubinRunner(output_dtype=torch.bfloat16).get_valid_tactics( + _tactic_inputs(bucket, n, k), None + ) + # Split-K accumulates with BF16 atomics: not run-to-run deterministic on either side. + serial = [t for t in tactics if t[0] != "base" or t[8] == 1] + step = max(1, len(serial) // RUN_TACTICS) + for tactic in [None, *serial[::step][:RUN_TACTICS]]: + got = mxfp8_gemm.run_with_tactic(act, weight, act_scale, weight_scale, tactic) + want = main_runner(inputs, tactic=-1 if tactic is None else tactic) + assert torch.equal(got, want), (n, k, m, tactic) + + +@torch.inference_mode() +def test_op_tunes_k3_runner_and_matches_main_untuned(monkeypatch): + _, mxfp8_gemm, _ = _mods() + k3_mxfp8_gemm = _entry() + from tensorrt_llm._torch.autotuner import AutoTuner + + for n, k in RUN_SHAPES: + for m in RUN_ROWS: + act, weight, act_scale, weight_scale = _operands(m, n, k, seed=1) + got = k3_mxfp8_gemm(act, weight, act_scale, weight_scale) + want = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + act, weight, act_scale, weight_scale, fine_grained_m=True + ) + assert torch.equal(got, want), (n, k, m) + + seen = [] + choose_one = AutoTuner.choose_one + + def spy(self, custom_op, runners, tuning_config, inputs, **kwargs): + seen.append((custom_op, [type(r) for r in runners], tuning_config)) + return choose_one(self, custom_op, runners, tuning_config, inputs, **kwargs) + + monkeypatch.setattr(AutoTuner, "choose_one", spy) + k3_mxfp8_gemm(*_operands(16, *RUN_SHAPES[0])) + assert seen == [ + ( + mxfp8_gemm.OP_NAME, + [mxfp8_gemm.K3MXFP8RubinRunner], + mxfp8_gemm.K3MXFP8RubinRunner.tuning_config, + ) + ] + + +@pytest.mark.parametrize("cell", _entry().CELLS, ids=lambda c: c.why[:48]) +@torch.inference_mode() +def test_k3_mxfp8_gemm_cells(cell): + """Both operands are E4M3 with power-of-two scales, so every product is exact; the kernel + sums them in FP32 (split-K tactics, up to M = 256, add their BF16-rounded partials with BF16 + atomics). Serial tactics: one BF16 rounding, at most 2 units of ``ULP`` of the row maximum. + Split-K: one rounding per partial plus one per atomic add; 8 units element-wise and 2 in RMS + allow a few of them (not yet measured on SM107).""" + k3_mxfp8_gemm = _entry() + n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"] + act, weight, act_scale, weight_scale = _operands(m, n, k, seed=2) + with validating(k3_mxfp8_gemm): + got = k3_mxfp8_gemm(act, weight, act_scale, weight_scale) + want = k3_mxfp8_gemm.reference(act, weight, act_scale, weight_scale) + if m <= 256: + assert_within_ulp(got, want, 8, 2) + else: + assert_within_ulp(got, want, 2, 1) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_norm_mlp_embedding.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_norm_mlp_embedding.py new file mode 100644 index 000000000000..1a991d5ee851 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/core/test_core_norm_mlp_embedding.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``layers`` norm, gated MLP, embedding, conv weight and two-stream fork against main's +modules (one GPU, one rank): bitwise.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as conv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as embedding, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as mlp, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as streams, +) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.modules.mamba.layernorm_gated import RMSNorm as GatedRMSNorm +from tensorrt_llm._torch.modules.multi_stream_utils import ( + maybe_execute_in_parallel, + with_multi_stream, +) +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm.mapping import Mapping + +ROWS = (1, 4, 16, 33, 128, 1000) + + +def _bf16(*shape, scale=1.0): + return (torch.randn(*shape, device="cuda") * scale).to(torch.bfloat16) + + +@pytest.mark.parametrize("hidden", [512, 1536, 7168]) +@torch.inference_mode() +def test_rms_norm_matches_main(hidden): + eps = 1e-6 + ref = RMSNorm(hidden_size=hidden, eps=eps, dtype=torch.bfloat16).cuda() + ref.weight.data.copy_(_bf16(hidden)) + prim = norm.RMSNormWeight(hidden, eps, dtype=torch.bfloat16, device="cuda") + prim.weight.data.copy_(ref.weight) + assert dict(prim.named_parameters()).keys() == dict(ref.named_parameters()).keys() + assert prim.variance_epsilon == ref.variance_epsilon == prim.eps + for m in ROWS: + x = _bf16(m, hidden, scale=3.0) + assert torch.equal(prim(x), ref(x)), m + r = _bf16(m, hidden) + x_ref, r_ref = ref(x.clone(), r.clone()) + x_prim, r_prim = prim(x.clone(), r.clone()) + assert torch.equal(x_prim, x_ref) and torch.equal(r_prim, r_ref), m + + +def test_gated_norm_and_conv_weights_register_main_names(): + gated = GatedRMSNorm(128, eps=1e-5) + prim = norm.RMSNormWeight(128, 1e-5) + assert dict(prim.named_parameters()).keys() == dict(gated.named_parameters()).keys() + assert prim.eps == gated.eps + from fla.modules import ShortConvolution + + ref_conv = ShortConvolution(hidden_size=1536, kernel_size=4, activation="silu") + prim_conv = conv.ShortConvWeight(1536, 4) + ref_params = {n: p.shape for n, p in ref_conv.named_parameters()} + assert {n: p.shape for n, p in prim_conv.named_parameters()} == ref_params + assert prim_conv.activation == ref_conv.activation + + +def _main_mlp(hidden, inter, activation, mapping=None): + config = ModelConfig(mapping=mapping or Mapping()) + with torch.device("cuda"): + return GatedMLP( + hidden_size=hidden, + intermediate_size=inter, + bias=False, + activation=activation, + dtype=torch.bfloat16, + config=config, + reduce_output=False, + ) + + +@pytest.mark.parametrize("act", ["situ", "situ_linear", "silu"]) +@torch.inference_mode() +def test_gated_mlp_matches_main(act): + hidden, inter = 2048, 1536 + if act == "silu": + activation, kwargs = torch.nn.functional.silu, {} + else: + linear_beta = 3.0 if act == "situ_linear" else None + activation = SituAndMul(beta=2.0, linear_beta=linear_beta, use_fused_activation=True) + kwargs = dict(beta=2.0, linear_beta=linear_beta) + ref = _main_mlp(hidden, inter, activation) + gate, up, down = ( + _bf16(inter, hidden, scale=0.02), + _bf16(inter, hidden, scale=0.02), + _bf16(hidden, inter, scale=0.02), + ) + ref.gate_up_proj.load_weights([{"weight": gate}, {"weight": up}]) + ref.down_proj.load_weights([{"weight": down}]) + prim = mlp.GatedMlp( + hidden, inter, activation="silu" if act == "silu" else "situ", device="cuda", **kwargs + ) + prim.load(gate, up, down) + assert torch.equal(prim.gate_up_proj.weight, ref.gate_up_proj.weight) + assert torch.equal(prim.down_proj.weight, ref.down_proj.weight) + for m in ROWS: + x = _bf16(m, hidden) + assert torch.equal(prim(x), ref(x)), m + + +@pytest.mark.parametrize("tp", [2, 4, 8]) +def test_gated_mlp_shards_match_main(tp): + hidden, inter = 1024, 3000 + gate, up, down = ( + torch.randn(inter, hidden), + torch.randn(inter, hidden), + torch.randn(hidden, inter), + ) + for rank in range(tp): + mapping = Mapping(world_size=tp, tp_size=tp, rank=rank) + ref = _main_mlp(hidden, inter, torch.nn.functional.silu, mapping) + ref.gate_up_proj.load_weights([{"weight": gate}, {"weight": up}]) + ref.down_proj.load_weights([{"weight": down}]) + prim = mlp.GatedMlp( + hidden, inter, activation="silu", mapping=mapping, reduce_output=False, device="cuda" + ) + prim.load(gate, up, down) + assert torch.equal(prim.gate_up_proj.weight, ref.gate_up_proj.weight) + assert torch.equal(prim.down_proj.weight, ref.down_proj.weight) + + +def test_subgroup_mapping_matches_main(): + for world, sub in ((8, 4), (16, 4), (16, 2)): + for rank in range(world): + mapping = Mapping(world_size=world, tp_size=world, rank=rank, gpus_per_node=4) + # GatedMLP keeps its subgroup mapping on the projections. + config = ModelConfig(mapping=mapping) + with torch.device("cuda"): + ref = GatedMLP( + hidden_size=64, + intermediate_size=64, + bias=False, + dtype=torch.bfloat16, + config=config, + overridden_tp_size=sub, + reduce_output=False, + ) + got = mlp.subgroup_mapping(mapping, sub) + want = ref.down_proj.mapping + assert ( + got.tp_size, + got.tp_rank, + got.tp_group, + got.world_size, + got.is_multi_node(), + ) == ( + want.tp_size, + want.tp_rank, + want.tp_group, + want.world_size, + want.is_multi_node(), + ) + + +@pytest.mark.parametrize("vocab", [1000, 163840]) +@pytest.mark.parametrize("tp", [1, 3, 4]) +@torch.inference_mode() +def test_embedding_matches_main(vocab, tp): + """Main's vocab-parallel ``Embedding`` (COLUMN, gather_output) before its all-reduce, or its + whole forward with one rank, on ids that include -1 and both ends of the vocabulary.""" + from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + hidden = 256 + table = _bf16(vocab, hidden) + ids = torch.cat( + [ + torch.randint(0, vocab, (37,), device="cuda"), + torch.tensor([-1, vocab - 1, 0], device="cuda"), + ] + ) + for rank in range(tp): + mapping = Mapping(world_size=tp, tp_size=tp, rank=rank) + with torch.device("cuda"): + ref = Embedding( + vocab, + hidden, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + reduce_output=False, + ) + ref.load_weights([{"weight": table}]) + prim = embedding.EmbeddingTable( + vocab, + hidden, + dtype=torch.bfloat16, + mapping=mapping, + vocab_parallel=True, + reduce_output=False, + device="cuda", + ) + prim.load(table) + assert torch.equal(prim.weight, ref.weight) + assert (prim.vocab_start, prim.vocab_end) == (ref.vocab_start_index, ref.vocab_end_index) + if tp == 1: + want = ref(ids) + assert torch.equal(embedding.gather_embedding(ids, prim.weight), want) + else: + want = pre_comm_embedding_ops( + ids, + ref.weight, + tp, + rank, + TensorParallelMode.COLUMN, + ref.vocab_start_index, + ref.vocab_end_index, + True, + ref.padding_size, + ) + assert torch.equal(prim(ids), want) + + +def test_fork_join_matches_maybe_execute_in_parallel(): + aux = torch.cuda.Stream() + x = torch.randn(64, 64, device="cuda") + seen = {} + + def fn0(): + seen["fn0"] = torch.cuda.current_stream() + return x @ x + + def fn1(): + seen["fn1"] = torch.cuda.current_stream() + return (x + 1, [x * 2]) + + for multi in (False, True): + with with_multi_stream(multi): + ref0, ref1 = maybe_execute_in_parallel( + fn0, fn1, torch.cuda.Event(), torch.cuda.Event(), aux + ) + ref_streams = dict(seen) + got0, got1 = streams.fork_join(fn0, fn1, torch.cuda.Event(), torch.cuda.Event(), aux) + assert seen == ref_streams + assert (seen["fn1"] == aux) == multi + torch.cuda.synchronize() + assert torch.equal(got0, ref0) and torch.equal(got1[0], ref1[0]) + assert torch.equal(got1[1][0], ref1[1][0]) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_k3_fp8_fused_projection_resmooth.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_k3_fp8_fused_projection_resmooth.py new file mode 100644 index 000000000000..858ce63c4a9f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_k3_fp8_fused_projection_resmooth.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The KDA q/k/v/g FP8 views and their fused qkvg projection share codes and scale grids. + +Resmoothing to UE8M0 is in place and not idempotent, so the shared storage must be +resmoothed exactly once (by the fused owner) whichever module is prepared first. +""" + +import weakref + +import pytest +import torch +from utils.util import getSMVersion + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.fp8_weight_read import ( # noqa: E501 + _Fp8BlockScaleWeightReadLinear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kda import ( # noqa: E501 + K3KDALinearAttention, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.layers import ( + Fp8BlockScaleLinear, +) +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103, 107) + or (getSMVersion() == 107 and not IS_CUTLASS_DSL_RUBIN_AVAILABLE), + reason="needs a supported block-scaled FP8 GPU and toolchain", +) +@pytest.mark.parametrize("owner_first", [False, True]) +def test_fused_checkpoint_projections_preserve_outputs(owner_first: bool) -> None: + """Shared checkpoint storage must keep FP8 codes paired with their scales.""" + sources = tuple(Fp8BlockScaleLinear(512, 128, device="cuda") for _ in range(4)) + for source in sources: + source.load( + torch.full((128, 512), 448, dtype=torch.float8_e4m3fn, device="cuda"), + torch.full((1, 4), 0.51, device="cuda"), + ) + + fused_source = K3KDALinearAttention._fuse_checkpoint_projections(sources) + K3KDALinearAttention._alias_projections_into_fused(sources, fused_source) + fused = _Fp8BlockScaleWeightReadLinear.from_linear(fused_source) + projections = tuple(_Fp8BlockScaleWeightReadLinear.from_linear(source) for source in sources) + for index, projection in enumerate(projections): + projection._fused_projection = (weakref.ref(fused), index * 128) + modules = (fused, *projections) if owner_first else (*projections, fused) + for module in modules: + module.post_load_weights() + + # 448 * 0.51 rounds to FP8 224 at scale 1. Resmoothing that representation + # again can produce 448 at scale 0.5, so stale scales would double outputs. + x = torch.full((8, 512), 1 / 512, dtype=torch.bfloat16, device="cuda") + for module in (fused, *projections): + actual = module(x) + torch.testing.assert_close(actual, torch.full_like(actual, 224), rtol=0, atol=0) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_decode_fused_fab.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_decode_fused_fab.py new file mode 100644 index 000000000000..4c52054e2d59 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_decode_fused_fab.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi-K3 KDA verify with the forget gate computed in the kernel (``kda_mtp_decode`` with ``fa``). + +The fused path runs the ``k3_kda_fab_fa`` entry (the fab kernel's GEMM 1 and reduction, BF16 +``[f_a | b]``) and ``kda_mtp_decode`` with ``fa`` / ``w_fb`` (the f_b MMA inside the verify +kernel), optionally as a programmatic dependent launch. It must reproduce the current +``k3_kda_fab_proj`` + ``kda_mtp_decode`` pair bit for bit: outputs, committed recurrent state, +conv windows and replay caches, over several verify rounds with mixed acceptance, eagerly and +under CUDA graph replay. +""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)), + reason="needs an SM100/SM103/SM107 GPU", +) + +H = 6 +K = 128 +W = 4 +HIDDEN = 7168 +MUTATED = ("state", "cs_q", "cs_k", "cs_v", "k_cache", "g_cache", "v_cache", "beta_cache") + + +def _n1(heads): + """[f_a | b] rows padded to a multiple of 8 (the decode weight's layout).""" + return (K + heads + 7) // 8 * 8 + + +# the default (H heads) layer's [f_a | b] rows and projection width +N1 = _n1(H) +DIM = H * K + + +def _fab_proj(): + """The ``k3_kda_fab_proj`` entry (its import needs CuTe DSL).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.k3_kda_fab_proj import ( # noqa: E501 + k3_kda_fab_proj, + ) + + return k3_kda_fab_proj + + +def _fab_fa(): + """The ``k3_kda_fab_fa`` entry (its import needs CuTe DSL).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.k3_kda_fab_fa import ( # noqa: E501 + k3_kda_fab_fa, + ) + + return k3_kda_fab_fa + + +def _ops(): + """The ``kda_mtp_decode`` entry (its import needs CuTe DSL).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.kda_mtp_decode import ( # noqa: E501 + kda_mtp_decode, + ) + + return kda_mtp_decode + + +def _layer(batch, num_spec, seed, heads=H): + """Weights, projections and a decode state that already went through several rounds.""" + g = torch.Generator(device="cpu").manual_seed(seed) + dev = "cuda" + H, N1, DIM = heads, _n1(heads), heads * K + + def rn(*shape, scale=1.0, dtype=torch.float32): + return (torch.randn(*shape, generator=g) * scale).to(dtype).to(dev) + + pool = max(2 * batch, 8) + m = batch * (num_spec + 1) + s = W - 1 + num_spec + x = torch.randn(m, HIDDEN, generator=g) + chan = torch.ones(HIDDEN) + chan[torch.randperm(HIDDEN, generator=g)[:16]] = 30.0 # activation outlier channels + w_bfa = torch.randn(N1, HIDDEN, generator=g) * 0.02 + w_bfa[K + H :] = 0 + # Per-row state scales spread over two orders of magnitude, like a long-running request. + state = torch.randn(pool, H, K, K, generator=g) * 0.05 + state *= torch.logspace(-1, 1, K).view(1, 1, K, 1) + perm = torch.randperm(pool, generator=g)[:batch].to(torch.int32) + accepted = torch.zeros(pool, dtype=torch.int32) + accepted[perm.long()] = torch.randint(0, num_spec + 1, (batch,), generator=g, dtype=torch.int32) + return { + "heads": heads, + "x": (x * chan).bfloat16().to(dev), + "w_bfa": w_bfa.bfloat16().to(dev), + "w_fb": rn(DIM, K, scale=0.05, dtype=torch.bfloat16), + "qkvg": rn(m, 4 * DIM, scale=0.5, dtype=torch.bfloat16), + "w_q": rn(DIM, W, scale=0.5), + "w_k": rn(DIM, W, scale=0.5), + "w_v": rn(DIM, W, scale=0.5), + "A_log": rn(H, scale=0.5), + "dt_bias": rn(DIM, scale=0.5), + "cs_q": rn(pool, s, DIM, dtype=torch.bfloat16).transpose(1, 2), + "cs_k": rn(pool, s, DIM, dtype=torch.bfloat16).transpose(1, 2), + "cs_v": rn(pool, s, DIM, dtype=torch.bfloat16).transpose(1, 2), + "state": state.to(dev), + "k_cache": rn(pool, num_spec, DIM, scale=0.08, dtype=torch.bfloat16), + "g_cache": -torch.rand(pool, num_spec, DIM, generator=g).to(dev) * 0.5, + "v_cache": rn(pool, num_spec, DIM, dtype=torch.bfloat16), + "beta_cache": torch.rand(pool, num_spec, H, generator=g).to(torch.bfloat16).to(dev), + "onorm_w": (0.75 + 0.5 * torch.rand(K, generator=g)).to(dev), + "slots": perm.to(dev), + "accepted": accepted.to(dev), + } + + +def _clone(lay): + out = dict(lay) + for key in MUTATED: + t = lay[key] + c = torch.empty_strided(t.shape, t.stride(), dtype=t.dtype, device=t.device) + c.copy_(t) + out[key] = c + return out + + +def _decode_kwargs(lay, batch, num_spec, *, g=None, beta=None, bfa=None, use_pdl=False): + m = batch * (num_spec + 1) + H = lay["heads"] + DIM = H * K + qkvg = lay["qkvg"] + kw = dict( + x_q=qkvg[:, :DIM].view(1, m, H, K), + x_k=qkvg[:, DIM : 2 * DIM].view(1, m, H, K), + x_v=qkvg[:, 2 * DIM : 3 * DIM].view(1, m, H, K), + w_q=lay["w_q"], + w_k=lay["w_k"], + w_v=lay["w_v"], + cs_q=lay["cs_q"], + cs_k=lay["cs_k"], + cs_v=lay["cs_v"], + g=None if g is None else g.view(1, m, H, K), + beta=(beta if bfa is None else bfa[:, K : K + H]).view(1, m, H), + A_log=lay["A_log"], + dt_bias=lay["dt_bias"], + recurrent_state=lay["state"], + k_cache=lay["k_cache"], + g_cache=lay["g_cache"], + v_cache=lay["v_cache"], + beta_cache=lay["beta_cache"], + ssm_state_indices=lay["slots"], + cu_seqlens=None, + num_spec=num_spec, + num_accepted_tokens=lay["accepted"], + lower_bound=-5.0, + scale=K**-0.5, + packed_token_layout=True, + onorm_g=qkvg[:, 3 * DIM :].view(1, m, H, K), + onorm_weight=lay["onorm_w"], + onorm_eps=1e-5, + fuse_output_norm=True, + use_pdl=use_pdl, + ) + if bfa is not None: + kw["fa"] = bfa[:, :K].view(1, m, K) + kw["w_fb"] = lay["w_fb"] + return kw + + +def _bits(t): + return t.view({1: torch.uint8, 2: torch.int16, 4: torch.int32}[t.element_size()]) + + +def _assert_same(name, actual, expected): + assert actual.shape == expected.shape, name + assert torch.equal(_bits(actual), _bits(expected)), ( + f"{name}: max |diff| {(actual.float() - expected.float()).abs().max().item():.3e}" + ) + + +def _next_acceptance(lay, batch, num_spec, seed): + g = torch.Generator(device="cpu").manual_seed(seed) + acc = torch.randint(0, num_spec + 1, (batch,), generator=g, dtype=torch.int32) + lay["accepted"][lay["slots"].long()] = acc.to("cuda") + + +def _reference(lay, batch, num_spec): + """The current pair: k3_kda_fab_proj (beta, forget gate) + kda_mtp_decode.""" + beta, gate = _fab_proj()(lay["x"], lay["w_bfa"], lay["w_fb"], lay["heads"]) + return _ops()(**_decode_kwargs(lay, batch, num_spec, g=gate, beta=beta)) + + +def _fused(lay, batch, num_spec, use_pdl): + bfa = _fab_fa()(lay["x"], lay["w_bfa"], lay["heads"], lay["w_fb"]) + return _ops()(**_decode_kwargs(lay, batch, num_spec, bfa=bfa, use_pdl=use_pdl)) + + +@pytest.mark.parametrize("m", [1, 7, 8, 32, 56, 64, 100, 128, 168, 224, 256]) +def test_fab_fa_matches_fab_beta(m): + lay = _layer((m + 7) // 8, 7, seed=m) + x = lay["x"][:m].contiguous() + beta, _ = _fab_proj()(x, lay["w_bfa"], lay["w_fb"], H) + bfa = _fab_fa()(x, lay["w_bfa"], H, lay["w_fb"]) + _assert_same("beta", bfa[:, K : K + H], beta) + # f_a is the fab tail's BF16 rounding of the same FP32 sums: close to the FP32 reference + ref = x.float() @ lay["w_bfa"].float().t() + torch.testing.assert_close(bfa[:, :K].float(), ref[:, :K], rtol=2**-7, atol=1e-3) + # deterministic run to run (CUDA-graph-safe workspace hand-off) + _assert_same("bfa rerun", _fab_fa()(x, lay["w_bfa"], H, lay["w_fb"]), bfa) + + +@pytest.mark.parametrize("m", [8, 16, 32, 64]) +def test_fab_fa_matches_fab_beta_helix8(m): + """12 heads per rank (helix8), the verify rows of 1 / 2 / 4 / 8 requests with seven drafts.""" + lay = _layer(m // 8, 7, seed=m, heads=12) + beta, _ = _fab_proj()(lay["x"], lay["w_bfa"], lay["w_fb"], 12) + bfa = _fab_fa()(lay["x"], lay["w_bfa"], 12, lay["w_fb"]) + _assert_same("beta", bfa[:, K : K + 12], beta) + ref = lay["x"].float() @ lay["w_bfa"].float().t() + torch.testing.assert_close(bfa[:, :K].float(), ref[:, :K], rtol=2**-7, atol=1e-3) + + +@pytest.mark.parametrize("use_pdl", [False, True]) +@pytest.mark.parametrize( + "heads,batch,num_spec", + [(H, 1, 7), (H, 4, 7), (H, 8, 7), (H, 16, 7), (H, 32, 6), (H, 32, 7), (H, 4, 5), (H, 16, 6)] + + [(12, 1, 7), (12, 2, 7), (12, 4, 7), (12, 8, 7)], +) +@torch.no_grad() +def test_fused_gate_verify_bitwise(heads, batch, num_spec, use_pdl): + lay = _layer(batch, num_spec, seed=100 * batch + num_spec, heads=heads) + ref, fused = _clone(lay), _clone(lay) + for rnd in range(3): + out_ref = _reference(ref, batch, num_spec) + out_fused = _fused(fused, batch, num_spec, use_pdl) + torch.cuda.synchronize() + assert torch.isfinite(out_ref.float()).all() + _assert_same(f"round {rnd} out", out_fused, out_ref) + for key in MUTATED: + _assert_same(f"round {rnd} {key}", fused[key], ref[key]) + for state in (ref, fused): + _next_acceptance(state, batch, num_spec, seed=rnd) + + +@pytest.mark.parametrize("batch,num_spec", [(4, 7), (16, 7), (32, 6)]) +@torch.no_grad() +def test_pdl_verify_bitwise(batch, num_spec): + """The PDL launch of the unchanged-gate path is bitwise the plain launch.""" + lay = _layer(batch, num_spec, seed=7 + batch) + beta, gate = _fab_proj()(lay["x"], lay["w_bfa"], lay["w_fb"], H) + plain, pdl = _clone(lay), _clone(lay) + for rnd in range(2): + out_plain = _ops()(**_decode_kwargs(plain, batch, num_spec, g=gate, beta=beta)) + out_pdl = _ops()(**_decode_kwargs(pdl, batch, num_spec, g=gate, beta=beta, use_pdl=True)) + _assert_same(f"round {rnd} out", out_pdl, out_plain) + for key in MUTATED: + _assert_same(f"round {rnd} {key}", pdl[key], plain[key]) + for state in (plain, pdl): + _next_acceptance(state, batch, num_spec, seed=50 + rnd) + + +@torch.no_grad() +def test_fused_gate_uses_w_fb(): + """Guard against a vacuous comparison: a different W_fb must change the result.""" + batch, num_spec = 4, 7 + lay = _layer(batch, num_spec, seed=3) + a, b = _clone(lay), _clone(lay) + b["w_fb"] = (lay["w_fb"].float() * 1.5).bfloat16() + out_a = _fused(a, batch, num_spec, use_pdl=False) + out_b = _fused(b, batch, num_spec, use_pdl=False) + assert not torch.equal(out_a, out_b) + + +@pytest.mark.parametrize("batch,num_spec", [(8, 7), (16, 7), (32, 6)]) +@torch.no_grad() +def test_fused_gate_cuda_graph(batch, num_spec): + """fab (f_a mode) on a side stream + the PDL verify kernel, captured and replayed.""" + lay = _layer(batch, num_spec, seed=11 + batch) + ref, graphed = _clone(lay), _clone(lay) + main, side = torch.cuda.Stream(), torch.cuda.Stream() + fork, join = torch.cuda.Event(), torch.cuda.Event() + x_static = lay["x"].clone() + + def step(state): + fork.record(main) + with torch.cuda.stream(side): + side.wait_event(fork) + bfa = _fab_fa()(x_static, state["w_bfa"], H, state["w_fb"]) + join.record(side) + main.wait_event(join) + return _ops()(**_decode_kwargs(state, batch, num_spec, bfa=bfa, use_pdl=True)) + + # Eager warm-up on a scratch copy (compiles, creates the fab workspace outside capture). + scratch = _clone(lay) + with torch.cuda.stream(main): + step(scratch) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=main): + out_graph = step(graphed) + torch.cuda.synchronize() + for rnd in range(3): + out_ref = _reference(ref, batch, num_spec) + graph.replay() + torch.cuda.synchronize() + _assert_same(f"replay {rnd} out", out_graph, out_ref) + for key in MUTATED: + _assert_same(f"replay {rnd} {key}", graphed[key], ref[key]) + # the next round: new tokens and acceptance, written into the captured buffers + x_new = (lay["x"].float() * (0.5 + 0.25 * rnd)).bfloat16() + x_static.copy_(x_new) + ref["x"] = x_new + for state in (ref, graphed): + _next_acceptance(state, batch, num_spec, seed=90 + rnd) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_compact.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_compact.py new file mode 100644 index 000000000000..992fe31bfa4b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_compact.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``trtllm::kda_mtp_decode`` two-CTAs-per-SM variants vs the one-CTA kernel, bit for bit. + +* The compact layout of ``kda_decode_mtp_kernel`` (``mtp_ops._compact_layout_enabled``, on + SM107) only re-places shared-memory / TMEM buffers; every value-producing instruction is + unchanged. +* The MXFP8 epilogue of ``kda_decode_mtp_2cta_kernel`` (``mtp_ops._mxfp8_two_cta_enabled``, on + SM107) applies ``kda_decode_mtp_kernel``'s quantization to the same normalized values. + +The tests force each selection on or off, so both kernels of a pair run on SM100 / SM103 too. + +So the outputs (E4M3 values and the UE8M0 scales of the valid rows for MXFP8 output), the committed +recurrent state, the conv windows and the replay caches must equal those of the default kernel, +over two chained verify rounds with mixed per-slot accepted counts (round 2 replays round-1 +drafts), from a zero or a populated recurrent state. Each test checks which kernel variant served +the calls, so that it cannot pass vacuously. + +Inputs follow the runtime contract: packed token layout (strided QKVG section views), bf16 state +pool (V-first, strided per-layer slice), bf16 conv / k / v / beta replay caches, fp32 gate cache, +fused gated RMSNorm; MXFP8 output is what SM107 requests for the prequantized o_proj. + +Requires 1 SM100/SM103/SM107 GPU and nvidia-cutlass-dsl. +""" + +from collections.abc import MutableMapping + +import pytest +import torch + +_HAVE_DEPS = True +_DEP_ERR = None +try: + import cuda.bindings.driver # noqa: F401 + import cutlass # noqa: F401 + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda import ( # noqa: E501 + mtp_ops as kda_ops, + ) + + # This GPU's default kernel selections (SM107: on), before any test forces them. + _DEFAULT_SELECTORS = { + name: getattr(kda_ops, name) + for name in ("_compact_layout_enabled", "_mxfp8_two_cta_enabled") + } +except ImportError as e: + _HAVE_DEPS = False + _DEP_ERR = str(e) + + +def _is_sm100_family(): + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability(0) in ((10, 0), (10, 3), (10, 7)) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), + pytest.mark.skipif(not _is_sm100_family(), reason="needs sm100/sm103/sm107"), + pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), +] + +K = V = 128 +W = 4 +POOL = 80 +MUTATED = ("cs_q", "cs_k", "cs_v", "recurrent_state", "k_cache", "g_cache", "v_cache", "beta_cache") +COMPACT = "_compact_layout_enabled" +MXFP8_2CTA = "_mxfp8_two_cta_enabled" + +# (heads per rank, requests, drafts): helix8 (12 heads), helix16 (6), dep16 (96) at batch 1 / 4 / +# 8 / 16 with seven drafts and batch 16 / 32 with five drafts. +DECODE_SHAPES = [ + pytest.param(h, n, s, id=f"{name}_b{n}_mtp{s}") + for name, h in (("helix8", 12), ("helix16", 6), ("dep16", 96)) + for n, s in ((1, 7), (4, 7), (8, 7), (16, 7), (16, 5), (32, 5)) +] + + +def _num_sms(): + return torch.cuda.get_device_properties(0).multi_processor_count + + +def _make_inputs(N, H, num_spec, seed, zero_state): + gen = torch.Generator().manual_seed(seed) + dev = "cuda" + d = H * K + T = N * (1 + num_spec) + S = W - 1 + num_spec + + def rn(*shape, dtype=torch.float32, scale=1.0): + return (torch.randn(*shape, generator=gen) * scale).to(dtype).to(dev) + + def conv(): + # [layers, slots, window, dim] -> dim-contiguous view of layer 1 + return rn(2, POOL, S, d, dtype=torch.bfloat16).transpose(-1, -2)[1] + + # fused [q | k | v | output-gate] projection: strided section views + qkvg = rn(T, 4 * d, dtype=torch.bfloat16, scale=0.5) + bfa = rn(T, 128 + H, dtype=torch.bfloat16) + state = rn(2, POOL, H, V, K, dtype=torch.bfloat16, scale=0.05) + if zero_state: + state.zero_() + slots = torch.randperm(POOL, generator=gen)[:N].to(torch.int32) + accepted = torch.zeros(POOL, dtype=torch.int32) + if not zero_state: + accepted[slots.long()] = torch.randint( + 0, num_spec + 1, (N,), generator=gen, dtype=torch.int32 + ) + heads_padded = (H + 7) // 8 * 8 + return dict( + x_q=qkvg[:, 0 * d : 1 * d].view(1, T, H, K), + x_k=qkvg[:, 1 * d : 2 * d].view(1, T, H, K), + x_v=qkvg[:, 2 * d : 3 * d].view(1, T, H, V), + onorm_g=qkvg[:, 3 * d : 4 * d].view(1, T, H, V), + g=rn(T, d, dtype=torch.bfloat16, scale=0.5).view(1, T, H, K), + beta=bfa[:, 128 : 128 + H].view(1, T, H), + w_q=rn(d, W, scale=0.5), + w_k=rn(d, W, scale=0.5), + w_v=rn(d, W, scale=0.5), + A_log=rn(H, scale=0.5), + dt_bias=rn(d, scale=0.5), + cs_q=conv(), + cs_k=conv(), + cs_v=conv(), + recurrent_state=state[1], + k_cache=rn(2, POOL, num_spec, d, dtype=torch.bfloat16, scale=0.08)[1], + g_cache=(-torch.rand(2, POOL, num_spec, d, generator=gen) * 0.5).to(dev)[1], + v_cache=rn(2, POOL, num_spec, d, dtype=torch.bfloat16)[1], + beta_cache=torch.rand(2, POOL, num_spec, heads_padded, generator=gen) + .to(torch.bfloat16) + .to(dev)[1, ..., :H], + ssm_state_indices=slots.to(dev), + cu_seqlens=None, + num_spec=num_spec, + num_accepted_tokens=accepted.to(dev), + lower_bound=-5.0, + scale=K**-0.5, + packed_token_layout=True, + onorm_weight=(0.75 + 0.5 * torch.rand(V, generator=gen)).to(dev), + onorm_eps=1e-5, + fuse_output_norm=True, + ) + + +def _clone(inp): + out = dict(inp) + for name in MUTATED: + t = inp[name] + base = t._base if t._base is not None else t + out[name] = torch.as_strided(base.clone(), t.shape, t.stride(), t.storage_offset()) + return out + + +def _clone_out(o): + # quantize_output=True returns (E4M3 activations, R128c4 UE8M0 scales) + return tuple(t.clone() for t in o) if isinstance(o, tuple) else o.clone() + + +def _scale_bytes(scales, rows, heads): + """The scale bytes of rows [0, rows) (four copies per (row, head)); the padding rows of the + R128c4 buffer are never written.""" + r = torch.arange(rows, device=scales.device)[:, None] + h = torch.arange(heads, device=scales.device)[None, :] + off = (r // 128) * heads * 512 + h * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + off = (off[..., None] + torch.arange(4, device=scales.device)).reshape(-1) + return scales.view(torch.uint8)[off] + + +def _assert_out_equal(a, b, heads): + if isinstance(a, tuple): + assert torch.equal(a[0].view(torch.uint8), b[0].view(torch.uint8)), "E4M3 values" + rows = a[0].shape[0] + assert torch.equal(_scale_bytes(a[1], rows, heads), _scale_bytes(b[1], rows, heads)), ( + "UE8M0 scales" + ) + else: + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), "bf16 output" + + +class _RecordingCache(MutableMapping): + """Pass-through view of the op's compile cache that records the keys looked up.""" + + def __init__(self, cache): + self.cache = cache + self.used = [] + + def __getitem__(self, key): + self.used.append(key) + return self.cache[key] + + def __setitem__(self, key, value): + self.cache[key] = value + + def __delitem__(self, key): + del self.cache[key] + + def __iter__(self): + return iter(self.cache) + + def __len__(self): + return len(self.cache) + + def __contains__(self, key): + return key in self.cache + + +def _variant(key): + """Kernel variant of a compile-cache key.""" + if key[0] == "2cta": + return "2cta_mxfp8" + if key[0] == "v2": + return "v2" + return "v1_compact" if key[-1] else "v1" + + +def _run_rounds(inp, quantize_output, selectors, monkeypatch, attrs=None): + """Two chained verify rounds with the kernel selections ``selectors`` (name -> forced on / + off, ``None``: this GPU's default) and the module settings ``attrs`` (name -> value); + returns the outputs, the mutated inputs and the kernel variants that served the calls.""" + for name, value in selectors.items(): + default = _DEFAULT_SELECTORS[name] + monkeypatch.setattr( + kda_ops, name, default if value is None else (lambda device, value=value: value) + ) + for name, value in (attrs or {}).items(): + monkeypatch.setattr(kda_ops, name, value) + recorder = _RecordingCache(kda_ops._compiled_cache) + monkeypatch.setattr(kda_ops, "_compiled_cache", recorder) + x = _clone(inp) + x["num_accepted_tokens"] = x["num_accepted_tokens"].clone() + x["quantize_output"] = quantize_output + o1 = _clone_out(kda_ops.kda_mtp_decode_impl(**x)) + # round 2: the sampler accepted a different number of drafts per slot + acc = x["num_accepted_tokens"] + slots = x["ssm_state_indices"].long() + acc[slots] = ( + acc[slots] + 1 + torch.arange(len(slots), device=acc.device, dtype=acc.dtype) % 3 + ) % (x["num_spec"] + 1) + o2 = _clone_out(kda_ops.kda_mtp_decode_impl(**x)) + torch.cuda.synchronize() + monkeypatch.setattr(kda_ops, "_compiled_cache", recorder.cache) + return o1, o2, x, sorted({_variant(k) for k in recorder.used}) + + +def _assert_bitwise(ref, new, heads): + ref1, ref2, ref_x, _ = ref + new1, new2, new_x, _ = new + _assert_out_equal(new1, ref1, heads) + _assert_out_equal(new2, ref2, heads) + for name in MUTATED: + assert torch.equal(new_x[name], ref_x[name]), name + # sanity: the kernel ran and produced finite values + act = ref2[0] if isinstance(ref2, tuple) else ref2 + assert torch.isfinite(act.float()).all() + + +@pytest.mark.parametrize("initial_state", ["zero", "populated"]) +@pytest.mark.parametrize("quantize_output", [False, True], ids=["bf16_out", "mxfp8_out"]) +@pytest.mark.parametrize( + "N,H,num_spec", [(40, 96, 3), (72, 96, 3), (16, 96, 2), (24, 16, 3), (64, 12, 3)] +) +def test_compact_layout_is_bitwise_identical( + N, H, num_spec, quantize_output, initial_state, monkeypatch +): + if H * N <= _num_sms(): + pytest.skip(f"H*N={H * N} <= {_num_sms()} SMs: single wave, compact layout not taken") + inp = _make_inputs(N, H, num_spec, seed=1234 + N, zero_state=initial_state == "zero") + # BF16 output reaches kda_decode_mtp_kernel only with the v2 kernel off + attrs = {"_MTP_V2": False} + ref = _run_rounds(inp, quantize_output, {MXFP8_2CTA: False, COMPACT: False}, monkeypatch, attrs) + new = _run_rounds(inp, quantize_output, {MXFP8_2CTA: False, COMPACT: True}, monkeypatch, attrs) + assert ref[3] == ["v1"] and new[3] == ["v1_compact"], (ref[3], new[3]) + _assert_bitwise(ref, new, H) + + +@pytest.mark.parametrize("initial_state", ["zero", "populated"]) +@pytest.mark.parametrize( + "H,N,num_spec", DECODE_SHAPES + [pytest.param(96, 16, 3, id="dep16_b16_mtp3")] +) +def test_mxfp8_two_cta_is_bitwise_identical(H, N, num_spec, initial_state, monkeypatch): + # (grids of up to SPLIT_MAX_TILES tiles: one CTA per tile vs the cluster pairs of the reference) + inp = _make_inputs(N, H, num_spec, seed=4321 + 7 * N + H, zero_state=initial_state == "zero") + ref = _run_rounds(inp, True, {COMPACT: False, MXFP8_2CTA: False}, monkeypatch) + new = _run_rounds(inp, True, {COMPACT: False, MXFP8_2CTA: True}, monkeypatch) + assert ref[3] == ["v1"] and new[3] == ["2cta_mxfp8"], (ref[3], new[3]) + _assert_bitwise(ref, new, H) + + +@pytest.mark.parametrize("initial_state", ["zero", "populated"]) +@pytest.mark.parametrize("quantize_output", [False, True], ids=["bf16_out", "mxfp8_out"]) +@pytest.mark.parametrize("H,N,num_spec", DECODE_SHAPES) +def test_decode_shapes_default_matches_one_cta_kernels( + H, N, num_spec, quantize_output, initial_state, monkeypatch +): + """The variants selected by default on this GPU reproduce the one-CTA kernels.""" + inp = _make_inputs(N, H, num_spec, seed=99 + 7 * N + H, zero_state=initial_state == "zero") + ref = _run_rounds(inp, quantize_output, {COMPACT: False, MXFP8_2CTA: False}, monkeypatch) + new = _run_rounds(inp, quantize_output, {COMPACT: None, MXFP8_2CTA: None}, monkeypatch) + assert ref[3] == (["v1"] if quantize_output else ["v2"]), ref[3] + # SM107: MXFP8 output runs the two-CTA kernel; other calls are unchanged + two_cta = quantize_output and torch.cuda.get_device_capability(0) == (10, 7) + assert new[3] == (["2cta_mxfp8"] if two_cta else ref[3]), new[3] + _assert_bitwise(ref, new, H) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_cute_parity.py new file mode 100644 index 000000000000..5f215b725298 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_cute_parity.py @@ -0,0 +1,1425 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parity: the ``kda_mtp_decode`` entry's CuTe kernel vs two references. + +References: + +1. A pure-torch fp32 CPU golden (vendored from the kernel drop's + ``cpu_reference`` self-check). +2. An FLA sequential reference built from the exact op sequence + ``K3KDALinearAttention.forward_verify`` uses in-tree: per-step fp32 causal + conv + SiLU followed by ``fla.ops.kda.fused_recurrent_kda`` with + ``use_qk_l2norm/use_gate/use_beta_sigmoid`` in kernel, ``lower_bound``, + ``state_v_first=True``. + +Agreement of all three establishes both that the kernel is internally +correct and that it computes the same function the model's sequential +verify path computes — i.e. it is a drop-in replacement. + +Round-2 tests exercise the kernel's replay mode (``num_accepted_tokens`` +mixed per request), chained from round-1 CPU-golden outputs. H=8 (K3 +state-TP4) and H=6 (TP16) are the production per-rank head counts outside +the drop's benchmark-tuned set {2, 12, 32}; the vendored v_row hoist fix +makes them compile, and this test validates them numerically. + +Requires: 1 GPU (sm100/sm103/sm107), fla-core, nvidia-cutlass-dsl, +cuda-bindings. Skips cleanly otherwise. +""" + +import pytest +import torch +import torch.nn.functional as F + +_HAVE_DEPS = True +_DEP_ERR = None +try: + import cuda.bindings.driver # noqa: F401 + import cutlass # noqa: F401 + from fla.ops.kda import fused_recurrent_kda # noqa: F401 +except ImportError as e: + _HAVE_DEPS = False + _DEP_ERR = str(e) + + +def _is_blackwell(): + if not torch.cuda.is_available(): + return False + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor in (100, 103, 107) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), + pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103/sm107"), + pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), +] + +# (N, H): the drop's benchmark-tuned shapes plus K3's production per-rank +# head counts (96 KDA heads: 12 at TP8, 8 at state-TP 12, 6 at TP16). +SHAPES = [ + (128, 2), + (32, 12), + (32, 32), + (32, 8), + (32, 6), +] +M = 2 # NUM_SPEC — spec-token count per verify round + + +def _silu(x): + return x * torch.sigmoid(x) + + +def make_conv_data(B, H, K=128, V=128, M=2, W=4, lower_bound=-5.0, scale=None, seed=2025): + """Random inputs + caches in the op's layout contract (from the drop).""" + torch.manual_seed(seed) + scale = K**-0.5 if scale is None else scale + device = "cuda" + dim = H * K + T = 2 * M + 1 + T_total = B * T + S = W - 1 + M + + x_q = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device=device) + x_k = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device=device) + x_v = torch.randn(1, T_total, H, V, dtype=torch.bfloat16, device=device) + g = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device=device) + beta = torch.randn(1, T_total, H, dtype=torch.bfloat16, device=device) + w_q = torch.randn(dim, W, dtype=torch.float32, device=device) + w_k = torch.randn(dim, W, dtype=torch.float32, device=device) + w_v = torch.randn(dim, W, dtype=torch.float32, device=device) + A_log = torch.randn(H, dtype=torch.float32, device=device) + dt_bias = torch.randn(dim, dtype=torch.float32, device=device) + # dim-contiguous extended conv caches: allocate [B, S, dim], transpose. + # raw conv inputs are cached in bf16 like the projections that produce them + cs_q = torch.randn(B, S, dim, dtype=torch.bfloat16, device=device).transpose(1, 2) + cs_k = torch.randn(B, S, dim, dtype=torch.bfloat16, device=device).transpose(1, 2) + cs_v = torch.randn(B, S, dim, dtype=torch.bfloat16, device=device).transpose(1, 2) + initial_state_kfirst = torch.randn(B, H, K, V, dtype=torch.float32, device=device) + initial_state_cute = initial_state_kfirst.permute(0, 1, 3, 2).contiguous() + cu_seqlens = torch.arange(0, B * T + 1, T, dtype=torch.int32, device=device) + ssm_state_indices = torch.arange(B, dtype=torch.int32, device=device) + num_accepted_tokens = torch.zeros(B, dtype=torch.int32, device=device) + # replay caches in the runtime's dtypes: k / v / beta bf16, gate fp32 + k_cache = torch.zeros(B, M, dim, dtype=torch.bfloat16, device=device) + g_cache = torch.zeros(B, M, dim, dtype=torch.float32, device=device) + v_cache = torch.zeros(B, M, H * V, dtype=torch.bfloat16, device=device) + beta_cache = torch.zeros(B, M, H, dtype=torch.bfloat16, device=device) + return { + "x_q": x_q, + "x_k": x_k, + "x_v": x_v, + "w_q": w_q, + "w_k": w_k, + "w_v": w_v, + "g": g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + "cs_q": cs_q, + "cs_k": cs_k, + "cs_v": cs_v, + "initial_state_kfirst": initial_state_kfirst, + "initial_state_cute": initial_state_cute, + "k_cache": k_cache, + "g_cache": g_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + "cu_seqlens": cu_seqlens, + "ssm_state_indices": ssm_state_indices, + "num_accepted_tokens": num_accepted_tokens, + "B": B, + "H": H, + "K": K, + "V": V, + "M": M, + "W": W, + "T": T, + "T_total": T_total, + "lower_bound": lower_bound, + "scale": scale, + } + + +def cpu_reference(data): + """fp32 pure-torch golden with replay semantics (from the drop).""" + B, H, K, V, M, W = (data["B"], data["H"], data["K"], data["V"], data["M"], data["W"]) + T = data["T"] + lower_bound = data["lower_bound"] + scale = data["scale"] + w_q = data["w_q"].float().cpu() + w_k = data["w_k"].float().cpu() + w_v = data["w_v"].float().cpu() + A_log = data["A_log"].float().cpu() + dt_bias = data["dt_bias"].float().cpu() + x_q = data["x_q"].float().cpu() + x_k = data["x_k"].float().cpu() + x_v = data["x_v"].float().cpu() + g = data["g"].float().cpu() + beta = data["beta"].float().cpu() + cs_q = data["cs_q"].float().cpu().clone() + cs_k = data["cs_k"].float().cpu().clone() + cs_v = data["cs_v"].float().cpu().clone() + # fp32 working copies; new entries are rounded to the caches' dtypes like the kernel does + k_cache = data["k_cache"].float().cpu().clone() + g_cache = data["g_cache"].float().cpu().clone() + v_cache = data["v_cache"].float().cpu().clone() + beta_cache = data["beta_cache"].float().cpu().clone() + k_dtype, v_dtype, beta_dtype = ( + data["k_cache"].dtype, + data["v_cache"].dtype, + data["beta_cache"].dtype, + ) + ht = data["initial_state_kfirst"].float().cpu().clone() + ht_commit = ht.clone() + out = torch.zeros(1, B * T, H, V, dtype=torch.float32) + + for n in range(B): + bos = n * T + slot = n + commit_len = int(data["num_accepted_tokens"][n].item()) + T_loop = commit_len + 1 + M + for h in range(H): + hk = h * K + hv = h * V + hist_q = cs_q[slot, hk : hk + K, : W - 1].clone() + hist_k = cs_k[slot, hk : hk + K, : W - 1].clone() + hist_v = cs_v[slot, hv : hv + V, : W - 1].clone() + h_state = ht[slot, h].clone() + for i_t in range(T_loop): + if i_t < commit_len: + # replayed drafts produce no output; their q is not cached + q_t = torch.zeros(K, dtype=torch.float32) + k_t = k_cache[slot, i_t, hk : hk + K] + gk_t = g_cache[slot, i_t, hk : hk + K] + v_t = v_cache[slot, i_t, hv : hv + V] + beta_t = beta_cache[slot, i_t, h] + xq_raw = cs_q[slot, hk : hk + K, W - 1 + i_t] + xk_raw = cs_k[slot, hk : hk + K, W - 1 + i_t] + xv_raw = cs_v[slot, hv : hv + V, W - 1 + i_t] + hist_q = torch.cat([hist_q[:, 1:], xq_raw.unsqueeze(-1)], dim=1) + hist_k = torch.cat([hist_k[:, 1:], xk_raw.unsqueeze(-1)], dim=1) + hist_v = torch.cat([hist_v[:, 1:], xv_raw.unsqueeze(-1)], dim=1) + else: + token = bos + i_t + xq_raw = x_q[0, token, h] + xk_raw = x_k[0, token, h] + xv_raw = x_v[0, token, h] + cq = (torch.cat([hist_q, xq_raw.unsqueeze(-1)], dim=-1) * w_q[hk : hk + K]).sum( + dim=-1 + ) + ck = (torch.cat([hist_k, xk_raw.unsqueeze(-1)], dim=-1) * w_k[hk : hk + K]).sum( + dim=-1 + ) + cv = (torch.cat([hist_v, xv_raw.unsqueeze(-1)], dim=-1) * w_v[hv : hv + V]).sum( + dim=-1 + ) + q_t = F.normalize(cq / (1.0 + torch.exp(-cq)), p=2, dim=-1) * scale + k_t = F.normalize(ck / (1.0 + torch.exp(-ck)), p=2, dim=-1) + v_t = cv / (1.0 + torch.exp(-cv)) + gr = g[0, token, h] + dt_bias[hk : hk + K] + gk_t = lower_bound * torch.sigmoid(gr * torch.exp(A_log[h])) + beta_t = torch.sigmoid(beta[0, token, h]) + hist_q = torch.cat([hist_q[:, 1:], xq_raw.unsqueeze(-1)], dim=1) + hist_k = torch.cat([hist_k[:, 1:], xk_raw.unsqueeze(-1)], dim=1) + hist_v = torch.cat([hist_v[:, 1:], xv_raw.unsqueeze(-1)], dim=1) + + decay = torch.exp(gk_t) + h_state = h_state * decay.unsqueeze(1) + sum_hk = (h_state * k_t.unsqueeze(1)).sum(dim=0) + v_new = (v_t - sum_hk) * beta_t + h_state = h_state + k_t.unsqueeze(1) * v_new.unsqueeze(0) + o_t = (h_state * q_t.unsqueeze(1)).sum(dim=0) + if i_t >= commit_len: + out[0, bos + i_t, h] = o_t + if i_t == commit_len: + ht_commit[slot, h] = h_state + cs_q[slot, hk : hk + K, : W - 1] = hist_q + cs_k[slot, hk : hk + K, : W - 1] = hist_k + cs_v[slot, hv : hv + V, : W - 1] = hist_v + if i_t > commit_len: + cache_pos = i_t - commit_len - 1 + k_cache[slot, cache_pos, hk : hk + K] = k_t.to(k_dtype).float() + g_cache[slot, cache_pos, hk : hk + K] = gk_t + v_cache[slot, cache_pos, hv : hv + V] = v_t.to(v_dtype).float() + beta_cache[slot, cache_pos, h] = beta_t.to(beta_dtype).float() + cs_q[slot, hk : hk + K, W - 1 + cache_pos] = xq_raw + cs_k[slot, hk : hk + K, W - 1 + cache_pos] = xk_raw + cs_v[slot, hv : hv + V, W - 1 + cache_pos] = xv_raw + + return { + "out": out, + "recurrent_state": ht_commit, + "k_cache": k_cache, + "g_cache": g_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + "cs_q": cs_q, + "cs_k": cs_k, + "cs_v": cs_v, + } + + +def cute_run( + data, + zero_accepted_hint=False, + packed_token_layout=False, + out=None, + fuse_output_norm=False, + quantize_output=False, + beta_cache_override=None, + state_override=None, +): + """Run the op, optionally substituting a pre-populated beta-cache view or a + caller-allocated recurrent-state tensor (e.g. a strided slice of a larger cache).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.kda_mtp_decode import ( # noqa: E501 + kda_mtp_decode, + ) + + if state_override is None: + state = data["initial_state_cute"].clone() + else: + state = state_override + state.copy_(data["initial_state_cute"]) + cs = {} + for name in ("cs_q", "cs_k", "cs_v"): + src = data[name] + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + cs[name] = dst + k_cache = data["k_cache"].clone() + g_cache = data["g_cache"].clone() + v_cache = data["v_cache"].clone() + beta_cache = data["beta_cache"].clone() if beta_cache_override is None else beta_cache_override + op_kwargs = dict( + x_q=data["x_q"], + x_k=data["x_k"], + x_v=data["x_v"], + w_q=data["w_q"], + w_k=data["w_k"], + w_v=data["w_v"], + cs_q=cs["cs_q"], + cs_k=cs["cs_k"], + cs_v=cs["cs_v"], + g=data["g"], + beta=data["beta"], + A_log=data["A_log"], + dt_bias=data["dt_bias"], + recurrent_state=state, + k_cache=k_cache, + g_cache=g_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=data["ssm_state_indices"], + cu_seqlens=data["cu_seqlens"], + num_spec=data["M"], + num_accepted_tokens=data["num_accepted_tokens"], + lower_bound=data["lower_bound"], + scale=data["scale"], + zero_accepted_hint=zero_accepted_hint, + packed_token_layout=packed_token_layout, + onorm_g=data.get("onorm_g"), + onorm_weight=data.get("onorm_weight"), + onorm_eps=data.get("onorm_eps", 1e-5), + fuse_output_norm=fuse_output_norm, + ) + output_scale = None + if quantize_output: + op_kwargs["fuse_output_norm"] = True + out, output_scale = kda_mtp_decode(**op_kwargs, quantize_output=True) + elif out is None: + out = kda_mtp_decode(**op_kwargs) + else: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda.mtp_ops import ( # noqa: E501 + kda_mtp_decode_impl, + ) + + out = kda_mtp_decode_impl(out=out, **op_kwargs) + return { + "out": out, + "output_scale": output_scale, + # committed state back in K-first layout for CPU-golden comparison + "recurrent_state": state.permute(0, 1, 3, 2).contiguous(), + "state_v_first": state, + "k_cache": k_cache, + "g_cache": g_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + "cs_q": cs["cs_q"], + "cs_k": cs["cs_k"], + "cs_v": cs["cs_v"], + } + + +def _pack_new_token_projection_views(data): + """Build production-like section views over compact fused GEMM outputs.""" + B, H, K, M, source_T = ( + data["B"], + data["H"], + data["K"], + data["M"], + data["T"], + ) + accepted = data["num_accepted_tokens"].cpu() + rows = torch.cat( + [ + torch.arange(n * source_T + int(accepted[n]), n * source_T + int(accepted[n]) + M + 1) + for n in range(B) + ] + ).cuda() + num_rows = B * (M + 1) + dim = H * K + + # q/k/v are adjacent sections of [q | k | v | onorm_g]. Their token + # stride is therefore 4 * dim, not dim. + qkvg = torch.empty(1, num_rows, 4 * dim, dtype=torch.bfloat16, device="cuda") + qkvg[..., :dim].copy_(data["x_q"][0, rows].reshape(1, num_rows, dim)) + qkvg[..., dim : 2 * dim].copy_(data["x_k"][0, rows].reshape(1, num_rows, dim)) + qkvg[..., 2 * dim : 3 * dim].copy_(data["x_v"][0, rows].reshape(1, num_rows, dim)) + qkvg[..., 3 * dim : 4 * dim].copy_(data["g"][0, rows].reshape(1, num_rows, dim)) + + # beta is the H-wide tail of the padded [f_a | beta] projection. + bfa_width = ((K + H + 7) // 8) * 8 + bfa = torch.empty(1, num_rows, bfa_width, dtype=torch.bfloat16, device="cuda") + bfa[..., K : K + H].copy_(data["beta"][0, rows].reshape(1, num_rows, H)) + + packed = dict(data) + packed.update( + x_q=qkvg[..., :dim].view(1, num_rows, H, K), + x_k=qkvg[..., dim : 2 * dim].view(1, num_rows, H, K), + x_v=qkvg[..., 2 * dim : 3 * dim].view(1, num_rows, H, K), + onorm_g=qkvg[..., 3 * dim : 4 * dim].view(1, num_rows, H, K), + g=data["g"][:, rows].contiguous(), + beta=bfa[..., K : K + H], + cu_seqlens=None, + T=M + 1, + T_total=num_rows, + ) + assert packed["x_q"].stride(1) == 4 * dim + assert packed["x_k"].stride(1) == 4 * dim + assert packed["x_v"].stride(1) == 4 * dim + assert packed["onorm_g"].stride(1) == 4 * dim + assert packed["beta"].stride(1) == bfa_width + return packed, rows + + +def _fla_sequential_reference(data, num_accepted): + """Per-request sequential conv+SiLU (fp32 torch) + fused_recurrent_kda. + + Mirrors ``K3KDALinearAttention.forward_verify``'s op sequence, extended with + the replay prefix: for request ``n`` with ``a = num_accepted[n]``, the + processed token sequence is ``a`` cached raw tokens (re-convolved from + the extended conv-cache slots) followed by the ``1 + M`` new tokens. + Returns out rows (new tokens only) and the committed state (after the + first new token, FLA/cute ``[B, H, V, K]`` layout). + """ + from fla.ops.kda import fused_recurrent_kda + + B, H, K, V, W = data["B"], data["H"], data["K"], data["V"], data["W"] + M, T = data["M"], data["T"] + dim = H * K + dev = data["x_q"].device + out = torch.zeros(1, B * T, H, V, dtype=torch.float32, device=dev) + committed = torch.zeros(B, H, V, K, dtype=torch.float32, device=dev) + + w_q, w_k, w_v = (data[k].float() for k in ("w_q", "w_k", "w_v")) + x_q, x_k, x_v = (data[k].float() for k in ("x_q", "x_k", "x_v")) + g_all, beta_all = data["g"], data["beta"] + + for n in range(B): + a = int(num_accepted[n]) + bos = n * T + hist = { + "q": data["cs_q"][n, :, : W - 1].float().clone(), + "k": data["cs_k"][n, :, : W - 1].float().clone(), + "v": data["cs_v"][n, :, : W - 1].float().clone(), + } + state = data["initial_state_cute"][n : n + 1].float().clone() + + for i_t in range(a + 1 + M): + if i_t < a: # replay a cached token (raw x from cache slots) + xq = data["cs_q"][n, :, W - 1 + i_t].float() + xk = data["cs_k"][n, :, W - 1 + i_t].float() + xv = data["cs_v"][n, :, W - 1 + i_t].float() + tok = None + g_t = data["g_cache"][n, i_t].float() + beta_t = data["beta_cache"][n, i_t].float() + replay = True + else: + tok = bos + i_t + xq = x_q[0, tok].reshape(dim) + xk = x_k[0, tok].reshape(dim) + xv = x_v[0, tok].reshape(H * V) + replay = False + + def conv_step(hist_s, x_raw, w): + window = torch.cat([hist_s, x_raw.unsqueeze(-1)], dim=-1) + y = (window * w).sum(dim=-1) + return y, window[:, 1:] + + cq, hist["q"] = conv_step(hist["q"], xq, w_q) + ck, hist["k"] = conv_step(hist["k"], xk, w_k) + cv, hist["v"] = conv_step(hist["v"], xv, w_v) + + if replay: + # Replayed tokens use the cached post-processed k/g/v/beta + # exactly as the kernel does (delta rule applied directly). + k_t = data["k_cache"][n, i_t].float().view(H, K) + gk_t = g_t.view(H, K) + v_t = data["v_cache"][n, i_t].float().view(H, V) + st = state[0] + decay = torch.exp(gk_t) + st = st * decay.unsqueeze(1) + sum_hk = torch.einsum("hvk,hk->hv", st, k_t) + v_new = (v_t - sum_hk) * beta_t.unsqueeze(-1) + st = st + torch.einsum("hk,hv->hvk", k_t, v_new) + state = st.unsqueeze(0) + else: + # fp32 hand-off into FLA: the comparison target is the + # mathematical function, not the model's bf16 dataflow. + q_in = _silu(cq).view(1, 1, H, K) + k_in = _silu(ck).view(1, 1, H, K) + v_in = _silu(cv).view(1, 1, H, V) + o_t, state = fused_recurrent_kda( + q=q_in, + k=k_in, + v=v_in, + g=g_all[0, tok].view(1, 1, H, K), + beta=beta_all[0, tok].view(1, 1, H).float(), + A_log=data["A_log"], + dt_bias=data["dt_bias"], + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=data["lower_bound"], + state_v_first=True, + ) + out[0, tok] = o_t[0, 0].float() + if i_t == a: # first new (golden) token -> committed state + committed[n] = state[0].float() + return out, committed + + +def _cute_layout_state(cpu_out): + # CPU golden reports K-first [B, H, K, V]; cute/FLA layout is [B,H,V,K] + return cpu_out["recurrent_state"].permute(0, 1, 3, 2).contiguous() + + +def _assert_close(name, a, b, atol, rtol=0.0): + diff = (a.float() - b.float()).abs() + denom = b.float().abs().clamp_min(1.0) + ok = (diff <= atol + rtol * denom).all() + assert ok, ( + f"{name}: max_abs={diff.max().item():.3e} " + f"(atol={atol}, worst rel={((diff / denom).max()):.3e})" + ) + + +@pytest.mark.parametrize("B,H", SHAPES, ids=lambda v: str(v)) +def test_round1_zero_accepted(B, H): + """Fresh verify round (no replay): kernel vs CPU golden vs FLA seq.""" + data = make_conv_data(B, H, M=M, seed=2025) + T = data["T"] + + cpu = {k: v.cuda() for k, v in cpu_reference(data).items()} + cute_out = cute_run(data) + fla_out, fla_committed = _fla_sequential_reference(data, data["num_accepted_tokens"].cpu()) + + new_rows = torch.cat([torch.arange(n * T, n * T + 1 + M) for n in range(B)]).cuda() + + # Kernel vs the fp32 golden (tight: fp32 accumulation). + _assert_close( + "out(cute vs cpu)", cute_out["out"][0, new_rows], cpu["out"][0, new_rows], atol=2e-2 + ) + _assert_close( + "state(cute vs cpu)", cute_out["recurrent_state"], cpu["recurrent_state"], atol=1e-4 + ) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): + _assert_close(f"{name}(cute vs cpu)", cute_out[name], cpu[name], atol=2e-2, rtol=4e-3) + + # Kernel vs the FLA sequential path (the in-tree fused verify math). + _assert_close( + "out(cute vs fla)", + cute_out["out"][0, new_rows].float(), + fla_out[0, new_rows], + atol=5e-2, + rtol=5e-2, + ) + _assert_close( + "state(cute vs fla)", cute_out["state_v_first"], fla_committed, atol=5e-3, rtol=5e-3 + ) + + +@pytest.mark.parametrize( + "B,H,M", + [(128, 2, 2), (32, 12, 2), (32, 6, 2), (2, 12, 7), (16, 12, 7)], + ids=lambda v: str(v), +) +def test_round2_replay(B, H, M): + """Replay round: mixed num_accepted per request, chained from a CPU- + golden round 1. Validates the kernel's cache-replay state math (the + path the drop's own self-check never exercised). ``M=7`` covers the + production MTP7 verify (T_loop up to 15) at the TEP8 shape.""" + data = make_conv_data(B, H, M=M, seed=7) + T = data["T"] + + # Round 1 (all zero accepted) on the CPU golden to produce the caches + # and committed state that seed round 2. + r1 = cpu_reference(data) + + # Round 2 inputs: fresh tokens, caches/state/conv from round 1. + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + # Preserve the contract's dim-contiguous (transposed) layout. + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + data2["initial_state_kfirst"] = r1["recurrent_state"].cuda().contiguous() + data2["initial_state_cute"] = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).contiguous() + # Mixed acceptance: 0, 1, 2 cycling across requests. + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + cute2 = cute_run(data2) + fla_out2, fla_committed2 = _fla_sequential_reference(data2, accept) + + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + + _assert_close("out2(cute vs cpu)", cute2["out"][0, rows], cpu2["out"][0, rows], atol=2e-2) + _assert_close( + "state2(cute vs cpu)", cute2["recurrent_state"], cpu2["recurrent_state"], atol=1e-4 + ) + _assert_close( + "out2(cute vs fla)", cute2["out"][0, rows].float(), fla_out2[0, rows], atol=5e-2, rtol=5e-2 + ) + _assert_close( + "state2(cute vs fla)", cute2["state_v_first"], fla_committed2, atol=5e-3, rtol=5e-3 + ) + + +def _round2_inputs(B, H, M): + """Round-2 inputs chained from a CPU-golden round 1 (caches, conv windows, committed state) + with mixed acceptance 0..M cycling across requests. Returns (data2, accept).""" + data = make_conv_data(B, H, M=M, seed=7) + r1 = cpu_reference(data) + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + data2["initial_state_kfirst"] = r1["recurrent_state"].cuda().contiguous() + data2["initial_state_cute"] = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).contiguous() + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + return data2, accept + + +@pytest.mark.parametrize("B,H,M", [(2, 12, 7), (16, 12, 7)], ids=lambda v: str(v)) +def test_round2_replay_bf16_math(B, H, M, monkeypatch): + """bf16 tensor-core operand path (``TRTLLM_KDA_MTP_BF16_MATH=1``) on the MTP7 replay round. + The conv pools and replay caches do not depend on the recurrence precision and must match + the 3xTF32 kernel bit for bit; the committed state and the outputs are held to the tolerance + of the bf16 tensor-core references (FLA, prefill).""" + data2, accept = _round2_inputs(B, H, M) + T = data2["T"] + monkeypatch.setenv("TRTLLM_KDA_MTP_BF16_MATH", "0") + ref = cute_run(data2) + monkeypatch.setenv("TRTLLM_KDA_MTP_BF16_MATH", "1") + fast = cute_run(data2) + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + + for name in ("cs_q", "cs_k", "cs_v"): + _assert_close(f"{name}(fast vs 3xTF32)", fast[name], ref[name], atol=0.0) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + _assert_close(f"{name}(fast vs 3xTF32)", fast[name], ref[name], atol=0.0) + state_err = (fast["recurrent_state"] - cpu2["recurrent_state"]).abs().max().item() + ref_err = (ref["recurrent_state"] - cpu2["recurrent_state"]).abs().max().item() + out_err = (fast["out"][0, rows].float() - cpu2["out"][0, rows].float()).abs().max().item() + print( + f"[bf16_math B={B} H={H} M={M}] state max err vs fp32 golden " + f"{state_err:.3e} (3xTF32 kernel {ref_err:.3e}), out max err {out_err:.3e}" + ) + # measured max |err| vs the fp32 golden at (16,12,7): bf16 1.6e-2 (3xTF32: 4e-5) + tol = 3e-2 + _assert_close( + "state(fast vs cpu)", fast["recurrent_state"], cpu2["recurrent_state"], atol=tol, rtol=tol + ) + _assert_close( + "out(fast vs cpu)", fast["out"][0, rows], cpu2["out"][0, rows], atol=5e-2, rtol=5e-2 + ) + + +@pytest.mark.parametrize("B,H,M", [(2, 12, 7), (16, 12, 7)], ids=lambda v: str(v)) +def test_round2_replay_strided_state_cache(B, H, M): + """Replay round with the recurrent state living in a strided slice of a larger buffer + (slot stride 3 * H * V * K), like the runtime's per-layer view of the SSM cache. Covers the + kernel's own global-memory views (TMA state load/store), which only a non-contiguous state + exercises; ``B=2`` runs the cluster-pair form, ``B=16`` the single-CTA form.""" + data = make_conv_data(B, H, M=M, seed=7) + T = data["T"] + r1 = cpu_reference(data) + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + data2["initial_state_kfirst"] = r1["recurrent_state"].cuda().contiguous() + data2["initial_state_cute"] = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).contiguous() + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + + state_shape = tuple(data2["initial_state_cute"].shape) # (B, H, V, K) + big = torch.zeros((B, 3) + state_shape[1:], dtype=torch.float32, device="cuda") + strided_state = big[:, 1] + assert not strided_state.is_contiguous() + + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + cute2 = cute_run(data2, state_override=strided_state) + + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + _assert_close("out2(cute vs cpu)", cute2["out"][0, rows], cpu2["out"][0, rows], atol=2e-2) + _assert_close( + "state2(cute vs cpu)", cute2["recurrent_state"], cpu2["recurrent_state"], atol=1e-4 + ) + # the neighbouring slices of the buffer must be untouched + assert torch.count_nonzero(big[:, 0]) == 0 and torch.count_nonzero(big[:, 2]) == 0 + + +@pytest.mark.parametrize("B,H,M", [(2, 12, 7), (16, 12, 7)], ids=lambda v: str(v)) +def test_round2_replay_bf16_state(B, H, M): + """Replay round with a bf16 recurrent-state pool: the kernel reads the bf16 state (exact in + tf32) and rounds the committed state to nearest bf16. The CPU golden starts from the same + bf16-rounded state; its fp32 result is compared against the kernel's bf16 output with a + half-ulp (2^-9) bf16 tolerance. ``B=2`` runs the cluster-pair form, ``B=16`` the single-CTA form.""" + data = make_conv_data(B, H, M=M, seed=7) + T = data["T"] + r1 = cpu_reference(data) + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + state_bf16 = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).to(torch.bfloat16).contiguous() + data2["initial_state_cute"] = state_bf16 + data2["initial_state_kfirst"] = state_bf16.permute(0, 1, 3, 2).float().contiguous() + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + cute2 = cute_run(data2) + assert cute2["state_v_first"].dtype == torch.bfloat16 + + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + _assert_close("out2(cute vs cpu)", cute2["out"][0, rows], cpu2["out"][0, rows], atol=2e-2) + _assert_close( + "state2(cute bf16 vs cpu fp32)", + cute2["recurrent_state"].float(), + cpu2["recurrent_state"], + atol=2e-3, + rtol=4e-3, + ) + + +@pytest.mark.parametrize("B,H,num_spec", [(4, 6, 2), (2, 12, 7)]) +def test_packed_projection_views_match_shifted_metadata_replay(B, H, num_spec): + """Compact row-strided projection views match the legacy shifted layout.""" + data = make_conv_data(B=B, H=H, M=num_spec, seed=10) + request_accepted = torch.arange(B, dtype=torch.int32, device="cuda") + request_accepted[-1] = num_spec + data["num_accepted_tokens"] = request_accepted + packed, legacy_rows = _pack_new_token_projection_views(data) + + # The production count is pool-level, while request rows may map to any + # state slot. Move every request's initial pools to a non-identity slot + # and scatter its accepted count to the same slot. + slots = torch.arange(B - 1, -1, -1, dtype=torch.int32, device="cuda") + packed["ssm_state_indices"] = slots + for name in ( + "initial_state_kfirst", + "initial_state_cute", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + request_order = packed[name] + slot_order = torch.empty_like(request_order) + slot_order[slots.long()] = request_order + packed[name] = slot_order + pool_accepted = torch.zeros(B, dtype=torch.int32, device="cuda") + pool_accepted[slots.long()] = request_accepted + packed["num_accepted_tokens"] = pool_accepted + + legacy = cute_run(data) + actual = cute_run(packed, packed_token_layout=True) + poison = torch.full( + (1, B * (num_spec + 1), H, data["V"]), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + poisoned = cute_run(packed, packed_token_layout=True, out=poison) + assert poisoned["out"].data_ptr() == poison.data_ptr() + assert torch.isfinite(poisoned["out"]).all() + _assert_close("poisoned packed out", poisoned["out"], actual["out"], atol=1e-5) + _assert_close( + "packed out", + actual["out"].reshape(-1, data["H"], data["V"]), + legacy["out"][0, legacy_rows], + atol=1e-5, + ) + + # Stage3a oracle: fuse only gated RMSNorm into the packed CuTe consumer. + # The raw recurrence output above is BF16, matching the explicit BF16 + # round-trip boundary retained in the fused epilogue. + from tensorrt_llm._torch.modules.mamba.layernorm_gated import rms_norm_gated_token_major + + fused_data = dict(packed) + fused_data["onorm_weight"] = torch.linspace( + 0.75, 1.25, data["V"], dtype=torch.float32, device="cuda" + ) + fused_data["onorm_eps"] = 1e-5 + expected_norm = rms_norm_gated_token_major( + actual["out"].reshape(-1, data["V"]), + packed["onorm_g"].reshape(-1, data["H"], data["V"]), + fused_data["onorm_weight"], + fused_data["onorm_eps"], + gate_activation="sigmoid", + ) + fused_norm = cute_run( + fused_data, + packed_token_layout=True, + fuse_output_norm=True, + ) + assert torch.isfinite(fused_norm["out"]).all() + fused_norm_diff = ( + fused_norm["out"].float().reshape_as(expected_norm) - expected_norm.float() + ).abs() + fused_norm_rel_l2 = ( + fused_norm_diff.norm() / expected_norm.float().norm().clamp_min(1e-12) + ).item() + print( + f" fused_norm B={B} H={H} M={num_spec}: " + f"max_abs={fused_norm_diff.max().item():.3e} rel_l2={fused_norm_rel_l2:.3e}" + ) + _assert_close( + "fused gated RMSNorm", + fused_norm["out"].reshape(-1, data["V"]), + expected_norm, + atol=2e-2, + rtol=1e-2, + ) + fused_fp8 = cute_run( + fused_data, + packed_token_layout=True, + quantize_output=True, + ) + expected_rows = B * (num_spec + 1) + assert fused_fp8["out"].dtype is torch.float8_e4m3fn + assert fused_fp8["out"].shape == (expected_rows, H * data["V"]) + assert torch.isfinite(fused_fp8["out"].float()).all() + assert fused_fp8["output_scale"].dtype is torch.uint8 + assert fused_fp8["output_scale"].shape == ((expected_rows + 127) // 128 * 128 * H * 4,) + expected_fp8, expected_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + expected_norm.reshape(expected_rows, H * data["V"]), + ) + assert torch.equal(fused_fp8["out"], expected_fp8) + rows = torch.arange(expected_rows, device="cuda", dtype=torch.int64).view(-1, 1) + scale_columns = torch.arange(H * 4, device="cuda", dtype=torch.int64).view(1, -1) + offsets = ( + ((rows // 128 * H + scale_columns // 4) * 32 + rows % 32) * 4 + (rows % 128) // 32 + ) * 4 + scale_columns % 4 + assert torch.equal( + fused_fp8["output_scale"][offsets.flatten()], + expected_scale[offsets.flatten()], + ) + for name in ( + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close( + f"{name}(packed vs shifted)", actual[name][slots.long()], legacy[name], atol=1e-5 + ) + _assert_close(f"{name}(norm vs raw)", fused_norm[name], actual[name], atol=1e-5) + _assert_close(f"{name}(fp8 vs raw)", fused_fp8[name], actual[name], atol=1e-5) + + +@pytest.mark.parametrize("B,H", [(32, 12)], ids=lambda v: str(v)) +def test_zero_accepted_hint_variant(B, H): + """The zero_accepted_hint fast variant matches the general variant.""" + data = make_conv_data(B, H, M=M, seed=11) + general = cute_run(data, zero_accepted_hint=False) + fast = cute_run(data, zero_accepted_hint=True) + for name in ( + "out", + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close(f"{name}(fast vs general)", fast[name], general[name], atol=1e-5) + + +def test_misaligned_state_indices_after_aligned_warmup(): + """Mixed context/generation metadata remains valid after op warmup.""" + data = make_conv_data(B=1, H=6, M=M, seed=13) + expected = cute_run(data) + + index_storage = torch.empty(2, dtype=torch.int32, device="cuda") + index_storage[0] = -1 + index_storage[1:].copy_(data["ssm_state_indices"]) + misaligned_indices = index_storage[1:] + assert misaligned_indices.is_contiguous() + assert misaligned_indices.data_ptr() % 16 != 0 + + misaligned_data = dict(data) + misaligned_data["ssm_state_indices"] = misaligned_indices + actual = cute_run(misaligned_data) + + for name in ( + "out", + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close(f"{name}(misaligned vs aligned)", actual[name], expected[name], atol=1e-5) + + +@pytest.mark.parametrize("field", ("cu_seqlens", "num_accepted_tokens")) +def test_misaligned_scalar_metadata_after_aligned_warmup(field): + """Every scalar CuTe metadata argument accepts an offset int32 view.""" + data = make_conv_data(B=1, H=6, M=M, seed=17) + expected = cute_run(data) + + source = data[field] + storage = torch.empty(source.numel() + 1, dtype=torch.int32, device="cuda") + storage[0] = -1 + storage[1:].copy_(source) + misaligned = storage[1:] + assert misaligned.is_contiguous() + assert misaligned.data_ptr() % 16 != 0 + + misaligned_data = dict(data) + misaligned_data[field] = misaligned + actual = cute_run(misaligned_data) + + for name in ( + "out", + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close( + f"{name}({field} misaligned vs aligned)", actual[name], expected[name], atol=1e-5 + ) + + +@pytest.mark.parametrize( + "attention_mode,parallel_size,dtype,pool_size,num_spec,expected", + ( + pytest.param("dep", 8, torch.float32, 1177, 5, 16, id="dep8-mtp5-fp32"), + pytest.param("dep", 16, torch.float32, 1177, 7, 16, id="dep16-mtp7-fp32"), + pytest.param("tep", 8, torch.float32, 1177, 7, 16, id="tep8-mtp7-fp32"), + pytest.param("tep", 16, torch.float32, 1177, 2, 16, id="tep16-mtp2-fp32"), + pytest.param("tep", 16, torch.float32, 1177, 5, 8, id="tep16-mtp5-fp32"), + pytest.param("tep", 16, torch.float32, 1177, 7, 8, id="tep16-mtp7-fp32"), + pytest.param("tep", 16, torch.float32, 248, 7, 16, id="tep16-mtp7-even-pool"), + pytest.param("tep", 32, torch.float32, 1177, 7, 4, id="tep32-mtp7-fp32"), + pytest.param("dep", 16, torch.bfloat16, 1177, 7, 16, id="dep16-mtp7-bf16"), + pytest.param("tep", 8, torch.bfloat16, 1177, 7, 8, id="tep8-mtp7-bf16"), + pytest.param("tep", 16, torch.bfloat16, 1177, 7, 4, id="tep16-mtp7-bf16"), + pytest.param("tep", 32, torch.bfloat16, 1177, 7, 2, id="tep32-mtp7-bf16"), + ), +) +def test_beta_cache_alignment_is_derived_from_layout_and_dtype( + attention_mode, parallel_size, dtype, pool_size, num_spec, expected +): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda.mtp_ops import ( # noqa: E501 + _beta_cache_assumed_align, + ) + + local_heads = 96 if attention_mode == "dep" else 96 // parallel_size + parent = torch.empty(2, pool_size, num_spec, local_heads, dtype=dtype, device="cuda") + beta_cache = parent[0] + + assert _beta_cache_assumed_align(beta_cache) == expected + + +def test_beta_cache_alignment_uses_padded_physical_stride(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda.mtp_ops import ( # noqa: E501 + _beta_cache_assumed_align, + ) + + parent = torch.empty(2, 1177, 7, 8, dtype=torch.float32, device="cuda") + beta_cache = parent[0, ..., :6] + + assert beta_cache.shape == (1177, 7, 6) + assert beta_cache.stride() == (56, 8, 1) + assert _beta_cache_assumed_align(beta_cache) == 16 + + +@pytest.mark.parametrize( + "attention_mode,parallel_size,num_spec,expected_alignment", + ( + pytest.param("dep", 16, 7, 16, id="dep16-mtp7"), + pytest.param("tep", 8, 7, 16, id="tep8-mtp7"), + pytest.param("tep", 16, 5, 8, id="tep16-mtp5"), + pytest.param("tep", 16, 7, 8, id="tep16-mtp7"), + pytest.param("tep", 32, 7, 4, id="tep32-mtp7"), + ), +) +def test_beta_cache_sibling_layer_after_aligned_warmup( + attention_mode, parallel_size, num_spec, expected_alignment +): + """A cached kernel accepts the next real per-layer beta-cache slice.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda.mtp_ops import ( # noqa: E501 + _beta_cache_assumed_align, + ) + + local_heads = 96 if attention_mode == "dep" else 96 // parallel_size + data = make_conv_data(B=1, H=local_heads, M=num_spec, seed=29) + + parent = torch.zeros(2, 1177, num_spec, local_heads, dtype=torch.float32, device="cuda") + layer0, layer1 = parent.unbind(0) + layer0[0].copy_(data["beta_cache"][0]) + layer1[0].copy_(data["beta_cache"][0]) + + assert _beta_cache_assumed_align(layer0) == expected_alignment + assert layer0.data_ptr() % 16 == 0 + assert layer1.data_ptr() % expected_alignment == 0 + if expected_alignment < 16: + assert layer1.data_ptr() % (2 * expected_alignment) != 0 + + expected = cute_run(data, beta_cache_override=layer0) + actual = cute_run(data, beta_cache_override=layer1) + + for name in ( + "out", + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close( + f"{name}({attention_mode}{parallel_size}-mtp{num_spec})", + actual[name], + expected[name], + atol=1e-5, + ) + + +def _dim_contiguous_clone(cache): + clone = torch.empty( + cache.shape[0], cache.shape[2], cache.shape[1], dtype=cache.dtype, device=cache.device + ).transpose(1, 2) + clone.copy_(cache) + return clone + + +def _verify_after_round_one(B, H, num_spec, seed): + """Inputs of a verify that follows a first verify run by the kernel, with + a different number of the first round's drafts accepted per request. + + The conv windows start at bf16 values, as they do when the cache manager + seeds them from its bf16 conv pool.""" + data = make_conv_data(B, H, M=num_spec, seed=seed) + for name in ("cs_q", "cs_k", "cs_v"): + data[name].copy_(data[name].bfloat16().float()) + first = cute_run(data) + data = make_conv_data(B, H, M=num_spec, seed=seed + 1) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): + data[name] = first[name] + data["initial_state_cute"] = first["state_v_first"] + data["num_accepted_tokens"] = torch.arange(B, dtype=torch.int32, device="cuda") % (num_spec + 1) + return data + + +@pytest.mark.parametrize("B,H,num_spec", [(32, 12, 2), (32, 6, 2), (4, 6, 7)], ids=lambda v: str(v)) +def test_commit_pending_drafts_matches_the_next_verify(B, H, num_spec): + """Committing the pending drafts first leaves the next verify's results + unchanged. + + This is what the cache manager does across a suspend. + ``kda_mtp_commit_pending_drafts`` folds the accepted drafts into the + state, ``store_kda_replay_conv_window`` puts the conv window after them in + the bf16 conv pool, and the resume seeds the replay caches from that pool + with nothing pending. The verify that follows must give the caches and + windows of the verify that replays the drafts itself bit for bit, and its + outputs and state to the kernel's accuracy: the tensor-core verify computes + a chunk in one WY-form pass, so splitting the replayed drafts and the new + tokens into two chunks rounds differently.""" + # Main's replay-only commit: the engine's cache manager (main's) runs it on suspend, + # over the replay caches the private verify kernel (cute_run) wrote. + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda._kda_kernels import ( # noqa: E501 + copy_kda_replay_conv_window, + store_kda_replay_conv_window, + ) + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( + kda_mtp_commit_pending_drafts, + ) + + data = _verify_after_round_one(B, H, num_spec, seed=41) + accepted = data["num_accepted_tokens"] + slots = data["ssm_state_indices"] + replayed = cute_run(data) + + state = data["initial_state_cute"].clone() + cs = [_dim_contiguous_clone(data[name]) for name in ("cs_q", "cs_k", "cs_v")] + kda_mtp_commit_pending_drafts( + state, + *cs, + data["k_cache"].clone(), + data["g_cache"].clone(), + data["v_cache"].clone(), + data["beta_cache"].clone(), + slots, + accepted, + num_spec, + conv_width=data["W"], + ) + for n in range(B): + unchanged = torch.equal(state[n], data["initial_state_cute"][n]) + assert unchanged == (int(accepted[n]) == 0), n + conv_pool = torch.zeros( + B, 3 * H * data["K"], data["W"] - 1, dtype=torch.bfloat16, device="cuda" + ) + store_kda_replay_conv_window(conv_pool, *cs, slots, accepted) + for cache in cs: + cache.zero_() + copy_kda_replay_conv_window(conv_pool, *cs, slots) + + resumed = dict(data) + resumed.update( + initial_state_cute=state, + cs_q=cs[0], + cs_k=cs[1], + cs_v=cs[2], + k_cache=torch.zeros_like(data["k_cache"]), + g_cache=torch.zeros_like(data["g_cache"]), + v_cache=torch.zeros_like(data["v_cache"]), + beta_cache=torch.zeros_like(data["beta_cache"]), + num_accepted_tokens=torch.zeros_like(accepted), + # The same token rows: the verify reads request n's new tokens from + # cu_seqlens[n] + num_accepted_tokens[n]. + cu_seqlens=data["cu_seqlens"] + torch.cat([accepted, accepted.new_zeros(1)]), + ) + resumed = cute_run(resumed) + + for name in ("k_cache", "g_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): + assert torch.equal(resumed[name], replayed[name]), name + _assert_close("out(commit vs replay)", resumed["out"], replayed["out"], atol=2e-2) + _assert_close( + "state(commit vs replay)", + resumed["state_v_first"], + replayed["state_v_first"], + atol=2e-4, + rtol=1e-4, + ) + + +def _replay_drafts_reference(data, slots, accepted): + """fp32 CPU replay of each selected slot's first ``accepted`` cached drafts onto its state, + the replay steps of ``cpu_reference``. Returns the K-first states of every slot.""" + H, K, V = data["H"], data["K"], data["V"] + state = data["initial_state_kfirst"].float().cpu().clone() + caches = { + name: data[name].float().cpu() for name in ("k_cache", "g_cache", "v_cache", "beta_cache") + } + for slot, count in zip(slots.tolist(), accepted.tolist()): + for h in range(H): + hk, hv = h * K, h * V + h_state = state[slot, h] + for t in range(count): + k_t = caches["k_cache"][slot, t, hk : hk + K] + decay = torch.exp(caches["g_cache"][slot, t, hk : hk + K]) + v_t = caches["v_cache"][slot, t, hv : hv + V] + beta_t = caches["beta_cache"][slot, t, h] + h_state = h_state * decay.unsqueeze(1) + v_new = (v_t - (h_state * k_t.unsqueeze(1)).sum(dim=0)) * beta_t + h_state = h_state + k_t.unsqueeze(1) * v_new.unsqueeze(0) + state[slot, h] = h_state + return state + + +@pytest.mark.parametrize( + "B,H,M", [(32, 12, 2), (32, 6, 2), (2, 12, 7), (16, 12, 7)], ids=lambda v: str(v) +) +def test_commit_pending_drafts_matches_the_replay_reference(B, H, M): + """The replay-only variant of the verify kernel commits, for each selected slot, the state + after its accepted drafts: the replay steps of a full verify, checked against the fp32 CPU + replay. It only reads the conv and replay caches, leaves a slot with nothing pending and the + unselected slots bit for bit, and covers both the cluster-pair and the one-CTA launch.""" + # Main's replay-only commit: the engine's cache manager (main's) runs it on suspend, + # over the replay caches the private verify kernel (cute_run) wrote. + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( + kda_mtp_commit_pending_drafts, + ) + + data, accept = _round2_inputs(B, H, M) + # Every other slot, in reverse order, with its own accepted count (0 included). + slots = torch.arange(B - 1, -1, -2, dtype=torch.int32) + accepted = accept[slots.long()] + expected = _replay_drafts_reference(data, slots, accepted).cuda() + + state = data["initial_state_cute"].clone() + cs = [_dim_contiguous_clone(data[name]) for name in ("cs_q", "cs_k", "cs_v")] + caches = [data[name].clone() for name in ("k_cache", "g_cache", "v_cache", "beta_cache")] + kda_mtp_commit_pending_drafts( + state, + *cs, + *caches, + slots.cuda(), + accepted.cuda(), + M, + conv_width=data["W"], + ) + + committed = state.permute(0, 1, 3, 2) + _assert_close( + "state(commit vs cpu replay)", committed[slots.long()], expected[slots.long()], atol=1e-4 + ) + for slot, count in zip(slots.tolist(), accepted.tolist()): + if count == 0: + assert torch.equal(state[slot], data["initial_state_cute"][slot]), slot + others = [slot for slot in range(B) if slot not in slots.tolist()] + assert torch.equal(state[others], data["initial_state_cute"][others]) + for name, cache in zip(("cs_q", "cs_k", "cs_v"), cs): + assert torch.equal(cache, data[name]), name + for name, cache in zip(("k_cache", "g_cache", "v_cache", "beta_cache"), caches): + assert torch.equal(cache, data[name]), name + + +def test_commit_pending_drafts_replays_the_per_token_records(): + """With the Kimi K3 verify kernels' draft records, a slot's accepted drafts + are replayed onto its pool state, in order, from its records (vn, + beta * k, decay), and the conv window after them is stored; other slots + are left alone.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda._kda_kernels import ( # noqa: E501 + commit_kda_pending_drafts, + ) + + torch.manual_seed(5) + num_slots, H, K, W, num_spec = 9, 6, 128, 4, 7 + + def small(*shape, scale): + # Few significant bits: every product and sum below is exact in fp32, + # so the reference does not depend on fused or separate rounding. + return torch.randint(-4, 5, shape, device="cuda").float() * scale + + state = small(num_slots, H, K, K, scale=2.0**-4) + vn = small(num_slots, num_spec, H, K, scale=2.0**-2) + beta_k = small(num_slots, num_spec, H, K, scale=2.0**-3) + decay = torch.randint(1, 3, (num_slots, num_spec, H, K), device="cuda").float() * 0.5 + state_tok = torch.zeros(num_slots, 3, num_spec, H, K, device="cuda") + records = state_tok.view(num_slots, -1) + size = num_spec * H * K + for i, part in enumerate((vn, beta_k, decay)): + records[:, i * size : (i + 1) * size] = part.reshape(num_slots, -1) + cs = [ + torch.randn(num_slots, W - 1 + num_spec, H * K, device="cuda") + .bfloat16() + .float() + .transpose(1, 2) + for _ in range(3) + ] + conv_pool = torch.zeros(num_slots, 3 * H * K, W - 1, dtype=torch.bfloat16, device="cuda") + slots = torch.tensor([7, 0, 3, 5, 1, 8], dtype=torch.int32, device="cuda") + accepted = torch.tensor([0, 1, 3, 7, 2, 5], dtype=torch.int32, device="cuda") + before = state.clone() + unread = torch.empty(0, device="cuda") + + commit_kda_pending_drafts( + state, + conv_pool, + *cs, + # k / g / v / beta replay caches: unread with per-token states. + unread, + unread, + unread, + unread, + slots, + accepted, + state_tok=state_tok, + ) + + for slot, count in zip(slots.tolist(), accepted.tolist()): + expected = before[slot].clone() + for t in range(count): + expected = ( + decay[slot, t][:, None, :] * expected + + vn[slot, t][:, :, None] * beta_k[slot, t][:, None, :] + ) + assert torch.equal(state[slot], expected), slot + for section, cache in enumerate(cs): + window = conv_pool[slot, section * H * K : (section + 1) * H * K].float() + assert torch.equal(window, cache[slot, :, count : count + W - 1]), (slot, section) + others = [slot for slot in range(num_slots) if slot not in slots.tolist()] + assert torch.equal(state[others], before[others]) + assert not conv_pool[others].any() + + +def _packed_fused_norm_inputs(B, H, num_spec, state_dtype): + """Runtime-contract inputs for the packed / fused-norm path: section-strided q / k / v / + output-gate views of one fused projection, the bf16 beta tail of [f_a | beta], replay caches + and conv windows chained from a CPU-golden round 1, mixed acceptance 0..num_spec and + non-identity state slots.""" + data, _ = _round2_inputs(B, H, num_spec) + packed, _ = _pack_new_token_projection_views(data) + slots = torch.arange(B - 1, -1, -1, dtype=torch.int32, device="cuda") + packed["ssm_state_indices"] = slots + for name in ( + "initial_state_cute", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + slot_order = torch.empty_like(packed[name]) + slot_order[slots.long()] = packed[name] + packed[name] = slot_order + pool_accepted = torch.zeros(B, dtype=torch.int32, device="cuda") + pool_accepted[slots.long()] = data["num_accepted_tokens"] + packed["num_accepted_tokens"] = pool_accepted + packed["initial_state_cute"] = packed["initial_state_cute"].to(state_dtype) + packed["onorm_weight"] = torch.linspace( + 0.75, 1.25, data["V"], dtype=torch.float32, device="cuda" + ) + packed["onorm_eps"] = 1e-5 + return packed + + +@pytest.mark.parametrize( + "B,H,num_spec", + [(1, 6, 6), (4, 6, 7), (16, 6, 6), (32, 6, 7), (128, 6, 6), (8, 12, 7)], + ids=lambda v: str(v), +) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +@pytest.mark.parametrize("strided_state", [False, True], ids=["pool", "strided"]) +def test_v2_kernel_matches_v1_bitwise(B, H, num_spec, state_dtype, strided_state, monkeypatch): + """The latency-scheduled verify kernel (``mtp_ops._MTP_V2``, packed rows + fused gated + RMSNorm) performs the same arithmetic as ``kda_decode_mtp_kernel`` in the same order: the + outputs, committed state, conv pools and replay caches must match bit for bit. Covers the + cluster-pair (HV * B <= 64) and single-CTA grids, fp32 / bf16 state pools and a state that + is a strided slice of a larger buffer (non-flat TMA view).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda import ( # noqa: E501 + mtp_ops, + ) + + packed = _packed_fused_norm_inputs(B, H, num_spec, state_dtype) + + def run(v2): + monkeypatch.setattr(mtp_ops, "_MTP_V2", v2) + override = None + if strided_state: + shape = tuple(packed["initial_state_cute"].shape) + big = torch.zeros((shape[0], 3) + shape[1:], dtype=state_dtype, device="cuda") + override = big[:, 1] + return cute_run( + packed, packed_token_layout=True, fuse_output_norm=True, state_override=override + ) + + ref = run(v2=False) + operand_checks = [] + operands_ok = mtp_ops._v2_operands_ok + + def record_operand_check(**kwargs): + operand_checks.append(operands_ok(**kwargs)) + return operand_checks[-1] + + monkeypatch.setattr(mtp_ops, "_v2_operands_ok", record_operand_check) + new = run(v2=True) + assert operand_checks == [True], "the v2 kernel did not take the call" + for name in ( + "out", + "state_v_first", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + assert torch.equal(new[name], ref[name]), f"{name}: v2 differs from v1" + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-x"])) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_mx_pdl.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_mx_pdl.py new file mode 100644 index 000000000000..40f0a6d36ef4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_mtp_decode_mx_pdl.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8-output KDA verify as a PDL launch with the early trigger (``TRTLLM_K3_KDA_MX_PDL``). + +The c240 o_proj push runs the MXFP8-output verify on the two-CTAs-per-SM kernel as a PDL launch +that releases its dependent right after ``griddepcontrol.wait``. Scheduling only: the E4M3 +values, the UE8M0 scales and every mutated buffer (recurrent state, conv windows, replay caches) +must equal those of the plain launch over two chained verify rounds with mixed per-slot accepted +counts, eagerly and under CUDA graph replay. + +Inputs follow the runtime contract: packed token layout (strided QKVG section views), bf16 state +pool (V-first, strided per-layer slice), bf16 conv / k / v / beta replay caches, fp32 gate cache, +fused gated RMSNorm with the MXFP8 operand of the Rubin o_proj. The two-CTA MXFP8 kernel runs by +default on SM107 only; the tests select it on SM100 / SM103 as well. + +Requires 1 SM100/SM103/SM107 GPU and nvidia-cutlass-dsl. +""" + +import pytest +import torch + +_HAVE_DEPS = True +_DEP_ERR = None +try: + import cuda.bindings.driver # noqa: F401 + import cutlass # noqa: F401 + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda import ( # noqa: E501 + mtp_ops as kda_ops, + ) +except ImportError as e: + _HAVE_DEPS = False + _DEP_ERR = str(e) + + +def _is_sm100_family(): + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability(0) in ((10, 0), (10, 3), (10, 7)) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), + pytest.mark.skipif(not _is_sm100_family(), reason="needs sm100/sm103/sm107"), + pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), +] + +K = V = 128 +W = 4 +POOL = 80 +MUTATED = ("cs_q", "cs_k", "cs_v", "recurrent_state", "k_cache", "g_cache", "v_cache", "beta_cache") + + +def _make_inputs(N, H, num_spec, seed, zero_state): + gen = torch.Generator().manual_seed(seed) + dev = "cuda" + d = H * K + T = N * (1 + num_spec) + S = W - 1 + num_spec + + def rn(*shape, dtype=torch.float32, scale=1.0): + return (torch.randn(*shape, generator=gen) * scale).to(dtype).to(dev) + + def conv(): + # [layers, slots, window, dim] -> dim-contiguous view of layer 1 + return rn(2, POOL, S, d, dtype=torch.bfloat16).transpose(-1, -2)[1] + + # fused [q | k | v | output-gate] projection: strided section views + qkvg = rn(T, 4 * d, dtype=torch.bfloat16, scale=0.5) + # [f_a | b | pad] projection rows (beta is a strided view) + bfa = rn(T, 128 + (H + 7) // 8 * 8, dtype=torch.bfloat16) + state = rn(2, POOL, H, V, K, dtype=torch.bfloat16, scale=0.05) + if zero_state: + state.zero_() + slots = torch.randperm(POOL, generator=gen)[:N].to(torch.int32) + accepted = torch.zeros(POOL, dtype=torch.int32) + if not zero_state: + accepted[slots.long()] = torch.randint( + 0, num_spec + 1, (N,), generator=gen, dtype=torch.int32 + ) + heads_padded = (H + 7) // 8 * 8 + return dict( + x_q=qkvg[:, 0 * d : 1 * d].view(1, T, H, K), + x_k=qkvg[:, 1 * d : 2 * d].view(1, T, H, K), + x_v=qkvg[:, 2 * d : 3 * d].view(1, T, H, V), + onorm_g=qkvg[:, 3 * d : 4 * d].view(1, T, H, V), + g=rn(T, d, dtype=torch.bfloat16, scale=0.5).view(1, T, H, K), + beta=bfa[:, 128 : 128 + H].view(1, T, H), + w_q=rn(d, W, scale=0.5), + w_k=rn(d, W, scale=0.5), + w_v=rn(d, W, scale=0.5), + A_log=rn(H, scale=0.5), + dt_bias=rn(d, scale=0.5), + cs_q=conv(), + cs_k=conv(), + cs_v=conv(), + recurrent_state=state[1], + k_cache=rn(2, POOL, num_spec, d, dtype=torch.bfloat16, scale=0.08)[1], + g_cache=(-torch.rand(2, POOL, num_spec, d, generator=gen) * 0.5).to(dev)[1], + v_cache=rn(2, POOL, num_spec, d, dtype=torch.bfloat16)[1], + beta_cache=torch.rand(2, POOL, num_spec, heads_padded, generator=gen) + .to(torch.bfloat16) + .to(dev)[1, ..., :H], + ssm_state_indices=slots.to(dev), + cu_seqlens=None, + num_spec=num_spec, + num_accepted_tokens=accepted.to(dev), + lower_bound=-5.0, + scale=K**-0.5, + packed_token_layout=True, + onorm_weight=(0.75 + 0.5 * torch.rand(V, generator=gen)).to(dev), + onorm_eps=1e-5, + fuse_output_norm=True, + ) + + +def _clone(inp): + out = dict(inp) + for name in MUTATED: + t = inp[name] + base = t._base if t._base is not None else t + out[name] = torch.as_strided(base.clone(), t.shape, t.stride(), t.storage_offset()) + out["num_accepted_tokens"] = inp["num_accepted_tokens"].clone() + return out + + +def _call_kwargs(x, use_pdl): + return {**x, "use_pdl": use_pdl} + + +def _next_round(x, rnd): + """The sampler accepted a different number of drafts per slot.""" + acc = x["num_accepted_tokens"] + slots = x["ssm_state_indices"].long() + acc[slots] = ( + acc[slots] + 1 + rnd + torch.arange(len(slots), device=acc.device, dtype=acc.dtype) % 3 + ) % (x["num_spec"] + 1) + + +def _scale_bytes(scales, rows, heads): + """The scale bytes of rows [0, rows) (four copies per (row, head)); the padding rows of the + R128c4 buffer are never written.""" + r = torch.arange(rows, device=scales.device)[:, None] + h = torch.arange(heads, device=scales.device)[None, :] + off = (r // 128) * heads * 512 + h * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + off = (off[..., None] + torch.arange(4, device=scales.device)).reshape(-1) + return scales.view(torch.uint8)[off] + + +@pytest.mark.parametrize("H,N,num_spec", [(12, 16, 5), (12, 12, 7), (12, 4, 7)]) +def test_mxfp8_pdl_early_trigger_bitwise(H, N, num_spec, monkeypatch): + """c240 TRTLLM_K3_KDA_MX_PDL: the MXFP8-output verify as a PDL launch with the early trigger + (two-CTA kernel) gives the plain launch's E4M3 values, scales and mutated buffers.""" + monkeypatch.setattr(kda_ops, "_mxfp8_two_cta_enabled", lambda device: True) + inp = _make_inputs(N, H, num_spec, seed=3 * N + H, zero_state=False) + rows = N * (1 + num_spec) + + def run(mx_pdl): + x = _clone(inp) + outs = [] + keys_before = set(kda_ops._compiled_cache) + for rnd in range(2): + with kda_ops.mxfp8_pdl_early_trigger(mx_pdl): + act, sf = kda_ops.kda_mtp_decode_impl( + **_call_kwargs(x, mx_pdl), quantize_output=True + ) + outs.append((act.view(torch.uint8).clone(), _scale_bytes(sf, rows, H).clone())) + _next_round(x, rnd) + torch.cuda.synchronize() + new_keys = [k for k in kda_ops._compiled_cache if k not in keys_before] + return outs, x, new_keys + + ref_outs, ref_x, _ = run(False) + new_outs, new_x, keys = run(True) + trig = [k for k in kda_ops._compiled_cache if k[0] == "2cta" and k[-1] is True] + assert trig, "the early-trigger two-CTA variant did not run" + for rnd, ((a, sa), (b, sb)) in enumerate(zip(new_outs, ref_outs)): + assert torch.equal(a, b), f"round {rnd} E4M3 values" + assert torch.equal(sa, sb), f"round {rnd} UE8M0 scales" + for name in MUTATED: + assert torch.equal(new_x[name], ref_x[name]), name + + +def test_mxfp8_pdl_early_trigger_cuda_graph(monkeypatch): + """The early-trigger launch in a CUDA graph, replayed, matches eager plain launches.""" + monkeypatch.setattr(kda_ops, "_mxfp8_two_cta_enabled", lambda device: True) + H, N, num_spec = 12, 16, 5 + inp = _make_inputs(N, H, num_spec, seed=99, zero_state=False) + rows = N * (1 + num_spec) + x_eager = _clone(inp) + act, sf = kda_ops.kda_mtp_decode_impl(**_call_kwargs(x_eager, False), quantize_output=True) + ref = (act.view(torch.uint8).clone(), _scale_bytes(sf, rows, H).clone()) + x_graph = _clone(inp) + snap = _clone(x_graph) + with kda_ops.mxfp8_pdl_early_trigger(True): + kda_ops.kda_mtp_decode_impl(**_call_kwargs(_clone(inp), True), quantize_output=True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_act, g_sf = kda_ops.kda_mtp_decode_impl( + **_call_kwargs(x_graph, True), quantize_output=True + ) + for name in MUTATED: + x_graph[name].copy_(snap[name]) + x_graph["num_accepted_tokens"].copy_(snap["num_accepted_tokens"]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(g_act.view(torch.uint8), ref[0]) + assert torch.equal(_scale_bytes(g_sf, rows, H), ref[1]) + for name in MUTATED: + assert torch.equal(x_graph[name], x_eager[name]), name diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_prefill_op.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_prefill_op.py new file mode 100644 index 000000000000..038481393a4b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kda_prefill_op.py @@ -0,0 +1,1364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parity tests for the Kimi K3 KDA prefill op and the layer's prefill chain. + +The layer runs on FP8 128x128 block-scale q/k/v/g/o weights (K3's checkpoint format), finalized +as the model loader does. The FLA reference (``KimiKDAReference``) shares the layer's FP8 +projections, so a comparison measures the convolution, recurrence and gated RMSNorm. +""" + +import importlib.util +import math +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("fla") + +from fla.modules import ShortConvolution # noqa: E402 + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + kda as kimi_kda_mixer, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.fp8_weight_read import ( # noqa: E402, E501 + _Fp8BlockScaleWeightReadLinear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kda import ( + K3KDALinearAttention, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda._kda_kernels import ( # noqa: E402, E501 + copy_kda_replay_conv_window, + fused_kda_post_conv, + is_kda_optimized_supported, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.layers import ( # noqa: E402, E501 # noqa: E402, E501 + Bf16Linear, + Fp8BlockScaleLinear, + ShortConvWeight, +) +from tensorrt_llm._torch.model_config import ModelConfig # noqa: E402 +from tensorrt_llm._torch.modules.mamba.causal_conv1d import causal_conv1d_fn # noqa: E402 +from tensorrt_llm.models.modeling_utils import QuantConfig # noqa: E402 +from tensorrt_llm.quantization import QuantAlgo # noqa: E402 +from tests.unittest._torch.modules.kimi_kda.kimi_kda_test_utils import KimiKDAReference + +NUM_HEADS = 96 +HEAD_DIM = 128 +CONV_KERNEL_SIZE = 4 +HIDDEN_SIZE = 7168 + + +def _has_supported_gpu() -> bool: + if not torch.cuda.is_available(): + return False + return is_kda_optimized_supported() + + +pytestmark = pytest.mark.skipif( + not _has_supported_gpu(), + reason="FlashInfer KDA prefill requires SM100/SM103/SM107", +) + + +def _kda_prefill_entry(): + """The ``ops/ssm/kda_prefill`` entry (its import needs CuTe DSL and FlashInfer).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.kda_prefill import ( # noqa: E501 + kda_prefill, + ) + + return kda_prefill + + +@torch.no_grad() +def _make_kda(*, num_heads: int = NUM_HEADS, seed: int = 0) -> K3KDALinearAttention: + """A layer on seeded FP8 checkpoint-like weights, finalized as the model loader does: the + fused QKVG and o_proj swapped to the K3 FP8 weight-read GEMM.""" + cfg = SimpleNamespace( + hidden_size=HIDDEN_SIZE, + rms_norm_eps=1e-5, + linear_attn_config={ + "num_heads": num_heads, + "head_dim": HEAD_DIM, + "short_conv_kernel_size": CONV_KERNEL_SIZE, + "use_full_rank_gate": True, + "gate_lower_bound": -5.0, + }, + ) + model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES)) + kda = K3KDALinearAttention(cfg, layer_idx=0, model_config=model_config).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(seed) + + def randn(shape, scale): + return torch.randn(shape, generator=generator, device="cuda") * scale + + for module in kda.modules(): + if isinstance(module, Fp8BlockScaleLinear): + codes = torch.randint(-32, 33, module.weight.shape, generator=generator, device="cuda") + module.weight.copy_(codes.to(module.weight.dtype)) + module.weight_scale.fill_(1 / (32 * math.sqrt(module.in_features))) + elif isinstance(module, Bf16Linear): + module.weight.copy_(randn(module.weight.shape, (3 * module.in_features) ** -0.5)) + elif isinstance(module, ShortConvWeight): + module.weight.copy_(randn(module.weight.shape, 0.3)) + kda.o_norm.weight.copy_(1 + randn(kda.o_norm.weight.shape, 0.1)) + kda.dt_bias.zero_() + kda.finalize_decode_weights() + kda._build_mtp_conv_weights() + for name in ("qkvg_proj", "o_proj"): + module = _Fp8BlockScaleWeightReadLinear.from_linear(getattr(kda, name)) + setattr(kda, name, module) + module.post_load_weights() + return kda + + +class _Columns(torch.nn.Module): + """Output columns ``[start, end)`` of a projection.""" + + def __init__(self, projection: torch.nn.Module, start: int, end: int) -> None: + super().__init__() + self.projection, self.start, self.end = projection, start, end + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.projection(x)[..., self.start : self.end] + + +@torch.no_grad() +def _make_reference(kda: K3KDALinearAttention) -> KimiKDAReference: + """The FLA reference on ``kda``'s weights: its q / k / v / g projections are the columns of the + layer's fused FP8 QKVG projection and its o_proj the layer's, the rest copies.""" + reference = KimiKDAReference( + hidden_size=HIDDEN_SIZE, + num_heads=kda.num_heads, + head_dim=HEAD_DIM, + conv_kernel_size=CONV_KERNEL_SIZE, + use_full_rank_gate=True, + gate_lower_bound=-5.0, + rms_norm_eps=1e-5, + dtype=torch.bfloat16, + ).to("cuda") + for name, parameter in reference.named_parameters(): + if not name.startswith(("q_proj.", "k_proj.", "v_proj.", "g_proj.", "o_proj.")): + parameter.copy_(kda.get_parameter(name)) + d = kda.proj_size + for index, name in enumerate(("q_proj", "k_proj", "v_proj", "g_proj")): + setattr(reference, name, _Columns(kda.qkvg_proj, index * d, (index + 1) * d)) + reference.o_proj = kda.o_proj + return reference + + +def _run_production_prefill( + attention: K3KDALinearAttention, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + *, + conv_pool: torch.Tensor | None = None, + state_pool: torch.Tensor | None = None, + slot_indices: torch.Tensor | None = None, + has_initial_states: torch.Tensor | None = None, + kda_snapshot=None, + inforward_stage=None, + single_sequence_length: int | None = None, + project_output: bool = True, +) -> torch.Tensor: + """The prefill output ``[batch, tokens, hidden]``, or with ``project_output=False`` the + normalized core ``[tokens, heads, head_dim]`` before o_proj.""" + batch_size, sequence_length, _ = hidden_states.shape + if cu_seqlens is None: + cu_seqlens = torch.arange(batch_size + 1, device="cuda", dtype=torch.long) * sequence_length + else: + batch_size = cu_seqlens.numel() - 1 + + projection_size = attention.proj_size + if conv_pool is None: + conv_pool = torch.zeros( + batch_size, + 3 * projection_size, + CONV_KERNEL_SIZE - 1, + dtype=torch.bfloat16, + device="cuda", + ) + if state_pool is None: + state_pool = torch.zeros( + batch_size, + attention.num_heads, + HEAD_DIM, + HEAD_DIM, + dtype=torch.float32, + device="cuda", + ) + if slot_indices is None: + slot_indices = torch.arange(batch_size, device="cuda", dtype=torch.long) + use_initial_states = has_initial_states is not None + if has_initial_states is None: + has_initial_states = torch.zeros(batch_size, device="cuda", dtype=torch.bool) + metadata = SimpleNamespace( + use_initial_states=use_initial_states, + has_initial_states=has_initial_states, + state_indices=slot_indices.to(torch.int32), + # causal_conv1d wants int32 offsets, unlike the int64 cu_seqlens above. + query_start_loc=cu_seqlens.to(torch.int32), + # Row indices of the dense fp32 copy a bf16 state pool is staged through. + _arange_buffer=torch.arange(batch_size + 1, dtype=torch.int32, device="cuda"), + kda_snapshot=kda_snapshot, + kda_single_sequence_length=single_sequence_length, + ) + core = attention.forward_prefill( + hidden_states.reshape(-1, HIDDEN_SIZE), + cu_seqlens.to(torch.int64), + metadata, + batch_size, + conv_pool, + state_pool, + slot_indices, + inforward_stage=inforward_stage, + ) + if not project_output: + return core + output = attention._project_output(core) + return output.reshape_as(hidden_states) + + +def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + actual_float = actual.float() + expected_float = expected.float() + cosine = torch.nn.functional.cosine_similarity( + actual_float.flatten(), expected_float.flatten(), dim=0 + ).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + assert cosine > 0.999 + assert relative_l2 < 3e-2 + + +@pytest.mark.parametrize("token_major", [False, True], ids=["channel_major", "token_major"]) +@pytest.mark.parametrize("sequence_length", [1, 127, 256]) +@torch.no_grad() +def test_fused_kda_post_conv_matches_reference(sequence_length: int, token_major: bool) -> None: + """The unfused convolution path produces contiguous raw Q/K/V without normalization.""" + torch.manual_seed(1) + num_heads, head_dim = 8, 128 + projection_size = num_heads * head_dim + shape = (3 * projection_size, sequence_length) + if token_major: + shape = shape[::-1] + packed = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + if token_major: + packed = packed.t() + + actual_q, actual_k, actual_v = fused_kda_post_conv( + packed, num_heads, head_dim, normalize_qk=False + ) + unpacked = packed.reshape(3, projection_size, sequence_length) + expected_q = unpacked[0].transpose(0, 1).reshape(1, sequence_length, num_heads, head_dim) + expected_k = unpacked[1].transpose(0, 1).reshape(1, sequence_length, num_heads, head_dim) + expected_v = unpacked[2].transpose(0, 1).reshape(1, sequence_length, num_heads, head_dim) + assert actual_q.is_contiguous() + assert actual_k.is_contiguous() + assert actual_v.is_contiguous() + torch.testing.assert_close(actual_q, expected_q, rtol=0, atol=0) + torch.testing.assert_close(actual_k, expected_k, rtol=0, atol=0) + assert torch.equal(actual_v, expected_v) + + +@torch.no_grad() +def test_fused_kda_post_conv_rejects_layout_without_unit_stride() -> None: + """Input with no unit stride along tokens or features is rejected.""" + projection_size = 2 * 128 + packed = torch.randn(3 * projection_size, 8, 2, dtype=torch.bfloat16, device="cuda")[..., 0] + with pytest.raises(ValueError, match="unit stride"): + fused_kda_post_conv(packed, 2, 128) + + +@torch.no_grad() +@pytest.mark.parametrize("num_tokens", [256, 301]) +@pytest.mark.parametrize("token_major", [False, True], ids=["channel_major", "token_major"]) +def test_project_packed_conv_input_layout(num_tokens: int, token_major: bool) -> None: + """The convolution input is the QKVG GEMM's [q | k | v] columns: a [T, 3D] view of its output + for fused_kda_conv_prefill, or repacked contiguous channel-major [3D, T] for + causal_conv1d_fn; the output gate is the view of the remaining columns.""" + torch.manual_seed(3) + # Only the projection layout is checked, so two heads keep the weights small. + kda = _make_kda(num_heads=2) + d = kda.proj_size + x = torch.randn(1, num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + + packed_conv, onorm_g = kda._project_packed_conv_input(x, token_major=token_major) + qkvg = kda.qkvg_proj(x).squeeze(0) + + if token_major: + assert packed_conv.shape == (num_tokens, 3 * d) and packed_conv.stride(1) == 1 + assert torch.equal(packed_conv, qkvg[:, : 3 * d]) + else: + assert packed_conv.shape == (3 * d, num_tokens) and packed_conv.is_contiguous() + assert torch.equal(packed_conv, qkvg[:, : 3 * d].t()) + assert torch.equal(onorm_g.squeeze(0), qkvg[:, 3 * d :]) + + +@torch.no_grad() +def test_prefill_output_matches_across_projection_layouts() -> None: + """Aligned and unaligned token counts produce the same prefill output. + + The odd-length request is the aligned one plus one trailing token, so its + shared prefix must match the aligned request token for token. + """ + torch.manual_seed(4) + kda = _make_kda() + hidden_states = torch.randn(1, 257, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + + aligned = _run_production_prefill(kda, hidden_states[:, :256].contiguous()) + unaligned = _run_production_prefill(kda, hidden_states) + + torch.testing.assert_close(unaligned[:, :256], aligned, rtol=2e-2, atol=2e-2) + + +@pytest.mark.parametrize("sequence_length", [1, 63, 65, 193]) +@pytest.mark.parametrize("beta_layout", ["contiguous", "fused_projection"]) +@torch.no_grad() +def test_fused_kda_post_conv_prepad_matches_separate_padding( + sequence_length: int, beta_layout: str +) -> None: + """The producer writes the same Q/K/V/beta prefix and exact zero tail as F.pad.""" + torch.manual_seed(1) + heads, head_dim = 6, 128 + packed = torch.randn(3 * heads * head_dim, sequence_length, device="cuda", dtype=torch.bfloat16) + if beta_layout == "fused_projection": + projection = torch.randn( + 1, sequence_length, head_dim + heads, device="cuda", dtype=torch.bfloat16 + ) + beta = projection[..., head_dim:] + assert beta.stride(-2) > heads + else: + beta = torch.randn(1, sequence_length, heads, device="cuda", dtype=torch.bfloat16) + + padded_tokens = ((sequence_length + 63) // 64) * 64 + baseline_qkv = fused_kda_post_conv(packed, heads, head_dim) + actual = fused_kda_post_conv(packed, heads, head_dim, padded_tokens=padded_tokens, beta=beta) + expected = tuple( + torch.nn.functional.pad(value, (0, 0, 0, 0, 0, padded_tokens - sequence_length)) + for value in baseline_qkv + ) + (torch.nn.functional.pad(beta.float(), (0, 0, 0, padded_tokens - sequence_length)),) + + for name, got, want in zip(("q", "k", "v", "beta"), actual, expected, strict=True): + context = f"sequence_length={sequence_length} beta_layout={beta_layout} tensor={name}" + assert got.shape == want.shape and got.is_contiguous(), context + got_bits = got.view(torch.int16 if got.element_size() == 2 else torch.int32) + want_bits = want.view(torch.int16 if want.element_size() == 2 else torch.int32) + assert torch.equal(got_bits, want_bits), context + + +@pytest.mark.parametrize("invalid", ["shape", "head_stride", "device"]) +def test_fused_kda_post_conv_rejects_invalid_beta(invalid: str) -> None: + """Reject inputs that would silently copy beta or read from another device.""" + packed = torch.empty(3 * 6 * 128, 65, device="cuda", dtype=torch.bfloat16) + beta = torch.empty(1, 65, 6, device="cuda", dtype=torch.bfloat16) + if invalid == "shape": + beta = beta[:, :-1] + elif invalid == "head_stride": + beta = torch.empty(1, 65, 12, device="cuda", dtype=torch.bfloat16)[..., ::2] + else: + beta = beta.cpu() + with pytest.raises(ValueError, match="beta must be"): + fused_kda_post_conv(packed, 6, 128, padded_tokens=128, beta=beta) + + +@torch.no_grad() +def test_packed_gdn_convolution_matches_three_fla_convolutions() -> None: + """The reused GDN kernel preserves KDA convolution and cache semantics.""" + torch.manual_seed(2) + dim, width, slots = 256, 4, 5 + sequence_lengths = [7, 5] + num_tokens = sum(sequence_lengths) + state_indices = torch.tensor([4, 1], dtype=torch.int32, device="cuda") + state_indices_long = state_indices.long() + has_initial_state = torch.tensor([True, False], device="cuda") + cu_seqlens = torch.tensor( + [0, sequence_lengths[0], num_tokens], dtype=torch.int32, device="cuda" + ) + cu_seqlens_long = cu_seqlens.long() + + projected = [ + torch.randn(1, num_tokens, dim, dtype=torch.bfloat16, device="cuda") for _ in range(3) + ] + convolutions = [ + ShortConvolution(dim, width, activation="silu").to(device="cuda", dtype=torch.bfloat16) + for _ in range(3) + ] + packed_weight = torch.cat( + [conv.weight.detach().squeeze(1) for conv in convolutions], dim=0 + ).contiguous() + conv_pool = torch.randn(slots, 3 * dim, width - 1, dtype=torch.bfloat16, device="cuda") + conv_pool_before = conv_pool.clone() + + packed = torch.cat(projected, dim=-1).squeeze(0).transpose(0, 1).contiguous() + causal_conv1d_fn( + packed, + packed_weight, + query_start_loc=cu_seqlens, + cache_indices=state_indices, + has_initial_state=has_initial_state, + conv_states=conv_pool, + activation="silu", + ) + + selected = conv_pool_before.index_select(0, state_indices_long) + selected[~has_initial_state] = 0 + expected_outputs = [] + expected_caches = [] + for section, (projected_section, convolution) in enumerate( + zip(projected, convolutions, strict=True) + ): + cache = torch.nn.functional.pad(selected[:, section * dim : (section + 1) * dim], (1, 0)) + output, final_cache = convolution( + projected_section, + cache=cache, + output_final_state=True, + cu_seqlens=cu_seqlens_long, + ) + expected_outputs.append(output) + expected_caches.append(final_cache[:, :, 1:]) + + actual_output = packed.transpose(0, 1).reshape(1, num_tokens, 3 * dim) + torch.testing.assert_close( + actual_output, + torch.cat(expected_outputs, dim=-1), + rtol=1e-2, + atol=1e-2, + ) + torch.testing.assert_close( + conv_pool.index_select(0, state_indices_long), + torch.cat(expected_caches, dim=1), + rtol=0, + atol=0, + ) + untouched = torch.ones(slots, dtype=torch.bool, device="cuda") + untouched[state_indices_long] = False + assert torch.equal(conv_pool[untouched], conv_pool_before[untouched]) + + +@torch.no_grad() +def test_copy_kda_replay_conv_window_preserves_slot_padding() -> None: + """Committed Q/K/V rows copy across both strided production layouts.""" + slots, dim, committed, num_spec = 6, 11, 3, 2 + slot_stride = 3 * dim * committed + 17 + storage = torch.arange(slots * slot_stride, dtype=torch.float32, device="cuda").to( + torch.bfloat16 + ) + conv_pool = torch.as_strided( + storage, + size=(slots, 3 * dim, committed), + stride=(slot_stride, committed, 1), + ) + + def replay_cache() -> torch.Tensor: + return torch.full( + (slots, committed + num_spec, dim), + -1, + dtype=torch.float32, + device="cuda", + ).transpose(-1, -2) + + q_cache, k_cache, v_cache = replay_cache(), replay_cache(), replay_cache() + state_indices = torch.tensor([4, 1], dtype=torch.int32, device="cuda") + copy_kda_replay_conv_window( + conv_pool, + q_cache, + k_cache, + v_cache, + state_indices, + ) + + state_indices_long = state_indices.long() + for section, cache in enumerate((q_cache, k_cache, v_cache)): + expected = conv_pool.index_select(0, state_indices_long)[ + :, section * dim : (section + 1) * dim + ].float() + torch.testing.assert_close( + cache.index_select(0, state_indices_long)[:, :, :committed], + expected, + ) + assert torch.equal( + cache.index_select(0, state_indices_long)[:, :, committed:], + torch.full( + (state_indices.numel(), dim, num_spec), + -1, + dtype=cache.dtype, + device=cache.device, + ), + ) + untouched = torch.ones(slots, dtype=torch.bool, device="cuda") + untouched[state_indices_long] = False + assert torch.equal(q_cache[untouched], replay_cache()[untouched]) + assert torch.equal(k_cache[untouched], replay_cache()[untouched]) + assert torch.equal(v_cache[untouched], replay_cache()[untouched]) + + +@torch.no_grad() +@pytest.mark.parametrize("num_heads", [NUM_HEADS, 12, 6], ids=["h96", "h12", "h6"]) +def test_optimized_prefill_matches_fla_reference(num_heads: int) -> None: + torch.manual_seed(0) + optimized = _make_kda(num_heads=num_heads) + reference = _make_reference(optimized) + + # Reuse the same module across changing token counts and batch sizes. + for batch_size, sequence_length in [(2, 256), (2, 512), (1, 1024)]: + hidden_states = ( + torch.randn( + batch_size, + sequence_length, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = _run_production_prefill(optimized, hidden_states) + expected = reference.forward_prefill(hidden_states) + _assert_close(actual, expected) + + hidden_states = torch.randn(1, 300, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + actual = _run_production_prefill(optimized, hidden_states) + expected = reference.forward_prefill(hidden_states) + _assert_close(actual, expected) + + sequence_lengths = [128, 256, 192] + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + + +@torch.no_grad() +@pytest.mark.parametrize( + "sequence_lengths", + [None, [100, 257, 63]], + ids=["padded_eqlen", "varlen"], +) +def test_fused_prefill_beta_sigmoid_matches_unfused_kernel( + sequence_lengths: list[int] | None, +) -> None: + """In-kernel beta sigmoid matches the legacy host launch.""" + from fla.modules.l2norm import l2norm_fwd + from fla.ops.common.gate import fused_beta_sigmoid + + torch.manual_seed(10) + batch_size = 1 + if sequence_lengths is None: + sequence_length = 300 + num_sequences = batch_size + cu_seqlens = None + else: + sequence_length = sum(sequence_lengths) + num_sequences = len(sequence_lengths) + cu_seqlens = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + A_log = torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.1 + dt_bias = torch.randn(NUM_HEADS * HEAD_DIM, dtype=torch.float32, device="cuda") * 0.1 + + for seed in (11, 12): + generator = torch.Generator(device="cuda").manual_seed(seed) + + def randn( + *shape: int, dtype: torch.dtype = torch.bfloat16, scale: float = 0.05 + ) -> torch.Tensor: + return ( + torch.randn(*shape, generator=generator, dtype=torch.float32, device="cuda") * scale + ).to(dtype) + + q = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM) + k = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM) + v = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM) + g = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM) + if seed == 11: + beta = randn(batch_size, sequence_length, NUM_HEADS, dtype=torch.float32) + else: + # Exercise the L2 epsilon path and saturated sigmoid tails. + q[:, ::17].zero_() + k[:, ::19].zero_() + beta = torch.linspace( + -20.0, + 20.0, + batch_size * sequence_length * NUM_HEADS, + dtype=torch.float32, + device="cuda", + ).reshape(batch_size, sequence_length, NUM_HEADS) + initial_state = randn( + num_sequences, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=torch.float32, scale=0.01 + ) + state_indices = torch.arange(num_sequences, dtype=torch.int32, device="cuda") + + common = dict( + v=v, + g=g, + scale=HEAD_DIM**-0.5, + state_indices=state_indices, + safe_gate=True, + lower_bound=-5.0, + use_gate_in_kernel=True, + A_log=A_log, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + ) + + q_unfused, _ = l2norm_fwd(q) + k_unfused, _ = l2norm_fwd(k) + beta_unfused = fused_beta_sigmoid(beta, scale=1.0).to(torch.bfloat16) + state_unfused = initial_state.clone() + output_unfused = _kda_prefill_entry()( + q=q_unfused, + k=k_unfused, + beta=beta_unfused, + state_pool=state_unfused, + use_beta_sigmoid_in_kernel=False, + **common, + ) + # The op reuses shape-keyed output buffers; retain the first result + # before the fused call overwrites it. + output_unfused = output_unfused.clone() + state_unfused = state_unfused.clone() + + state_fused = initial_state.clone() + output_fused = _kda_prefill_entry()( + q=q_unfused, + k=k_unfused, + beta=beta, + state_pool=state_fused, + use_beta_sigmoid_in_kernel=True, + **common, + ) + + _assert_close(output_fused, output_unfused) + _assert_close(state_fused, state_unfused) + + +# fused_k123's K1 warps meet at this barrier after the cross-warp scan, then +# read the scanned prefix rows and the chunk total from sPartialLast. +_K1_SCAN_BARRIER = " k1_internal_barrier()\n" +_K1_SCAN_READS = ( + _K1_SCAN_BARRIER + + "\n" + + " for vi in cutlass.range_constexpr(VEC):\n" + + " rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi]\n" +) +# K1 warp 1 reads the prefix row 0 and the chunk total. +_LATE_K1_WARP = ( + " if k1_warp == 1:\n" + " late_t = _test_globaltimer()\n" + " late_end = late_t + cutlass.Int64(20000)\n" + " while late_t < late_end:\n" + " late_t = _test_globaltimer()\n" +) +_GLOBALTIMER = """ + +@dsl_user_op +def _test_globaltimer(*, loc=None, ip=None): + return cutlass.Int64( + llvm.inline_asm( + T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) +""" + + +@torch.no_grad() +@pytest.mark.parametrize( + "sequence_lengths", + [None, [100, 257, 63]], + ids=["padded_eqlen", "varlen"], +) +def test_fused_prefill_does_not_depend_on_k1_warp_timing( + sequence_lengths: list[int] | None, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A K1 warp that reads the scanned partial sums late gets the same bits. + + A copy of fused_k123 stalls K1 warp 1 for 20 us right after the scan + barrier, so the other K1 warps run ahead into the next chunk before it + reads sPartialLast. Outputs and states must equal the unpatched kernel's. + """ + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.kda import ( # noqa: E501 + fused_k123, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm import ( + _kda_prefill as kimi_k3_ops, + ) + + source = Path(fused_k123.__file__).read_text() + assert source.count(_K1_SCAN_READS) == 1, "fused_k123's K1 scan reads moved" + late_source = source.replace( + _K1_SCAN_READS, _K1_SCAN_BARRIER + _LATE_K1_WARP + _K1_SCAN_READS[len(_K1_SCAN_BARRIER) :] + ) + late_path = tmp_path / "fused_k123_late_k1_warp.py" + late_path.write_text(late_source + _GLOBALTIMER) + spec = importlib.util.spec_from_file_location(late_path.stem, late_path) + late_module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, late_path.stem, late_module) + spec.loader.exec_module(late_module) + + torch.manual_seed(10) + if sequence_lengths is None: + sequence_length = 300 + num_sequences = 1 + cu_seqlens = None + else: + sequence_length = sum(sequence_lengths) + num_sequences = len(sequence_lengths) + cu_seqlens = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + + def randn(*shape: int, dtype: torch.dtype = torch.bfloat16, scale: float = 0.05): + return (torch.randn(*shape, dtype=torch.float32, device="cuda") * scale).to(dtype) + + def l2norm(x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.normalize(x.float(), dim=-1).to(torch.bfloat16) + + shape = (1, sequence_length, NUM_HEADS, HEAD_DIM) + inputs = dict( + q=l2norm(randn(*shape)), + k=l2norm(randn(*shape)), + v=randn(*shape), + g=randn(*shape), + beta=randn(1, sequence_length, NUM_HEADS, dtype=torch.float32), + scale=HEAD_DIM**-0.5, + state_indices=torch.arange(num_sequences, dtype=torch.int32, device="cuda"), + safe_gate=True, + lower_bound=-5.0, + use_gate_in_kernel=True, + A_log=torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.1, + dt_bias=torch.randn(NUM_HEADS * HEAD_DIM, dtype=torch.float32, device="cuda") * 0.1, + cu_seqlens=cu_seqlens, + use_beta_sigmoid_in_kernel=True, + ) + initial_state = randn( + num_sequences, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=torch.float32, scale=0.01 + ) + + def run() -> tuple[torch.Tensor, torch.Tensor]: + state = initial_state.clone() + output = _kda_prefill_entry()(state_pool=state, **inputs).clone() + return output, state + + expected_output, expected_state = run() + monkeypatch.setattr(kimi_k3_ops, "_fused_make_host", late_module.make_host_function) + monkeypatch.setattr(kimi_k3_ops, "_fused_k123_cache", {}) + output, state = run() + + torch.testing.assert_close(output, expected_output, rtol=0, atol=0) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) + + +@torch.no_grad() +def test_kda_mixer_empty_prefill(): + """An empty token payload still publishes the conv state for its slot.""" + optimized = _make_kda() + hidden_states = torch.empty(1, 0, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + # Seed the pool with a sentinel: the output shape alone would still match if the + # zero-token state write-back were skipped, so only the pool can catch that. + conv_pool = torch.full( + (1, 3 * NUM_HEADS * HEAD_DIM, CONV_KERNEL_SIZE - 1), + 7.0, + dtype=torch.bfloat16, + device="cuda", + ) + # The engine never runs a forward without tokens, so o_proj (whose GEMMs need a row) is not + # part of this contract. + core = _run_production_prefill( + optimized, hidden_states, conv_pool=conv_pool, project_output=False + ) + assert core.shape == (0, NUM_HEADS, HEAD_DIM) + # Without an initial state a zero-token sequence leaves an all-zero causal window. + assert torch.equal(conv_pool, torch.zeros_like(conv_pool)) + + +@torch.no_grad() +def test_kda_mixer_empty_prefill_preserves_initial_conv_state(): + """A zero-token payload carrying an initial state leaves that state untouched.""" + optimized = _make_kda() + hidden_states = torch.empty(1, 0, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + conv_seed = ( + torch.randn( + 1, + 3 * NUM_HEADS * HEAD_DIM, + CONV_KERNEL_SIZE - 1, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + conv_pool = conv_seed.clone() + core = _run_production_prefill( + optimized, + hidden_states, + conv_pool=conv_pool, + has_initial_states=torch.ones(1, dtype=torch.bool, device="cuda"), + project_output=False, + ) + assert core.shape == (0, NUM_HEADS, HEAD_DIM) + # No token shifts into the window, so the write-back must republish it verbatim. + assert torch.equal(conv_pool, conv_seed) + + +@torch.no_grad() +def test_kda_prefill_op_partial_final_chunk_large_batch(): + """Large packed batches with short final sequences match the FLA reference. + + Cover the autotuner warmup split [8191, 1] and a mixed-length 8192-token batch. + """ + optimized = _make_kda() + reference = _make_reference(optimized) + for sequence_lengths in ([8191, 1], [8000, 150, 42]): + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + + +@torch.no_grad() +@pytest.mark.parametrize( + "sequence_lengths", + ([1], [30], [65], [6, 12], [64, 64], [1, 2, 3], [64, 64, 63], [6, 12, 20, 25]), + ids=( + "single-token", + "short-single-seq", + "partial-single-seq", + "two-short-seqs", + "two-aligned-seqs", + "three-short-seqs", + "three-seqs-partial-tail", + "four-short-seqs", + ), +) +def test_kda_prefill_small_varlen_matches_reference(sequence_lengths: list[int]) -> None: + """Short packed prompts (shorter than the convolution window, partial chunks) on the prefill + op match the reference.""" + optimized = _make_kda() + reference = _make_reference(optimized) + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + + +@torch.no_grad() +def test_kda_prefill_continuation_matches_one_shot_prefill() -> None: + """Chunked prefill: the first chunks fresh, the rest from the pool rows they left (mixed + initial states in one batch, permuted slots), gives the outputs and final rows of one prefill + of the whole sequences; an unused slot stays untouched.""" + torch.manual_seed(3) + kda = _make_kda() + lengths, first = [128, 192, 100, 77], [64, 64, 64, 13] + slots = 5 + slot_indices = torch.tensor([3, 0, 4, 1], dtype=torch.int32, device="cuda") + hidden = [ + torch.randn(1, n, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 for n in lengths + ] + projection_size = NUM_HEADS * HEAD_DIM + + def pools(): + conv = torch.randn( + slots, 3 * projection_size, CONV_KERNEL_SIZE - 1, dtype=torch.bfloat16, device="cuda" + ) + state = torch.randn( + slots, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=torch.float32, device="cuda" + ) + return conv, state + + def cumulative(seq_lengths): + return torch.tensor( + [0, *torch.tensor(seq_lengths).cumsum(0).tolist()], dtype=torch.long, device="cuda" + ) + + def prefill(chunks, conv, state, has_initial_states): + seq_lengths = [chunk.shape[1] for chunk in chunks] + output = _run_production_prefill( + kda, + torch.cat(chunks, dim=1), + cumulative(seq_lengths), + conv_pool=conv, + state_pool=state, + slot_indices=slot_indices, + has_initial_states=torch.tensor(has_initial_states, device="cuda"), + ) + return list(output.split(seq_lengths, dim=1)) + + conv_seed, state_seed = pools() + whole_conv, whole_state = conv_seed.clone(), state_seed.clone() + whole = prefill(hidden, whole_conv, whole_state, [False] * 4) + + conv, state = conv_seed.clone(), state_seed.clone() + head = prefill([h[:, :n] for h, n in zip(hidden, first)], conv, state, [False] * 4) + tail = prefill([h[:, n:] for h, n in zip(hidden, first)], conv, state, [True] * 4) + + for i in range(len(lengths)): + _assert_close(torch.cat([head[i], tail[i]], dim=1), whole[i]) + rows = slot_indices.long() + # The two batches run the QKVG GEMM on other row counts, so the convolution rows (raw + # projection outputs) may differ in their last bits. + _assert_close(conv[rows], whole_conv[rows]) + _assert_close(state[rows], whole_state[rows]) + unused_slot = 2 + assert torch.equal(conv[unused_slot], conv_seed[unused_slot]) + assert torch.equal(state[unused_slot], state_seed[unused_slot]) + + +@torch.no_grad() +def test_kda_prefill_rejects_unaligned_slot_indices() -> None: + """The prefill op addresses the pool through 16-byte aligned indices (the metadata's state + indices); an offset int32 view is an error, not a different kernel.""" + kda = _make_kda(num_heads=2) + slot_storage = torch.tensor([-1, 3, 0], dtype=torch.int32, device="cuda") + slot_indices = slot_storage[1:] + assert slot_indices.data_ptr() % 16 != 0 + hidden_states = torch.randn(1, 128, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + with pytest.raises(RuntimeError, match="kda_prefill entry"): + _run_production_prefill( + kda, + hidden_states, + torch.tensor([0, 64, 128], dtype=torch.long, device="cuda"), + conv_pool=torch.zeros(4, 3 * kda.proj_size, 3, dtype=torch.bfloat16, device="cuda"), + state_pool=torch.zeros(4, 2, HEAD_DIM, HEAD_DIM, dtype=torch.float32, device="cuda"), + slot_indices=slot_indices, + ) + + +@torch.no_grad() +@pytest.mark.parametrize("sequence_lengths", [[64, 64, 64, 63], [300]], ids=["varlen", "single"]) +def test_fused_conv_prefill_matches_unfused_chain_with_initial_states( + monkeypatch: pytest.MonkeyPatch, sequence_lengths: list[int] +) -> None: + """features.kda_fused_conv_prefill (fused_kda_conv_prefill on the token-major QKVG output) + carries continuation state like extract_transpose_prefill_slice + causal_conv1d_fn + + fused_kda_post_conv: the same outputs and recurrent rows, the same convolution rows bit for + bit, for mixed initial states, permuted slots and an unused slot; with and without + features.kda_prepad_prefill.""" + torch.manual_seed(5) + # The fused kernel's channel-contiguous taps are built with the layer's constants. + monkeypatch.setattr(kimi_kda_mixer, "_KDA_FUSED_CONV_PREFILL", True) + kda = _make_kda() + num_sequences = len(sequence_lengths) + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], dtype=torch.long, device="cuda" + ) + hidden_states = ( + torch.randn(1, sum(sequence_lengths), HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + * 0.05 + ) + slots = 5 + slot_indices = torch.tensor([3, 0, 4, 1][:num_sequences], dtype=torch.int32, device="cuda") + has_initial_states = torch.tensor([True, False, True, False][:num_sequences], device="cuda") + conv_seed = ( + torch.randn( + slots, 3 * kda.proj_size, CONV_KERNEL_SIZE - 1, dtype=torch.bfloat16, device="cuda" + ) + * 0.05 + ) + state_seed = ( + torch.randn(slots, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=torch.float32, device="cuda") * 0.05 + ) + + def run(fused_conv: bool, prepad: bool): + """Prefill from fresh copies of the seeded pools; return output and pools.""" + monkeypatch.setattr(kimi_kda_mixer, "_KDA_FUSED_CONV_PREFILL", fused_conv) + monkeypatch.setattr(kimi_kda_mixer, "_KDA_PREPAD_PREFILL", prepad) + conv_pool, state_pool = conv_seed.clone(), state_seed.clone() + output = _run_production_prefill( + kda, + hidden_states, + cumulative_lengths, + conv_pool=conv_pool, + state_pool=state_pool, + slot_indices=slot_indices, + has_initial_states=has_initial_states, + single_sequence_length=sequence_lengths[0] if num_sequences == 1 else None, + ) + return output, conv_pool, state_pool + + unfused = run(False, False) + rows = slot_indices.long() + unused_slot = 2 + for fused_conv, prepad in ((True, False), (True, True), (False, True)): + actual = run(fused_conv, prepad) + _assert_close(actual[0], unfused[0]) + # Convolution states are copied verbatim out of the projection. + torch.testing.assert_close(actual[1], unfused[1], rtol=0, atol=0) + _assert_close(actual[2].index_select(0, rows), unfused[2].index_select(0, rows)) + assert torch.equal(actual[1][unused_slot], conv_seed[unused_slot]) + assert torch.equal(actual[2][unused_slot], state_seed[unused_slot]) + + +@torch.no_grad() +def test_kda_prefill_op_shape_growth_and_cu_dtype_transitions(): + """Changing token counts and caller offset dtypes preserve reference parity. + + The production adapter prepares int64 packed offsets for FlashInfer from + either caller dtype. Repeated growth and shrinkage exercise cached kernels. + """ + torch.manual_seed(0) + optimized = _make_kda() + reference = _make_reference(optimized) + cases = [ + ([517, 654], torch.long), # small batch, int64 cu (dump-replay-like) + ([8191], torch.int32), # growth + dtype conversion + ([1171], torch.long), # shrink + flip back + ([150, 900, 333, 640], torch.int32), # multi-seq int32 + ] + for sequence_lengths, cu_dtype in cases: + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=cu_dtype, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + + +@torch.no_grad() +def test_kda_prefill_context_batch_sizes_share_one_k4_compile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Context batches of different sizes run one compiled K4. + + A bf16 state pool reaches K4 as a dense fp32 copy with one row per context + sequence, so the pool's row count follows the batch size; an fp32 pool is + addressed in place. Every batch on either pool matches the reference, and + the two pools end with the same state rows. + """ + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm import ( + _kda_prefill as prefill_ops, + ) + + monkeypatch.setattr(prefill_ops, "_k4p_cache", {}) + torch.manual_seed(5) + optimized = _make_kda() + reference = _make_reference(optimized) + slots = 6 + projection_size = NUM_HEADS * HEAD_DIM + for sequence_lengths in ([200], [64, 130, 77], [96, 33], [150, 64, 40, 129, 10]): + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn(1, sum(sequence_lengths), HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + * 0.05 + ) + slot_indices = torch.arange( + slots - len(sequence_lengths), slots, dtype=torch.int32, device="cuda" + ) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + final_rows = {} + for pool_dtype in (torch.bfloat16, torch.float32): + state_pool = torch.zeros( + slots, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=pool_dtype, device="cuda" + ) + conv_pool = torch.zeros( + slots, + 3 * projection_size, + CONV_KERNEL_SIZE - 1, + dtype=torch.bfloat16, + device="cuda", + ) + actual = _run_production_prefill( + optimized, + hidden_states, + cumulative_lengths, + conv_pool=conv_pool, + state_pool=state_pool, + slot_indices=slot_indices, + ) + _assert_close(actual, expected) + final_rows[pool_dtype] = state_pool.index_select(0, slot_indices).float() + _assert_close(final_rows[torch.bfloat16], final_rows[torch.float32]) + assert len(prefill_ops._k4p_cache) == 1 + + +def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + return ((actual.float() - expected.float()).norm() / (expected.float().norm() + 1e-12)).item() + + +@pytest.mark.parametrize( + "pool_dtype", [torch.float32, torch.bfloat16], ids=["fp32_pool", "bf16_pool"] +) +@torch.no_grad() +def test_kda_prefill_midchunk_snapshot_matches_prefix_prefill(pool_dtype: torch.dtype) -> None: + """A mid-chunk snapshot holds the state and convolution window of its sequence's prefix. + + Sequences of 200, 130 and 96 tokens snapshot after 128 tokens (sequence 0) and 64 tokens + (sequence 2) into their own slots; sequence 1 takes none. Each snapshot row must match a + prefill of that prefix alone, and lie much closer to it than to the state up to one chunk + later (sequence 2 ends 32 tokens after its snapshot). The snapshots leave the batch's outputs + and final states as a run without them. + """ + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kda import ( # noqa: E501 + build_kda_snapshot_plan, + ) + + torch.manual_seed(7) + optimized = _make_kda() + projection_size = NUM_HEADS * HEAD_DIM + lengths = [200, 130, 96] + starts = [0, 200, 330] + snapshots = {0: (6, 128), 2: (7, 64)} # sequence -> (slot, prefix tokens) + slots = 8 + hidden_states = ( + torch.randn(1, sum(lengths), HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + ) + + def pools(): + state = torch.zeros(slots, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=pool_dtype, device="cuda") + conv = torch.zeros( + slots, 3 * projection_size, CONV_KERNEL_SIZE - 1, dtype=torch.bfloat16, device="cuda" + ) + return state, conv + + def cumulative(seq_lengths): + return torch.tensor( + [0, *torch.tensor(seq_lengths).cumsum(0).tolist()], dtype=torch.long, device="cuda" + ) + + slot_indices = torch.tensor([1, 3, 5], dtype=torch.int32, device="cuda") + plan = build_kda_snapshot_plan( + [(seq, slot, offset) for seq, (slot, offset) in snapshots.items()], + lengths, + torch.device("cuda"), + ) + state, conv = pools() + output = _run_production_prefill( + optimized, + hidden_states, + cumulative(lengths), + conv_pool=conv, + state_pool=state, + slot_indices=slot_indices, + kda_snapshot=plan, + ) + plain_state, plain_conv = pools() + plain_output = _run_production_prefill( + optimized, + hidden_states, + cumulative(lengths), + conv_pool=plain_conv, + state_pool=plain_state, + slot_indices=slot_indices, + ) + _assert_close(output, plain_output) + _assert_close(state[slot_indices.long()], plain_state[slot_indices.long()]) + _assert_close(conv[slot_indices.long()], plain_conv[slot_indices.long()]) + + def prefix_rows(extra_tokens: int): + """Final state and convolution rows of each snapshotted sequence's first + ``offset + extra_tokens`` tokens (at most the whole sequence).""" + prefix_lengths = [ + min(offset + extra_tokens, lengths[seq]) for seq, (_, offset) in snapshots.items() + ] + prefix_hidden = torch.cat( + [ + hidden_states[:, starts[seq] : starts[seq] + n] + for seq, n in zip(snapshots, prefix_lengths) + ], + dim=1, + ) + prefix_state, prefix_conv = pools() + prefix_slots = torch.arange(len(snapshots), dtype=torch.int32, device="cuda") + _run_production_prefill( + optimized, + prefix_hidden, + cumulative(prefix_lengths), + conv_pool=prefix_conv, + state_pool=prefix_state, + slot_indices=prefix_slots, + ) + return prefix_state[: len(snapshots)], prefix_conv[: len(snapshots)] + + expected_state, expected_conv = prefix_rows(0) + later_state, _ = prefix_rows(64) + snapshot_slots = torch.tensor([slot for slot, _ in snapshots.values()], device="cuda") + _assert_close(state[snapshot_slots], expected_state) + _assert_close(conv[snapshot_slots], expected_conv) + for j in range(len(snapshots)): + assert _relative_l2(state[snapshot_slots[j]], expected_state[j]) < 0.1 * _relative_l2( + later_state[j], expected_state[j] + ) + untouched = [s for s in range(slots) if s not in {1, 3, 5, *snapshot_slots.tolist()}] + assert torch.count_nonzero(state[untouched]) == 0 + + +@pytest.mark.parametrize( + "pool_dtype", [torch.float32, torch.bfloat16], ids=["fp32_pool", "bf16_pool"] +) +@torch.no_grad() +def test_kda_prefill_inforward_snapshot_swaps_live_and_staging_rows( + pool_dtype: torch.dtype, +) -> None: + """In-forward snapshots (``MambaHybridCacheManagerV2`` with block reuse): each planned + sequence's live slot ends with the state and convolution window of its prefix up to the + snapshot point, and its staging row with the chunk-end ones, which + ``finish_inforward_snapshots`` copies back after committing the prefix blocks. + + Sequences of 200, 130 and 96 tokens (live slots 1, 3, 5) snapshot after 128 tokens + (sequence 0, staging row 1) and 64 tokens (sequence 2, staging row 0). The outputs and the + unplanned sequence's rows match a run without snapshots. + """ + from dataclasses import replace + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kda import ( # noqa: E501 + build_kda_snapshot_plan, + ) + + torch.manual_seed(11) + optimized = _make_kda() + projection_size = NUM_HEADS * HEAD_DIM + lengths = [200, 130, 96] + starts = [0, 200, 330] + live = [1, 3, 5] + planned = {0: (128, 1), 2: (64, 0)} # sequence -> (prefix tokens, staging row) + slots = 8 + hidden_states = ( + torch.randn(1, sum(lengths), HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + ) + + def pools(rows=slots): + state = torch.zeros(rows, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=pool_dtype, device="cuda") + conv = torch.zeros( + rows, 3 * projection_size, CONV_KERNEL_SIZE - 1, dtype=torch.bfloat16, device="cuda" + ) + return state, conv + + def cumulative(seq_lengths): + return torch.tensor( + [0, *torch.tensor(seq_lengths).cumsum(0).tolist()], dtype=torch.long, device="cuda" + ) + + slot_indices = torch.tensor(live, dtype=torch.int32, device="cuda") + plan = build_kda_snapshot_plan( + [(seq, live[seq], offset) for seq, (offset, _) in planned.items()], + lengths, + torch.device("cuda"), + ) + plan = replace( + plan, + stage_rows=torch.tensor([row for _, row in planned.values()], device="cuda"), + ) + state, conv = pools() + stage_state, stage_conv = pools(len(planned)) + output = _run_production_prefill( + optimized, + hidden_states, + cumulative(lengths), + conv_pool=conv, + state_pool=state, + slot_indices=slot_indices, + kda_snapshot=plan, + inforward_stage=(stage_state, stage_conv), + ) + plain_state, plain_conv = pools() + plain_output = _run_production_prefill( + optimized, + hidden_states, + cumulative(lengths), + conv_pool=plain_conv, + state_pool=plain_state, + slot_indices=slot_indices, + ) + _assert_close(output, plain_output) + _assert_close(state[3], plain_state[3]) + _assert_close(conv[3], plain_conv[3]) + + prefix_lengths = [offset for offset, _ in planned.values()] + prefix_hidden = torch.cat( + [ + hidden_states[:, starts[seq] : starts[seq] + n] + for seq, n in zip(planned, prefix_lengths) + ], + dim=1, + ) + prefix_state, prefix_conv = pools() + _run_production_prefill( + optimized, + prefix_hidden, + cumulative(prefix_lengths), + conv_pool=prefix_conv, + state_pool=prefix_state, + slot_indices=torch.arange(len(planned), dtype=torch.int32, device="cuda"), + ) + for j, (seq, (_, row)) in enumerate(planned.items()): + slot = live[seq] + _assert_close(state[slot], prefix_state[j]) + _assert_close(conv[slot], prefix_conv[j]) + _assert_close(stage_state[row], plain_state[slot]) + _assert_close(stage_conv[row], plain_conv[slot]) + assert _relative_l2(state[slot], prefix_state[j]) < 0.1 * _relative_l2( + plain_state[slot], prefix_state[j] + ) + untouched = [s for s in range(slots) if s not in live] + assert torch.count_nonzero(state[untouched]) == 0 + assert torch.count_nonzero(conv[untouched]) == 0 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kimi_kda_fused_verify_parity.py new file mode 100644 index 000000000000..91f0e82d3806 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kda/test_kimi_kda_fused_verify_parity.py @@ -0,0 +1,716 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime-level parity: K3KDALinearAttention fused verify vs a sequential FLA reference. + +Simulates two chained speculative-verification rounds in both worlds: + +* Sequential world (``_reference_verify``): per step, the depthwise convolution on the gathered + window and ``fla.ops.kda.fused_recurrent_kda``, each step's conv window and state written to + intermediate buffers, plus the legacy promotion (copy the accepted step's conv window / SSM + state into the live pools). +* Fused world: ``K3KDALinearAttention.forward_verify`` on per-slot replay caches (the + ``kda_mtp_decode`` entry), in-place state commit after the golden token, and + only the accepted-draft count recorded between rounds. + +Both worlds share the layer's FP8 projections (the fused QKVG and o_proj weight-read GEMMs) and +take identical hidden states; the fused world runs its gate projection on the aux stream. With +mixed per-request acceptance between rounds, matching round-2 outputs proves the packed row/slot +bookkeeping, conv-window seeding and the in-kernel gated RMSNorm epilogue reproduce the +promoted-state semantics. + +Requires 1 Blackwell GPU, fla-core, nvidia-cutlass-dsl. Skips otherwise. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +_HAVE_DEPS = True +_DEP_ERR = None +try: + import cuda.bindings.driver # noqa: F401 + import cutlass # noqa: F401 + from fla.ops.kda import fused_recurrent_kda # noqa: F401 + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.fp8_weight_read import ( # noqa: E501 + _Fp8BlockScaleWeightReadLinear, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kda import ( + K3KDALinearAttention, + ) + + # The model module transitively imports the optional deps above, so it + # must stay behind the guard too or collection fails instead of skipping. + from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig + from tensorrt_llm._torch.modules.mamba.layernorm_gated import rms_norm_gated_token_major +except ImportError as e: + _HAVE_DEPS = False + _DEP_ERR = str(e) + + +def _is_blackwell(): + if not torch.cuda.is_available(): + return False + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor in (100, 103, 107) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), + pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103/sm107"), + pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), +] + +HIDDEN = 512 +H = 8 # per-rank head count outside the drop's tuned set — general variant +K = 128 +W = 4 +M = 2 # draft tokens per round +LB = -5.0 + + +@torch.no_grad() +def _make_runtime(seed, aux_stream=None, *, num_heads=H): + """A finalized layer on seeded FP8 checkpoint-like weights, with its fused QKVG and o_proj + swapped to the K3 FP8 weight-read GEMM as the model loader does.""" + # A real KimiLinearConfig (not a SimpleNamespace) so the runtime sees the + # same config surface it does in production. ``linear_attn_config`` carries + # the per-layer KDA params the runtime reads plus the (unused here) + # kda_layers/full_attn_layers schedule the config's own validation requires. + cfg = KimiLinearConfig( + hidden_size=HIDDEN, + rms_norm_eps=1e-5, + linear_attn_config=dict( + kda_layers=[1], + full_attn_layers=[], + num_heads=num_heads, + head_dim=K, + short_conv_kernel_size=W, + use_full_rank_gate=True, + gate_lower_bound=LB, + ), + ) + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.quantization import QuantAlgo + + model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES)) + rt = K3KDALinearAttention( + cfg, + layer_idx=0, + aux_stream=aux_stream, + model_config=model_config, + ).to("cuda") + gen = torch.Generator(device="cuda").manual_seed(seed) + for name, p in rt.named_parameters(): + if p.dtype == torch.float8_e4m3fn: + p.copy_(torch.randint(-32, 33, p.shape, generator=gen, device="cuda").to(p.dtype)) + elif name.endswith("weight_scale"): + p.fill_(0.001) + elif name.endswith("input_scale"): + p.fill_(1.0) + elif name.endswith("A_log"): + p.copy_(torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.5) + elif name.endswith("dt_bias"): + p.copy_(torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.1) + else: + p.copy_( + (torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.03).to( + p.dtype + ) + ) + # Weight-load finalize as in production (after the random init above, which the fused + # projections and kernel constants snapshot). + rt.finalize_decode_weights() + rt._build_mtp_conv_weights() + for name in ("qkvg_proj", "o_proj"): + module = _Fp8BlockScaleWeightReadLinear.from_linear(getattr(rt, name)) + setattr(rt, name, module) + module.post_load_weights() + return rt + + +def _make_pools(B, seed, *, num_heads=H): + gen = torch.Generator(device="cuda").manual_seed(seed) + d = num_heads * K + conv_pool = ( + torch.randn(B, 3 * d, W - 1, generator=gen, device="cuda", dtype=torch.float32) * 0.5 + ).to(torch.bfloat16) + ssm_pool = torch.randn(B, num_heads, K, K, generator=gen, device="cuda", dtype=torch.float32) + ssm_pool *= torch.linspace(0.5, 1.5, K, device="cuda").view(1, 1, K, 1) + return conv_pool, ssm_pool + + +def _make_fused_layer_cache(B, conv_pool, *, num_heads=H): + """Replay caches shaped like PythonMambaCacheManager's KDA allocation, + with the committed conv window seeded from the base pool (the prefill + seeding contract: the base pool stores committed columns directly).""" + d = num_heads * K + S = W - 1 + M + + def _conv_cache(section): + cache = torch.zeros(B, S, d, device="cuda", dtype=torch.float32).transpose(-1, -2) + cache[:, :, : W - 1] = conv_pool[:, section * d : (section + 1) * d].float() + return cache + + return SimpleNamespace( + kda_conv_q=_conv_cache(0), + kda_conv_k=_conv_cache(1), + kda_conv_v=_conv_cache(2), + kda_k_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.bfloat16), + kda_g_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.float32), + kda_v_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.bfloat16), + kda_beta_cache=torch.zeros(B, M, num_heads, device="cuda", dtype=torch.bfloat16), + prev_num_accepted_tokens=torch.zeros(B, dtype=torch.int32, device="cuda"), + has_kda_replay_caches=True, + intermediate_conv_window=None, + intermediate_ssm=None, + ) + + +def _make_seq_layer_cache(B, *, num_heads=H): + d = num_heads * K + return SimpleNamespace( + kda_k_cache=None, + has_kda_replay_caches=False, + intermediate_conv_window=torch.zeros( + B, M + 1, 3 * d, W - 1, device="cuda", dtype=torch.bfloat16 + ), + intermediate_ssm=torch.zeros(B, M + 1, num_heads, K, K, device="cuda", dtype=torch.float32), + ) + + +@torch.no_grad() +def _reference_verify(rt, x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices): + """Per-step FLA verification on gathered copies of the live rows (the pools stay read-only), + the layer's own projections in front and its gated output RMSNorm behind. Each step's conv + window and state go to ``layer_cache``'s intermediate buffers. Returns the normalized core + ``[tokens, H, D]``.""" + from fla.ops.kda import fused_recurrent_kda + + d, hd, heads = rt.proj_size, rt.head_dim, rt.num_heads + batch = x2d.shape[0] // num_steps + x = x2d.view(batch, num_steps, -1) + q_in, k_in, v_in, onorm_g = rt.qkvg_proj(x).split(d, dim=-1) + bfa = F.linear(x, rt._bfa_proj_weight) + beta = bfa[..., hd : hd + heads].float() + gate = F.linear(bfa[..., :hd], rt.f_b_proj.weight).unflatten(-1, (heads, hd)) + rows = slot_indices.long() + window = conv_pool.index_select(0, rows) # [B, 3D, W - 1] + state = ssm_pool.index_select(0, rows).float() + conv_weight = ( + torch.cat([rt.q_conv1d.weight, rt.k_conv1d.weight, rt.v_conv1d.weight]).squeeze(1).float() + ) # [3D, W] + outputs = [] + for t in range(num_steps): + step = torch.cat([q_in[:, t], k_in[:, t], v_in[:, t]], dim=-1).to(window.dtype) + full = torch.cat([window, step.unsqueeze(-1)], dim=-1) # [B, 3D, W] + window = full[..., 1:] + conv = F.silu((full.float() * conv_weight).sum(-1)).to(x2d.dtype) + q_t, k_t, v_t = ( + part.unflatten(-1, (heads, hd)).unsqueeze(1) for part in conv.split(d, dim=-1) + ) + o_t, state = fused_recurrent_kda( + q=q_t, + k=k_t, + v=v_t, + g=gate[:, t : t + 1], + beta=beta[:, t : t + 1], + A_log=rt.A_log, + dt_bias=rt.dt_bias, + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=rt.gate_lower_bound, + state_v_first=True, + ) + outputs.append(o_t) + layer_cache.intermediate_conv_window[:batch, t] = window.to( + layer_cache.intermediate_conv_window.dtype + ) + layer_cache.intermediate_ssm[:batch, t] = state.to(layer_cache.intermediate_ssm.dtype) + o = torch.cat(outputs, dim=1) # [B, T, H, D] + return rms_norm_gated_token_major( + o.reshape(-1, hd), + onorm_g.reshape(-1, heads, hd), + rt.o_norm.weight, + rt.o_norm.eps, + gate_activation="sigmoid", + ).reshape(-1, heads, hd) + + +def _promote_sequential(layer_cache, conv_pool, ssm_pool, accept): + """The manager's legacy promotion: accepted step's states -> pools.""" + B = conv_pool.shape[0] + rows = torch.arange(B, device="cuda") + conv_pool.copy_(layer_cache.intermediate_conv_window[rows, accept]) + ssm_pool.copy_(layer_cache.intermediate_ssm[rows, accept]) + + +def _capture_mtp_verify(runtime): + """Record the kwargs of every fused verify launch, then run it.""" + calls = [] + mtp_verify = runtime._dispatch.mtp_verify + + def _capturing_mtp_verify(**kwargs): + calls.append(kwargs) + return mtp_verify(**kwargs) + + runtime._dispatch.mtp_verify = _capturing_mtp_verify + return calls + + +def _capture_mtp_verify_fused_fb(runtime): + """Record the kwargs of every fused verify launch with the forget gate from ``f_a``, then + run it.""" + calls = [] + mtp_verify_fused_fb = runtime._dispatch.mtp_verify_fused_fb + + def _capturing_mtp_verify_fused_fb(**kwargs): + calls.append(kwargs) + return mtp_verify_fused_fb(**kwargs) + + runtime._dispatch.mtp_verify_fused_fb = _capturing_mtp_verify_fused_fb + return calls + + +def _assert_fused_epilogue_call(call, runtime): + """The fused verify ran with the in-kernel output norm on packed rows.""" + assert call["fuse_output_norm"] + assert call["packed_token_layout"] + assert call["cu_seqlens"] is None + assert not call["mxfp8_output"] + onorm_g = call["onorm_g"] + assert onorm_g.stride(-1) == 1 and onorm_g.stride(-2) == K + # The gate and q/k/v stay strided views of the fused QKVG output. + row_stride = 4 * runtime.proj_size + assert onorm_g.stride(1) == row_stride + for name in ("x_q", "x_k", "x_v"): + assert call[name].stride(1) == row_stride + assert call[name].stride(-1) == 1 + assert call["beta"].stride(-1) == 1 + + +def _rep(name, a, b): + a, b = a.float(), b.float() + cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() + rel = ((a - b).norm() / (b.norm() + 1e-12)).item() + print(f" {name}: cos={cos:.6f} rel_l2={rel:.3e}") + return cos > 0.999 and rel < 3e-2 + + +def _make_kvcm_replay_manager(num_heads): + """A V2 hybrid manager whose KVCM pools hold the KDA states and replay caches, with a host tier to offload to.""" + from tensorrt_llm._torch.modules.kimi_kda.cache_manager import KimiK3HybridCacheManager + from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig + from tensorrt_llm.mapping import Mapping + + return KimiK3HybridCacheManager( + mamba_d_state=K, + mamba_d_conv=W, + mamba_num_heads=num_heads, + mamba_n_groups=num_heads, + mamba_head_dim=K, + mamba_num_layers=2, + mamba_layer_mask=[True, True], + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig( + max_tokens=128, host_cache_size=64 << 20, enable_block_reuse=False + ), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=0, + num_kv_heads=0, + head_dim=K, + tokens_per_block=32, + max_seq_len=64, + max_batch_size=4, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + spec_config=MTPDecodingConfig(max_draft_len=M), + layer_mask=[False, False], + conv_state_layout="q_k_v", + kda_replay_num_spec=M, + ) + + +def _offload_and_resume(manager, request_ids): + """Suspend the requests' KV caches, push their pages to the host tier with newer caches, and resume them; returns + the slots they resume in.""" + from tensorrt_llm.runtime.kv_cache_manager_v2 import CACHE_LEVEL1, LayerId, _introspection + + stream = torch.cuda.current_stream().cuda_stream + caches = [manager.kv_cache_map[rid] for rid in request_ids] + for cache in caches: + cache.suspend() + group = manager.impl.get_layer_group_id(LayerId(0)) + pool = next( + pool + for pool in manager.impl.pool_group_descs + if any(variant.layer_group_id == group for variant in pool.slot_desc.variants) + ) + pressure = [] + try: + for _ in range(pool.num_slots + len(caches)): + other = manager.impl.create_kv_cache() + pressure.append(other) + assert other.resume(stream) + other.suspend() + finally: + for other in pressure: + other.close() + # The suspended requests' pages must have been evicted to the host tier, or + # the resume below would not exercise the offload path at all. + assert max(_introspection.storage_utilization(manager.impl, CACHE_LEVEL1)) > 0 + new_slots = [] + for cache in reversed(caches): + assert cache.resume(stream) + new_slots.append(cache.get_ssm_block_base_index(group)) + return torch.tensor(list(reversed(new_slots)), dtype=torch.int32, device="cuda") + + +@torch.no_grad() +@pytest.mark.parametrize( + ("num_heads", "large_state_stride", "kvcm_replay"), + [ + (H, False, False), + (16, False, False), + (64, False, False), + (64, True, False), + (H, False, True), + ], + ids=["h8", "h16", "h64", "large-state-stride", "h8-kvcm-offload"], +) +def test_fused_vs_sequential_two_rounds(request, num_heads, large_state_stride, kvcm_replay): + """With ``kvcm_replay``, the fused world keeps its states and replay caches in a V2 manager's KVCM pools, and + the requests are offloaded to the host tier and resumed in other slots between the two rounds. + + The [f_a | b] rows go through ``k3_kda_fab`` in f_a mode, the forget gate computed inside the + verify kernel (``kda_mtp_decode`` with ``fa`` / ``w_fb``, launched with PDL), where that kernel + takes the head count (up to 16 heads), else through ``F.linear`` [f_a | b] -> f_b and + ``kda_mtp_decode`` with the fused output norm.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + kda as kimi_kda_mixer, + ) + from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream + + torch.manual_seed(0) + B = 4 + T = M + 1 + rt_seq = _make_runtime(seed=1, num_heads=num_heads) + rt_fused = _make_runtime(seed=1, aux_stream=torch.cuda.Stream(), num_heads=num_heads) + assert rt_fused.qkvg_proj is not None + assert rt_fused._bfa_proj_weight is not None + assert rt_fused._onorm_w_f32 is not None + mtp_calls = _capture_mtp_verify(rt_fused) + fb_calls = _capture_mtp_verify_fused_fb(rt_fused) + slot_indices = torch.arange(B, dtype=torch.int32, device="cuda") + + conv_pool_seq, ssm_pool_seq = _make_pools(B, seed=2, num_heads=num_heads) + conv_pool_fused = conv_pool_seq.clone() + if large_state_stride: + # Only four small states are populated. The gaps reproduce V2's + # coalesced layer layout, with the last slot beyond INT32_MAX elements. + state_stride = ((2**31 // (B - 1)) // (K * K) + 1) * K * K + storage_size = (B - 1) * state_stride + num_heads * K * K + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < storage_size * ssm_pool_seq.element_size() + 2**30: + pytest.skip("large-stride regression needs 9 GiB of free GPU memory") + storage = torch.empty(storage_size, dtype=ssm_pool_seq.dtype, device="cuda") + ssm_pool_fused = storage.as_strided(ssm_pool_seq.shape, (state_stride, K * K, K, 1)) + ssm_pool_fused.copy_(ssm_pool_seq) + else: + ssm_pool_fused = ssm_pool_seq.clone() + cache_seq = _make_seq_layer_cache(B, num_heads=num_heads) + cache_fused = _make_fused_layer_cache(B, conv_pool_fused, num_heads=num_heads) + fused_slots = slot_indices + replay_manager = _make_kvcm_replay_manager(num_heads) if kvcm_replay else None + if replay_manager is not None: + request.addfinalizer(replay_manager.shutdown) + request_ids = list(range(101, 101 + B)) + replay_manager.add_dummy_requests(request_ids, token_nums=[32] * B, is_gen=False) + fused_slots = torch.tensor( + replay_manager.get_state_indices(request_ids), dtype=torch.int32, device="cuda" + ) + cache_fused = replay_manager.mamba_layer_cache(0) + conv_pool_fused = cache_fused.conv + ssm_pool_fused = cache_fused.temporal + conv_pool_fused.index_copy_(0, fused_slots.long(), conv_pool_seq) + ssm_pool_fused.index_copy_(0, fused_slots.long(), ssm_pool_seq) + replay_manager.on_state_transfer_complete(request_ids) + + gen = torch.Generator(device="cuda").manual_seed(3) + + def tokens(scale=0.5): + return ( + torch.randn(B * T, HIDDEN, generator=gen, device="cuda", dtype=torch.float32) * scale + ).to(torch.bfloat16) + + ok = True + # ---- Round 1 (no pending drafts) ---- + x1 = tokens() + expect_fused_fb = kimi_kda_mixer._use_fab_fa( + x1.view(B, T, -1), rt_fused._bfa_proj_weight, rt_fused.f_b_proj.weight, num_heads + ) + # The f_a-mode fab kernel serves the K3 Helix head shards (up to 16 heads). + assert expect_fused_fb == (num_heads <= 16) + projected = rt_fused._project_verify_inputs(x1.view(B, T, -1), B * T, gate_from_fa=True) + beta_view = projected[4] + # beta is a view of the [f_a | b] projection in f_a mode and on the F.linear chain. + assert beta_view.stride(-2) == rt_fused._bfa_proj_weight.shape[0] + assert beta_view.stride(-1) == 1 + assert (projected[6] is not None) == expect_fused_fb + out1_seq = rt_seq._project_output( + _reference_verify(rt_seq, x1, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices) + ) + with with_multi_stream(True): + out1_fused = rt_fused._project_output( + rt_fused.forward_verify(x1, T, cache_fused, ssm_pool_fused, fused_slots) + ) + if expect_fused_fb: + assert not mtp_calls and len(fb_calls) == 1 + call = fb_calls[-1] + assert call["use_pdl"] + # f_a and beta are strided views of the f_a-mode [f_a | b] output + assert call["fa"].shape == (1, B * T, K) + assert call["fa"].stride(1) == rt_fused._bfa_proj_weight.shape[0] + assert call["beta"].stride(1) == rt_fused._bfa_proj_weight.shape[0] + assert call["onorm_g"].stride(-2) == K + assert call["onorm_g"].stride(1) == 4 * rt_fused.proj_size + else: + assert not fb_calls + _assert_fused_epilogue_call(mtp_calls[-1], rt_fused) + print("round 1:") + ok &= _rep("out", out1_fused, out1_seq) + + # ---- Acceptance: 0, 1, 2, 0 drafts across the 4 requests ---- + accept = torch.tensor([0, 1, 2, 0], dtype=torch.long, device="cuda") + _promote_sequential(cache_seq, conv_pool_seq, ssm_pool_seq, accept) + cache_fused.prev_num_accepted_tokens.index_copy_(0, fused_slots.long(), accept.to(torch.int32)) + if replay_manager is not None: + fused_slots = _offload_and_resume(replay_manager, request_ids) + + # ---- Round 2 (fused path replays the accepted drafts) ---- + x2 = tokens() + out2_seq = rt_seq._project_output( + _reference_verify(rt_seq, x2, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices) + ) + core2_fused = x2.new_empty(B * T, num_heads, K) + with with_multi_stream(True): + result2_fused = rt_fused.forward_verify( + x2, T, cache_fused, ssm_pool_fused, fused_slots, output=core2_fused + ) + assert result2_fused is core2_fused + assert len(mtp_calls) + len(fb_calls) == 2 + assert len(fb_calls) == (2 if expect_fused_fb else 0) + if not expect_fused_fb: + _assert_fused_epilogue_call(mtp_calls[-1], rt_fused) + out2_fused = rt_fused._project_output(core2_fused) + print("round 2 (mixed replay):") + ok &= _rep("out", out2_fused, out2_seq) + + # Committed pool state cross-check: fused pool holds the state after + # round-2's golden token; reproduce it in the sequential world by + # promoting with accept=0 (golden only). + _promote_sequential( + cache_seq, conv_pool_seq, ssm_pool_seq, torch.zeros(B, dtype=torch.long, device="cuda") + ) + ok &= _rep("committed ssm", ssm_pool_fused.index_select(0, fused_slots.long()), ssm_pool_seq) + + assert ok + + +def _run_one_verify_round(runtime, x, batch, seed, **verify_kwargs): + conv_pool, ssm_pool = _make_pools(batch, seed=seed) + layer_cache = _make_fused_layer_cache(batch, conv_pool) + slot_indices = torch.arange(batch, dtype=torch.int32, device="cuda") + return runtime.forward_verify(x, M + 1, layer_cache, ssm_pool, slot_indices, **verify_kwargs) + + +@torch.no_grad() +def test_fused_verify_requires_finalized_weights(): + """The verify runs on the kernel constants finalize_decode_weights builds; without them it + raises instead of computing the output gate another way.""" + batch = 2 + runtime = _make_runtime(seed=5) + runtime._onorm_w_f32 = None + x = (torch.randn(batch * (M + 1), HIDDEN, device="cuda") * 0.5).to(torch.bfloat16) + with pytest.raises(RuntimeError, match="finalize_decode_weights"): + _run_one_verify_round(runtime, x, batch, seed=6) + + +class _PrequantizedProjectionStub(torch.nn.Module): + """o_proj stand-in that records the prequantized pair it receives.""" + + supports_prequantized_input = True + + def __init__(self, out_features): + super().__init__() + self.out_features = out_features + self.calls = [] + + def forward(self, x): + return x.new_zeros(x.shape[0], self.out_features) + + def forward_prequantized(self, activation, activation_scale): + self.calls.append((activation, activation_scale)) + return torch.zeros( + activation.shape[0], self.out_features, dtype=torch.bfloat16, device="cuda" + ) + + +@torch.no_grad() +@pytest.mark.parametrize("allow, with_output", [(True, False), (False, False), (True, True)]) +def test_fused_verify_mxfp8_output_gating(allow, with_output): + """MXFP8 output only when allowed, unbuffered and accepted by o_proj.""" + batch = 2 + rows = batch * (M + 1) + runtime = _make_runtime(seed=9) + stub = _PrequantizedProjectionStub(HIDDEN) + runtime.o_proj = stub + calls = [] + + def _fake_mtp_verify(**kwargs): + calls.append(kwargs) + if kwargs["mxfp8_output"]: + activation = torch.zeros(rows, H * K, device="cuda").to(torch.float8_e4m3fn) + return activation, torch.zeros(128 * H * 4, dtype=torch.uint8, device="cuda") + return torch.zeros(1, rows, H, K, dtype=torch.bfloat16, device="cuda") + + fb_calls = [] + + def _fake_mtp_verify_fused_fb(**kwargs): + fb_calls.append(kwargs) + return torch.zeros(1, rows, H, K, dtype=torch.bfloat16, device="cuda") + + runtime._dispatch.mtp_verify = _fake_mtp_verify + runtime._dispatch.mtp_verify_fused_fb = _fake_mtp_verify_fused_fb + x = torch.zeros(rows, HIDDEN, dtype=torch.bfloat16, device="cuda") + output = x.new_empty(rows, H, K) if with_output else None + core = _run_one_verify_round( + runtime, x, batch, seed=10, output=output, allow_prequantized_output=allow + ) + out = runtime._project_output(core) + + expect_mxfp8 = allow and not with_output + assert len(calls) + len(fb_calls) == 1 + if fb_calls: + # A BF16 core with the forget gate computed from f_a in the verify kernel. + assert not expect_mxfp8 + else: + assert calls[-1]["fuse_output_norm"] + assert calls[-1]["mxfp8_output"] is expect_mxfp8 + assert len(stub.calls) == int(expect_mxfp8) + assert out.shape == (rows, HIDDEN) + if with_output: + assert core is output + + +@torch.no_grad() +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability(0) != (10, 7), + reason="the MXFP8 verify epilogue feeds the Rubin (SM107) CuTe MXFP8 GEMM", +) +def test_fused_verify_mxfp8_output_matches_bf16_projection(): + """SM107: epilogue-quantized o_proj input matches the BF16 core path.""" + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("needs the Rubin CuTe DSL") + batch = 4 + runtime = _make_runtime(seed=11) + assert isinstance(runtime.o_proj, _Fp8BlockScaleWeightReadLinear) + assert runtime.o_proj.supports_prequantized_input + calls = _capture_mtp_verify(runtime) + fb_calls = _capture_mtp_verify_fused_fb(runtime) + x = (torch.randn(batch * (M + 1), HIDDEN, device="cuda") * 0.5).to(torch.bfloat16) + + bf16_out = runtime._project_output(_run_one_verify_round(runtime, x, batch, seed=12)) + # The BF16 core: the forget gate from f_a in the verify kernel where the f_a-mode fab kernel + # serves these rows, else the fused-epilogue verify with BF16 output. + bf16_mxfp8_flags = [call["mxfp8_output"] for call in calls] + assert bf16_mxfp8_flags == ([] if fb_calls else [False]) + assert len(fb_calls) + len(calls) == 1 + mxfp8_core = _run_one_verify_round(runtime, x, batch, seed=12, allow_prequantized_output=True) + mxfp8_out = runtime._project_output(mxfp8_core) + + assert [call["mxfp8_output"] for call in calls[len(bf16_mxfp8_flags) :]] == [True] + assert len(fb_calls) + len(calls) == 2 + assert mxfp8_core.activation.dtype is torch.float8_e4m3fn + assert mxfp8_core.scale.dtype is torch.uint8 + assert _rep("mxfp8 vs bf16 o_proj input", mxfp8_out, bf16_out) + + +def _is_sm107(): + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 7) + + +@pytest.mark.skipif(not _is_sm107(), reason="SM107 MXFP8 verify epilogue") +@pytest.mark.parametrize("bf16_rows", [False, True]) +@torch.no_grad() +def test_fused_verify_with_fp8_weight_read_o_proj(bf16_rows, monkeypatch): + """FP8-resident verify on SM107 with the K3 FP8 weight-read o_proj: the verify kernel hands + o_proj its MXFP8 operand, or BF16 rows where o_proj runs them on the low-M GEMM.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + fp8_weight_read as kimi, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_fp8w_lowm_gemm, + ) + from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream + + B = 4 + T = M + 1 + monkeypatch.setattr(kimi, "_K3_FP8W_LOWM_GEMM", True) + monkeypatch.setitem( + k3_fp8w_lowm_gemm._SM107_BF16_OVER_PREQUANTIZED_MAX_M, + (HIDDEN, H * K), + B * T if bf16_rows else 0, + ) + rt_seq = _make_runtime(seed=1) + rt_fused = _make_runtime(seed=1, aux_stream=torch.cuda.Stream()) + assert rt_fused.o_proj.lowm_weight_scale is not None + assert rt_fused.o_proj.takes_prequantized_input(B * T) != bf16_rows + mtp_calls = _capture_mtp_verify(rt_fused) + fb_calls = _capture_mtp_verify_fused_fb(rt_fused) + slot_indices = torch.arange(B, dtype=torch.int32, device="cuda") + conv_pool_seq, ssm_pool_seq = _make_pools(B, seed=2) + conv_pool_fused = conv_pool_seq.clone() + ssm_pool_fused = ssm_pool_seq.clone() + cache_seq = _make_seq_layer_cache(B) + cache_fused = _make_fused_layer_cache(B, conv_pool_fused) + gen = torch.Generator(device="cuda").manual_seed(3) + x = (torch.randn(B * T, HIDDEN, generator=gen, device="cuda") * 0.5).to(torch.bfloat16) + + out_seq = rt_seq._project_output( + _reference_verify(rt_seq, x, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices) + ) + with with_multi_stream(True): + core_fused = rt_fused.forward_verify( + x, T, cache_fused, ssm_pool_fused, slot_indices, allow_prequantized_output=True + ) + out_fused = rt_fused._project_output(core_fused) + assert len(mtp_calls) + len(fb_calls) == 1 + if fb_calls: + # BF16 rows with the forget gate computed from f_a in the verify kernel. + assert bf16_rows + else: + assert mtp_calls[0]["mxfp8_output"] != bf16_rows + assert _rep("out", out_fused, out_seq) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_embedding_gather.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_embedding_gather.py new file mode 100644 index 000000000000..9f6f1b3ee697 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_embedding_gather.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""K3's token-embedding gather (``ops/embedding/k3_embedding_gather``, ``features.embedding_gather``) +is bitwise ``F.embedding``: every cell of ``k3_embedding_gather.CELLS`` under ``validating`` against +the entry's reference (out-of-range ids give zero rows), and small tables at 0..32 ids.""" + +import importlib +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 + +k3_embedding_gather = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + "kimi_k3_nvfp4__sm_107__default.ops.embedding.k3_embedding_gather" +).k3_embedding_gather + +_ARCH = {(10, 0): Arch.SM_100, (10, 3): Arch.SM_103, (10, 7): Arch.SM_107} +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +certified = pytest.mark.skipif( + not torch.cuda.is_available() + or _ARCH.get(torch.cuda.get_device_capability()) not in k3_embedding_gather.ARCHS, + reason="k3_embedding_gather is certified on its ARCHS", +) + + +@pytest.fixture(scope="module") +def k3_table(): + """The replicated K3 table shape, [163840, 7168] BF16.""" + gen = torch.Generator(device="cuda").manual_seed(0) + spec = k3_embedding_gather.CELLS[0].spec + return torch.randn(spec["vocab"], spec["hidden"], device="cuda", generator=gen).to( + torch.bfloat16 + ) + + +@certified +@pytest.mark.parametrize( + "cell", k3_embedding_gather.CELLS, ids=[c.why[:48] for c in k3_embedding_gather.CELLS] +) +def test_certified_cells(cell, k3_table): + spec = cell.spec + assert tuple(k3_table.shape) == (spec["vocab"], spec["hidden"]) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + rows = spec["rows"] + ids = torch.randint(0, spec["vocab"], (rows,), device="cuda", generator=gen).to(torch.int32) + ids[0] = spec["vocab"] - 1 + if spec["padded"]: + ids[1::3] = -1 + with validating(k3_embedding_gather): + out = k3_embedding_gather(ids, k3_table) + ref = k3_embedding_gather.reference(ids, k3_table) + assert out.shape == (rows, spec["hidden"]) and out.dtype == torch.bfloat16 + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), cell.why + if spec["padded"]: + assert not out[1::3].view(torch.int16).any(), f"{cell.why}: padded ids are not zero rows" + + +@pytest.mark.parametrize("num_ids", [0, 1, 7, 14, 32]) +def test_embedding_gather_matches_embedding(num_ids): + gen = torch.Generator(device="cuda").manual_seed(num_ids) + weight = torch.randn(5000, 7168, device="cuda", generator=gen).to(torch.bfloat16) + ids = torch.randint(0, 5000, (num_ids,), device="cuda", generator=gen).to(torch.int32) + if num_ids: + ids[0] = 4999 + out = k3_embedding_gather(ids, weight) + torch.cuda.synchronize() + assert out.shape == (num_ids, 7168) + assert torch.equal(out.view(torch.int16), F.embedding(ids, weight).view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_fp4_mla_packed_reduce.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_fp4_mla_packed_reduce.py new file mode 100644 index 000000000000..e8ed7b4368b5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_fp4_mla_packed_reduce.py @@ -0,0 +1,663 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""K3's FP4 MLA decode task plan and split reduction (``k3_rubin/kernels/fp4_mla_packed_reduce``). + +The host planners (the rule, wave and makespan (LPT) plans, the task order and the modelled +makespan), the CuTe DSL plan kernel against them (optionally fused with the step's generation +lengths), and the split reduction against an FP64 reference, eager and replayed. +""" + +import math + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="requires an SM100-or-newer CUDA device", +) + + +def _plan_mod(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + fp4_mla_packed_reduce, + ) + + return fp4_mla_packed_reduce + + +@pytest.mark.parametrize( + "tiles,capacity,wave_tasks,expected_tiles_per_task,expected_begins", + [ + # Even batch under one wave (helix16 B4, 6 m-tiles -> 17 wave tasks): one wave of 32-tile tasks. + ([128, 128, 128, 128], 24, 17, 32, [0, 4, 8, 12, 16]), + # One long request among short ones takes most of the slots (two-wave target, fit by the fallback share). + ([1, 79, 3, 176], 24, 17, 13, [0, 1, 8, 9, 23]), + # Full machine, longest 1.5x above the wave share: two-wave target, the short requests also halve. + ([128] * 15 + [200], 48, 17, 67, list(range(0, 32, 2)) + [33]), + # Full machine, longest within 1.5x of the wave share: no split at all. + ([128] * 15 + [130], 48, 17, 130, list(range(17))), + # Strong skew in a full machine: the long request gets twelve tasks. + ([128] * 15 + [1024], 48, 17, 86, list(range(0, 32, 2)) + [42]), + # Strong skew in a small batch: two-wave target, fit by the fallback share. + ([128, 128, 128, 1024], 24, 17, 71, [0, 2, 4, 6, 21]), + # A single request fills one wave with 16 tasks, or honours the minimum. + ([128], 16, 17, 8, [0, 16]), + ([8], 16, 17, 2, [0, 4]), + # No spare slots: one task per request, however long. + ([5, 900], 2, 17, 900, list(range(3))), + ], +) +def test_plan_packed_tasks_balances_by_kv_length( + tiles: list[int], + capacity: int, + wave_tasks: int, + expected_tiles_per_task: int, + expected_begins: list[int], +) -> None: + tiles_per_task, begins = _plan_mod().plan_packed_tasks(tiles, capacity, wave_tasks) + assert tiles_per_task == expected_tiles_per_task + assert begins == expected_begins + assert begins[-1] <= capacity + for request_tiles, begin, end in zip(tiles, begins, begins[1:]): + assert end > begin + assert (end - begin - 1) * tiles_per_task < request_tiles <= (end - begin) * tiles_per_task + + +def test_plan_packed_tasks_one_wave_for_helix() -> None: + plan = _plan_mod().plan_packed_tasks + # helix16 bs16 (6 m-tiles, 17 wave tasks), per-rank KV skewed 3.5x: the two-wave + # target cuts 12-tile pieces (42 tasks); a Helix rank keeps one wave of 28-tile tasks. + tiles = [16] * 12 + [40, 48, 56, 56] + assert plan(tiles, 48, 17) == (12, list(range(0, 26, 2)) + [28, 32, 37, 42]) + assert plan(tiles, 48, 17, one_wave=True) == (28, list(range(13)) + [14, 16, 18, 20]) + # Even Helix batches are unchanged. + assert plan([128] * 4, 24, 17, one_wave=True) == plan([128] * 4, 24, 17) + + +@pytest.mark.parametrize( + "tiles,capacity,wave_tasks,expected_tiles_per_task,expected_begins", + [ + # Even batch under one wave (helix16 B4, 6 m-tiles -> 17 wave tasks): one wave of 32-tile tasks. + ([128, 128, 128, 128], 24, 17, 32, [0, 4, 8, 12, 16]), + # One long request among short ones: 16-tile tasks; the one task past the wave is a short tail. + ([1, 79, 3, 176], 24, 17, 16, [0, 1, 6, 7, 18]), + # Full machine, longest 1.5x above the wave share: 100-tile tasks, the tails in a second wave. + ([128] * 15 + [200], 48, 17, 100, list(range(0, 34, 2))), + # Full machine, longest within 1.5x of the wave share: no split at all. + ([128] * 15 + [130], 48, 17, 130, list(range(17))), + # Strong skew in a full machine: eleven 94-tile tasks for the long request, short tails behind. + ([128] * 15 + [1024], 48, 17, 94, list(range(0, 32, 2)) + [41]), + # Strong skew in a small batch: 94-tile tasks, 11 for the long request, fill one wave. + ([128, 128, 128, 1024], 24, 17, 94, [0, 2, 4, 6, 17]), + # A single request fills one wave with 16 tasks; a short one is not worth splitting. + ([128], 16, 17, 8, [0, 16]), + ([8], 16, 17, 8, [0, 1]), + # No spare slots: one task per request, however long. + ([5, 900], 2, 17, 900, list(range(3))), + ], +) +def test_plan_packed_tasks_wave_plan_balances_by_kv_length( + tiles: list[int], + capacity: int, + wave_tasks: int, + expected_tiles_per_task: int, + expected_begins: list[int], +) -> None: + tiles_per_task, begins = _plan_mod().plan_packed_tasks( + tiles, capacity, wave_tasks, wave_plan=True + ) + assert tiles_per_task == expected_tiles_per_task + assert begins == expected_begins + assert begins[-1] <= capacity + for request_tiles, begin, end in zip(tiles, begins, begins[1:]): + assert end > begin + assert (end - begin - 1) * tiles_per_task < request_tiles <= (end - begin) * tiles_per_task + + +def test_plan_packed_tasks_simulates_helix_ranks() -> None: + plan = _plan_mod().plan_packed_tasks + # helix16 bs16 (6 m-tiles, 17 wave tasks), per-rank KV skewed 3.5x: the Helix rule targets one + # wave of 28-tile tasks (21 tasks, a second wave); the simulation keeps the 16 requests whole. + tiles = [16] * 12 + [40, 48, 56, 56] + assert plan(tiles, 48, 17, one_wave=True, wave_plan=False) == ( + 28, + list(range(13)) + [14, 16, 18, 20], + ) + assert plan(tiles, 48, 17, one_wave=True, wave_plan=True) == (56, list(range(17))) + # Even Helix batches are unchanged. + assert plan([128] * 4, 24, 17, one_wave=True, wave_plan=True) == plan( + [128] * 4, 24, 17, wave_plan=True + ) + + +def test_plan_packed_tasks_simulates_the_waves() -> None: + reduce_mod = _plan_mod() + plan = reduce_mod.plan_packed_tasks + # DEP16 rank on Rubin (3 m-tiles, 35 wave tasks), production-like skew: the two-wave target + # (93 tiles) rounds every request up to 74 tasks, a third wave; 206-tile tasks fill one wave. + tiles = [1848, 1250, 977, 743, 586, 430, 313, 196, 106] + assert plan(tiles, 9 * 9, 35, wave_plan=True) == ( + 206, + [0, 9, 16, 21, 25, 28, 31, 33, 34, 35], + ) + # Nine equal requests: 36 one-wave tasks spill one task; three 239-tile pieces each fit one wave. + assert plan([715] * 9, 9 * 9, 35, wave_plan=True) == (239, list(range(0, 28, 3))) + simulate = reduce_mod.plan_simulated_makespan + cost = reduce_mod.PLAN_TASK_COST_TILES + reduce_mod.PLAN_SPLIT_COST_TILES + # 36 equal 179-tile tasks on 35 slots: the last one waits for a whole task. + assert simulate([716] * 9, 179, 35) == 2 * (179 + cost) + # The 36th task (a 178-tile tail) starts when the first tail frees its slot. + assert simulate([715] * 9, 179, 35) == (178 + cost) + (178 + cost) + # 2-tile tails free their slots early, so a spilled tail does not extend the makespan. + assert simulate([181] * 18, 179, 35) == 179 + cost + # Requests that keep one task pay no split cost. + assert simulate([100, 300], 300, 35) == 300 + reduce_mod.PLAN_TASK_COST_TILES + + +def test_plan_task_order_runs_full_tasks_then_longest_tails() -> None: + order = _plan_mod().plan_task_order + # Tails: request 0 -> 4 tiles, request 1 -> 10 (a full task), request 2 -> 3, request 3 -> 4. + full = [(0, 0), (0, 1), (1, 0), (1, 1), (3, 0)] + tails = [(1, 2), (0, 2), (3, 1), (2, 0)] + assert order([24, 30, 3, 14], 10, wave_plan=True) == full + tails + # One task per request: longest first, ties in request order. + assert order([5, 9, 5, 2], 10, wave_plan=True) == [(1, 0), (0, 0), (2, 0), (3, 0)] + + +def test_plan_without_wave_plan_keeps_the_rule_and_request_order() -> None: + reduce_mod = _plan_mod() + plan = reduce_mod.plan_packed_tasks + # Without the wave simulation: the rule's length. + assert plan([1, 79, 3, 176], 24, 17, wave_plan=False) == (13, [0, 1, 8, 9, 23]) + assert plan([128] * 15 + [1024], 48, 17, wave_plan=False) == ( + 86, + list(range(0, 32, 2)) + [42], + ) + order = reduce_mod.plan_task_order([24, 30, 3, 14], 10, wave_plan=False) + assert order == [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (3, 0), (3, 1)] + + +def test_plan_packed_tasks_honours_the_runtime_profile_cap() -> None: + plan = _plan_mod().plan_packed_tasks + # An even full machine would not split, but the profile only covers 512 tiles. + assert plan([1024] * 16, 48, 17, max_tiles_per_task=512) == (512, list(range(0, 34, 2))) + # The cap never pushes a batch past its slots when it is ceil(longest / (splits - 1)). + assert plan([1024, 3, 1024, 700], 4 * 6, 17, max_tiles_per_task=205)[1][-1] <= 24 + + +def test_plan_packed_tasks_random_batches_fit_their_budget() -> None: + import random + + rng = random.Random(1234) + plan = _plan_mod().plan_packed_tasks + for _ in range(2000): + requests = rng.randint(1, 40) + tiles = [rng.choice([1, 2, 3, 8, 64, 128, 500, 1024, 4096]) for _ in range(requests)] + splits = rng.randint(1, 16) + wave_tasks = rng.choice([17, 53, 106]) + tiles_per_task, begins = plan(tiles, requests * splits, wave_tasks) + assert begins[-1] <= requests * splits + assert all(end > begin for begin, end in zip(begins, begins[1:])) + assert min(2, max(tiles)) <= tiles_per_task <= max(tiles) + assert all( + math.ceil(t / tiles_per_task) == e - b for t, b, e in zip(tiles, begins, begins[1:]) + ) + + +# A production-like DEP16 rank: 9 requests, 1.65M KV tokens. +_DEP16_RANK_KV_LENS = [473000, 320000, 250000, 190000, 150000, 110000, 80000, 50000, 27000] + + +@pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 7) if torch.cuda.is_available() else True, + reason="the CuTe DSL plan kernel is compiled for Rubin SM107", +) +@pytest.mark.parametrize( + "kv_lens,splits,wave_tasks,query_len,helix,wave_plan", + [ + # The rule plan: tasks in request order. + ([32768] * 15 + [262144], 4, 17, 8, False, False), + ([131, 20000, 600, 45000], 10, 17, 8, False, False), + ([131, 20000, 600, 45000], 10, 17, 8, True, False), + ([300, 70000, 2200], 16, 106, 1, False, False), + ([5, 5, 5], 2, 106, 4, False, False), + # The wave plan. + ([32768] * 15 + [262144], 4, 17, 8, False, True), + ([131, 20000, 600, 45000], 10, 17, 8, False, True), + ([131, 20000, 600, 45000], 10, 17, 8, True, True), + ([300, 70000, 2200], 16, 106, 1, False, True), + ([5, 5, 5], 2, 106, 4, False, True), + # DEP16 rank shapes (3 m-tiles, 35 wave tasks): the simulated task length. + (_DEP16_RANK_KV_LENS, 9, 35, 4, False, True), + ([183000] * 9, 9, 35, 4, False, True), + ([103000] * 16, 6, 35, 4, False, True), + # The rule plan on the same shapes, off and on Helix. + (_DEP16_RANK_KV_LENS, 9, 35, 4, False, False), + ([131, 20000, 600, 45000], 10, 17, 8, True, False), + ], +) +def test_packed_plan_kernel_matches_host_mirror( + kv_lens: list[int], splits: int, wave_tasks: int, query_len: int, helix: bool, wave_plan: bool +) -> None: + reduce_mod = _plan_mod() + requests = len(kv_lens) + tasks = requests * splits + kv = torch.tensor(kv_lens, dtype=torch.int32, device="cuda") + bounds = None + bounds_list = None + if helix: + # Per-token Helix bounds: ragged, last token of each request decides. + bounds_list = [ + max(1, kv_len - (query_len - 1 - q) * 37) + for kv_len in kv_lens + for q in range(query_len) + ] + bounds = torch.tensor(bounds_list, dtype=torch.int32, device="cuda") + tiles = reduce_mod.request_kv_tiles(kv_lens, query_len, bounds_list) + cap = max(tiles) + expected_tiles, expected_begins = reduce_mod.plan_packed_tasks( + tiles, tasks, wave_tasks, cap, one_wave=helix, wave_plan=wave_plan + ) + plan = torch.full( + (reduce_mod.packed_plan_length(requests, tasks),), -7, dtype=torch.int32, device="cuda" + ) + reduce_mod.launch_packed_plan( + kv, + plan, + num_requests=requests, + num_splits=splits, + wave_tasks=wave_tasks, + max_tiles_per_task=cap, + query_len=query_len, + helix_kv_bounds=bounds, + wave_plan=wave_plan, + ) + torch.cuda.synchronize() + plan = plan.tolist() + assert plan[0] == expected_tiles + assert plan[1] == expected_begins[-1] + assert plan[2 : 3 + requests] == expected_begins + tables = 3 + requests + request_of, split_of, count_of = ( + plan[tables : tables + tasks], + plan[tables + tasks : tables + 2 * tasks], + plan[tables + 2 * tasks : tables + 3 * tasks], + ) + # Tasks are listed in grid order; each request still owns begins[r + 1] - begins[r] of them. + order = reduce_mod.plan_task_order(tiles, expected_tiles, wave_plan) + assert len(order) == expected_begins[-1] + for task, (request, split) in enumerate(order): + assert (request_of[task], split_of[task], count_of[task]) == ( + request, + split, + expected_begins[request + 1] - expected_begins[request], + ) + # Unused slots point past every request's tiles and are never scheduled. + for task in range(expected_begins[-1], tasks): + assert request_of[task] == requests - 1 and split_of[task] * expected_tiles >= tiles[-1] + + +@pytest.mark.parametrize("helix", [False, True]) +def test_packed_plan_kernel_writes_generation_lengths(helix: bool) -> None: + reduce_mod = _plan_mod() + kv_lens = [40000, 1200, 9000, 70000] + prompt_lens = [1, 3, 1, 8] + generation_len = 8 + corrected = [kv - prompt + generation_len for kv, prompt in zip(kv_lens, prompt_lens)] + requests, splits, wave_tasks, query_len = len(kv_lens), 6, 17, generation_len + tasks = requests * splits + bounds = None + bounds_list = None + if helix: + bounds_list = [ + max(1, kv_len - (query_len - 1 - q) * 37) + for kv_len in corrected + for q in range(query_len) + ] + bounds = torch.tensor(bounds_list, dtype=torch.int32, device="cuda") + tiles = reduce_mod.request_kv_tiles(corrected, query_len, bounds_list) + cap = max(tiles) + expected_tiles, expected_begins = reduce_mod.plan_packed_tasks( + tiles, tasks, wave_tasks, cap, one_wave=helix + ) + kv = torch.tensor(kv_lens, dtype=torch.int32, device="cuda") + prompt = torch.tensor(prompt_lens, dtype=torch.int32, device="cuda") + corrected_out = torch.full((requests,), -1, dtype=torch.int32, device="cuda") + generation_out = torch.full((requests,), -1, dtype=torch.int32, device="cuda") + plan = torch.full( + (reduce_mod.packed_plan_length(requests, tasks),), -7, dtype=torch.int32, device="cuda" + ) + reduce_mod.launch_packed_plan( + kv, + plan, + num_requests=requests, + num_splits=splits, + wave_tasks=wave_tasks, + max_tiles_per_task=cap, + query_len=query_len, + helix_kv_bounds=bounds, + generation_lengths=(prompt, corrected_out, generation_out, generation_len), + ) + torch.cuda.synchronize() + # One launch writes the lengths every FP4 MLA kernel of the step reads and + # plans on the corrected ones. + assert corrected_out.tolist() == corrected + assert generation_out.tolist() == [generation_len] * requests + plan = plan.tolist() + assert plan[0] == expected_tiles + assert plan[1] == expected_begins[-1] + assert plan[2 : 3 + requests] == expected_begins + + +def test_plan_packed_tasks_lpt_launches_longest_first() -> None: + import random + + rng = random.Random(99) + mod = _plan_mod() + for _ in range(300): + requests = rng.randint(1, 40) + tiles = [ + rng.choice([1, 2, 3, 8, 43, 64, 128, 500, 704, 1024, 3516]) for _ in range(requests) + ] + splits = rng.randint(1, 16) + m_tiles = rng.choice([1, 5, 6]) + wave_clusters = rng.choice([74, 106]) + capacity = requests * splits + tiles_per_task, begins, order = mod.plan_packed_tasks_lpt( + tiles, capacity, max(1, wave_clusters // m_tiles), m_tiles, wave_clusters + ) + assert begins[-1] == len(order) <= capacity + assert all( + math.ceil(t / tiles_per_task) == e - b for t, b, e in zip(tiles, begins, begins[1:]) + ) + # Every task of every request exactly once, longest first. + assert sorted(order) == [ + (request, split) + for request in range(requests) + for split in range(begins[request + 1] - begins[request]) + ] + lengths = [min(tiles_per_task, tiles[r] - s * tiles_per_task) for r, s in order] + assert min(lengths) > 0 and lengths == sorted(lengths, reverse=True) + + +def test_plan_packed_tasks_lpt_shortens_the_modelled_schedule() -> None: + mod = _plan_mod() + mixed = [ + math.ceil(kv / 256) for kv in (20000, 40000, 80000, 100000, 150000, 200000, 250000, 600000) + ] + for tiles, m_tiles, expected_tiles in ( + # dep16 B8 MTP7 at 180K: two full tasks per request and eight 54-tile + # tails that fill the ten cluster slots the full tasks leave free. + ([704] * 8, 6, 325), + # dep16 B16 MTP5 at 175K: one task per request leaves 26 of 106 slots idle. + ([684] * 16, 5, 555), + (mixed, 6, None), + ): + capacity = 8 * len(tiles) + wave_tasks = 106 // m_tiles + tiles_per_task, _, _ = mod.plan_packed_tasks_lpt(tiles, capacity, wave_tasks, m_tiles, 106) + balanced, _ = mod.plan_packed_tasks(tiles, capacity, wave_tasks, wave_plan=False) + chosen = mod.lpt_makespan(tiles, tiles_per_task, m_tiles, 106, capacity) + assert chosen <= mod.lpt_makespan(tiles, balanced, m_tiles, 106, capacity) + if expected_tiles is not None: + assert tiles_per_task == expected_tiles + # A one-tile gain is inside the margin: helix16 B16 MTP5 keeps whole requests. + assert mod.plan_packed_tasks_lpt([43] * 16, 128, 21, 5, 106, one_wave=True)[0] == 43 + # Batches past the simulation limit keep the balanced length. + tiles = [43] * 80 + assert ( + mod.plan_packed_tasks_lpt(tiles, 160, 17, 6, 106)[0] + == mod.plan_packed_tasks(tiles, 160, 17, wave_plan=False)[0] + ) + + +@pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 7) if torch.cuda.is_available() else True, + reason="the CuTe DSL plan kernel is compiled for Rubin SM107", +) +@pytest.mark.parametrize( + "kv_lens,splits,query_len,helix,m_tiles", + [ + ([180000] * 8, 8, 8, False, 6), + ([175000] * 16, 8, 6, False, 5), + ([20000, 40000, 80000, 100000, 150000, 200000, 250000, 600000], 8, 8, False, 6), + ([131, 20000, 600, 45000], 10, 8, True, 6), + ([11000] * 32, 8, 8, True, 6), + ([300, 70000, 2200], 16, 1, False, 1), + # Past the simulation limit: balanced length, longest first. + ([5000 + 300 * i for i in range(80)], 2, 8, False, 6), + ], +) +def test_packed_plan_kernel_lpt_matches_host_mirror( + kv_lens: list[int], splits: int, query_len: int, helix: bool, m_tiles: int +) -> None: + reduce_mod = _plan_mod() + requests = len(kv_lens) + tasks = requests * splits + wave_clusters = 106 + wave_tasks = wave_clusters // m_tiles + kv = torch.tensor(kv_lens, dtype=torch.int32, device="cuda") + bounds = None + bounds_list = None + if helix: + bounds_list = [ + max(1, kv_len - (query_len - 1 - q) * 37) + for kv_len in kv_lens + for q in range(query_len) + ] + bounds = torch.tensor(bounds_list, dtype=torch.int32, device="cuda") + tiles = reduce_mod.request_kv_tiles(kv_lens, query_len, bounds_list) + cap = 640 + expected_tiles, expected_begins, expected_order = reduce_mod.plan_packed_tasks_lpt( + tiles, tasks, wave_tasks, m_tiles, wave_clusters, cap, one_wave=helix + ) + plan = torch.full( + (reduce_mod.packed_plan_length(requests, tasks),), + -7, + dtype=torch.int32, + device="cuda", + ) + reduce_mod.launch_packed_plan( + kv, + plan, + num_requests=requests, + num_splits=splits, + wave_tasks=wave_tasks, + max_tiles_per_task=cap, + query_len=query_len, + helix_kv_bounds=bounds, + lpt_m_tiles=m_tiles, + lpt_wave_clusters=wave_clusters, + ) + torch.cuda.synchronize() + plan = plan.tolist() + assert plan[0] == expected_tiles + assert plan[1] == expected_begins[-1] + assert plan[2 : 3 + requests] == expected_begins + tables = 3 + requests + request_of, split_of, count_of = ( + plan[tables : tables + tasks], + plan[tables + tasks : tables + 2 * tasks], + plan[tables + 2 * tasks : tables + 3 * tasks], + ) + for task, (request, split) in enumerate(expected_order): + assert (request_of[task], split_of[task]) == (request, split) + assert count_of[task] == expected_begins[request + 1] - expected_begins[request] + for task in range(expected_begins[-1], tasks): + assert request_of[task] == requests - 1 and split_of[task] == 1 << 15 + + +def test_request_kv_tiles_follow_helix_bounds() -> None: + tiles = _plan_mod().request_kv_tiles + assert tiles([1, 256, 257, 0], 1) == [1, 1, 2, 1] + # Helix: the last query token's bound of each request decides. + assert tiles([5000, 5000], 2, [100, 512, 5000, 1000]) == [2, 4] + + +@pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 7) if torch.cuda.is_available() else True, + reason="the CuTe DSL reducer is compiled for Rubin SM107", +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("with_stats", [False, True]) +@pytest.mark.parametrize( + "kv_tiles,splits,wave_tasks,num_heads,query_len", + [ + ([128, 128, 128, 128], 6, 17, 96, 8), + ([4096], 16, 106, 12, 8), + ([3, 300, 3], 3, 106, 6, 6), + ([7, 1024], 2, 26, 128, 4), + ([1, 1, 1, 1, 1], 1, 17, 96, 1), + # Skewed batches use an uneven number of tasks per request, and + # requests that end up with one task are stats-only here. + ([1, 79, 3, 176], 6, 17, 96, 8), + ([128] * 15 + [1024], 3, 106, 12, 8), + # Many single-task requests: stats-only rows across many reducer CTAs. + ([9] * 40, 2, 17, 96, 8), + ], +) +def test_packed_split_reduce_matches_reference_and_replays( + dtype: torch.dtype, + with_stats: bool, + kv_tiles: list[int], + splits: int, + wave_tasks: int, + num_heads: int, + query_len: int, +) -> None: + reduce_mod = _plan_mod() + torch.manual_seed(7) + requests = len(kv_tiles) + valid_m = num_heads * query_len + physical_m = (valid_m + 127) // 128 * 128 + tasks = splits * requests + # KV lengths just below the tile boundary; the kernels' plan is the mirror's plan. + kv_lens = torch.tensor([t * 256 - 3 for t in kv_tiles], dtype=torch.int32, device="cuda") + tiles_per_task, begins = reduce_mod.plan_packed_tasks(kv_tiles, tasks, wave_tasks) + plan = torch.empty( + reduce_mod.packed_plan_length(requests, tasks), dtype=torch.int32, device="cuda" + ) + reduce_mod.launch_packed_plan( + kv_lens, + plan, + num_requests=requests, + num_splits=splits, + wave_tasks=wave_tasks, + max_tiles_per_task=max(kv_tiles), + query_len=query_len, + ) + torch.cuda.synchronize() + assert plan[0].item() == tiles_per_task and plan[2 : 3 + requests].tolist() == begins + assert reduce_mod.request_kv_tiles(kv_lens.tolist(), query_len) == kv_tiles + partial = torch.randn((tasks, physical_m, 512), device="cuda") + maxima = torch.randn((tasks, physical_m), device="cuda") * 12 + sums = torch.rand((tasks, physical_m), device="cuda") * 16 + 0.5 + # Each request's first task is empty for every row; row 0 of request 0 is + # empty everywhere; unplanned task slots hold garbage that must be ignored. + for begin in begins[:-1]: + sums[begin] = 0 + sums[begins[0] : begins[1], 0] = 0 + maxima.masked_fill_(sums == 0, -torch.inf) + partial.masked_fill_((sums == 0).unsqueeze(-1), torch.nan) + sums[begins[-1] :] = 3.0 + maxima[begins[-1] :] = 1e9 + partial[begins[-1] :] = torch.nan + q_scale = torch.tensor(1.25, device="cuda") + kv_scale = torch.tensor(0.75, device="cuda") + sm_scale, output_scale = 0.1, 0.25 + output = torch.empty((requests * valid_m, 512), dtype=dtype, device="cuda") + stats = torch.empty((requests * valid_m, 2), device="cuda") if with_stats else None + single_task = torch.tensor( + [end - begin == 1 for begin, end in zip(begins, begins[1:])], device="cuda" + ).repeat_interleave(valid_m) + direct_output = torch.randn((requests * valid_m, 512), device="cuda").to(dtype) + + def launch() -> None: + # Rows of single-task requests were stored by the decode kernel itself. + output.copy_(direct_output) + reduce_mod.reduce_packed_splits( + partial, + maxima, + sums, + output, + q_scale, + kv_scale, + num_requests=requests, + num_splits=splits, + valid_m=valid_m, + physical_m=physical_m, + sm_scale=sm_scale, + output_scale=output_scale, + plan=plan, + stats=stats, + sm_count=212, + ) + + def check() -> None: + scale = sm_scale * math.log2(math.e) / (q_scale.double() * kv_scale.double()) + expected_rows = [] + max_rows = [] + sum_rows = [] + for begin, end in zip(begins, begins[1:]): + row_sum = sums[begin:end, :valid_m].double() + active = row_sum != 0 + row_max = maxima[begin:end, :valid_m].double() * scale + global_max = row_max.masked_fill(~active, -torch.inf).amax(0) + weights = torch.where(active, torch.exp2(row_max - global_max), 0.0) + values = partial[begin:end, :valid_m].double() + numerator = ( + torch.where(active.unsqueeze(-1), values, 0.0) * weights.unsqueeze(-1) + ).sum(0) + denominator = (row_sum * weights).sum(0) + expected_rows.append( + torch.where( + denominator.unsqueeze(-1) != 0, + numerator / denominator.unsqueeze(-1) * output_scale, + 0.0, + ) + ) + max_rows.append(global_max) + sum_rows.append(denominator) + expected = torch.cat(expected_rows).to(dtype) + expected = torch.where(single_task.unsqueeze(-1), direct_output, expected) + global_max = torch.cat(max_rows) + denominator = torch.cat(sum_rows) + tolerance = torch.finfo(dtype).eps + torch.testing.assert_close(output, expected, rtol=4 * tolerance, atol=tolerance) + if stats is not None: + expected_max = torch.where(denominator != 0, global_max / math.log2(math.e), -torch.inf) + torch.testing.assert_close( + stats[:, 0], expected_max.reshape(-1).float(), rtol=1e-5, atol=1e-5 + ) + torch.testing.assert_close( + stats[:, 1], denominator.reshape(-1).float(), rtol=2e-5, atol=2e-6 + ) + + launch() + torch.cuda.synchronize() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + # Replays follow the live scales and payloads at the same addresses. + q_scale.fill_(0.5) + kv_scale.fill_(2.0) + partial.mul_(0.75) + sums.mul_(1.5) + output.fill_(torch.nan) + graph.replay() + torch.cuda.synchronize() + check() diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_route_tail_sort_tables.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_route_tail_sort_tables.py new file mode 100644 index 000000000000..db758f47180e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_route_tail_sort_tables.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The CuTe routing tail's ``moe_sort`` tables (``ops.moe.k3_route_tail``, ``TRTLLM_K3_ROUTE_TAIL``), +which the gate hands to the routed experts as ``SortTables`` in place of +their own ``moe_sort``: for the tail's top-k, the same tiles as ``trtllm::moe_sort`` (count, +expert, row limit) and every routed row of this rank's experts in a slot of its expert's rows +there, with the inverse map back to the row.""" + +import pytest +import torch + +K3_RUBIN = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) +NUM_EXPERTS = 896 +TOP_K = 16 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.mark.parametrize("num_tokens", [1, 8, 32, 256]) +@pytest.mark.parametrize( + "local_expert_offset, local_num_experts", + [(0, 896), (224, 224), (672, 224), (280, 56)], +) +def test_route_tail_tables_are_moe_sort_tables(num_tokens, local_expert_offset, local_num_experts): + route_tail = pytest.importorskip(f"{K3_RUBIN}.kernels.k3_route_tail") + entry = pytest.importorskip(f"{K3_RUBIN}.ops.moe.k3_route_tail").k3_route_tail + device = torch.device("cuda") + if not route_tail.supports_device(device): + pytest.skip("k3_route_tail runs on SM100 / SM103 / SM107") + generator = torch.Generator(device=device).manual_seed(num_tokens) + logits = torch.randn( + num_tokens, NUM_EXPERTS, device=device, dtype=torch.float32, generator=generator + ) + bias = 0.1 * torch.randn(NUM_EXPERTS, device=device, dtype=torch.float32, generator=generator) + + outs = entry(logits, bias, 2.5, local_expert_offset, local_num_experts, route_tail.TILE) + topk_ids, topk_weights = outs[0], outs[1] + ref = torch.ops.trtllm.moe_sort( + token_selected_experts=topk_ids, + token_final_scales=topk_weights, + num_experts=NUM_EXPERTS, + top_k=TOP_K, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + tile_tokens_dim=route_tail.TILE, + ) + torch.cuda.synchronize() + tile_expert, tile_limit, e2p, p2e, total, num_tiles = outs[2:8] + ref_tile_expert, ref_tile_limit, ref_e2p, _, ref_total, ref_num_tiles = ref + + assert int(num_tiles.item()) == int(ref_num_tiles.item()) + assert int(total.item()) == int(ref_total.item()) + tiles = int(ref_num_tiles.item()) + assert torch.equal(tile_expert[:tiles], ref_tile_expert[:tiles]) + assert torch.equal(tile_limit[:tiles], ref_tile_limit[:tiles]) + + experts = topk_ids.reshape(-1) + local = (experts >= local_expert_offset) & (experts < local_expert_offset + local_num_experts) + rows = local.nonzero().flatten() + slots = e2p.reshape(-1)[rows].long() + ref_slots = ref_e2p.reshape(-1)[rows].long() + for expert in experts[rows].unique(): + mine = experts[rows] == expert + assert torch.equal(slots[mine].sort().values, ref_slots[mine].sort().values) + assert torch.equal(p2e[slots].long(), rows) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_small_batch_moe.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_small_batch_moe.py new file mode 100644 index 000000000000..4c5542e30e9a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_small_batch_moe.py @@ -0,0 +1,470 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi K3's copy of the small-batch NVFP4 routed-experts kernel (``ops.moe.k3_small_batch_moe``). + +Against a float64 reference built from the dequantized operands: + +* FC1 + SiTU + NVFP4 requantization + FC2 + combine: the output matches the routing-weighted + down-projections of the reference intermediate after the same NVFP4 requantization (a few + elements may round to the neighbouring FP4 value). +* The red combine (the default) is within the BF16 accumulation bound of the slice combine. +* The slice combine does not depend on the schedule (bitwise for other persistent CTA counts), + and the token-indexed layout gives the permuted layout's bits. +* Every output element is written and the counters are zero after every launch, also over CUDA + graph replays. +""" + +from typing import Dict, Optional, Tuple + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_small_batch_moe_op as _op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.moe import ( + k3_small_batch_moe as _entry, +) +from tensorrt_llm._torch.utils import unswizzle_sf +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 + +SF_VEC = 16 + +SMALL_BATCH_MOE_SITU_BETA = 2.5 +SMALL_BATCH_MOE_SITU_LINEAR_BETA = 0.7 + + +def random_nvfp4_experts( + num: int, rows: int, cols: int, gen: torch.Generator +) -> Tuple[torch.Tensor, torch.Tensor]: + """Random e2m1 weights ``[num, rows, cols / 2]`` and E4M3 block scales in + [0.5, 2) (128x4 layout).""" + w = torch.randint( + 0, 256, (num, rows, cols // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + sf = torch.rand((num, rows, cols // 16), device="cuda", generator=gen) * 1.5 + 0.5 + return w, sf.to(torch.float8_e4m3fn).view(torch.uint8) + + +def small_batch_moe_operands( + ids: torch.Tensor, + scales: torch.Tensor, + hidden: int, + inter_size: int, + num_experts: int, + num_local: int, + offset: int, + gen: torch.Generator, + weight_gen: Optional[torch.Generator] = None, +) -> Dict[str, object]: + """Operands of ``k3_small_batch_moe`` on a rank + whose ``num_local`` experts start at ``offset``, for the top-k ``ids`` / + ``scales`` ``[num_tokens, top_k]``: the ``moe_sort`` tables, NVFP4 tokens + drawn from ``gen`` and NVFP4 experts drawn from ``weight_gen`` (default + ``gen``).""" + num_tokens, top_k = ids.shape + t2g, t2lim, e2p, p2e, _, nnet = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=offset, + local_num_experts=num_local, + tile_tokens_dim=128, + ) + xb = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16, generator=gen) + x_gsf = xb.abs().max().float() / (448 * 6) + x, x_sf = torch.ops.trtllm.fp4_quantize(xb, 1 / x_gsf, 16, False, False) + weight_gen = gen if weight_gen is None else weight_gen + w1, w1_sf = random_nvfp4_experts(num_local, 2 * inter_size, hidden, weight_gen) + w2, w2_sf = random_nvfp4_experts(num_local, hidden, inter_size, weight_gen) + ramp = torch.linspace(0.5, 1.5, num_local, device="cuda") + global_sf = torch.tensor([448 * 6 / 64.0], dtype=torch.float32, device="cuda") + return dict( + hidden=hidden, + inter_size=inter_size, + num_local=num_local, + top_k=top_k, + offset=offset, + ids=ids, + scales=scales, + t2g=t2g, + t2lim=t2lim, + p2e=p2e, + e2p=e2p, + nnet=nnet, + x=x, + x_sf=x_sf.view(num_tokens, hidden // 16), + w1=w1, + w1_sf=w1_sf, + w2=w2, + w2_sf=w2_sf, + alpha1=(ramp * x_gsf * 0.004).float(), + alpha2=(ramp * 0.003 / global_sf[0]).float(), + global_sf=global_sf, + ) + + +def run_small_batch_moe( + d: Dict[str, object], + out: torch.Tensor, + counters: torch.Tensor, + *, + token_indexed: bool, + deterministic: bool, + num_ctas: int = 0, + die_major: bool = False, + published: bool = False, + push=None, +) -> Dict[str, object]: + """``k3_small_batch_moe`` on the operands ``d`` + of :func:`small_batch_moe_operands`: on their ``moe_sort`` tables or, with + ``published``, on their top-k (the routing word holds the token count + before the launch); ``push``: the ``LamportWorkspaceView`` of a one-shot + MNNVL all-reduce call the kernel pushes its output into. Returns the call's arguments.""" + tables = (d["t2g"], d["t2lim"], d["p2e"], d["nnet"]) + e2p = d["e2p"] + if published: + _op.small_batch_moe_route_word(counters).fill_(d["ids"].size(0)) + tables, e2p = (None,) * 4, d["ids"].reshape(-1) + kwargs = dict( + input=d["x"].view(torch.float4_e2m1fn_x2), + input_scale=d["x_sf"].view(torch.uint8), + fc1_weight=d["w1"].view(torch.float4_e2m1fn_x2), + fc1_weight_scale=d["w1_sf"], + fc1_alpha=d["alpha1"], + fc2_weight=d["w2"].view(torch.float4_e2m1fn_x2), + fc2_weight_scale=d["w2_sf"], + fc2_alpha=d["alpha2"], + output=out, + tile_idx_to_group_idx=tables[0], + tile_idx_to_mn_limit=tables[1], + permuted_idx_to_expanded_idx=tables[2], + expanded_idx_to_permuted_idx=e2p, + num_non_exiting_tiles=tables[3], + global_sf=d["global_sf"], + token_final_scales=d["scales"], + counters=counters, + top_k=d["top_k"], + situ_beta=SMALL_BATCH_MOE_SITU_BETA, + situ_linear_beta=SMALL_BATCH_MOE_SITU_LINEAR_BETA, + token_indexed=token_indexed, + num_ctas=num_ctas, + deterministic_combine=deterministic, + die_major=die_major, + published_routing=published, + expert_offset=d["offset"], + push_mc=push.mc_base if push is not None else 0, + push_flags=push.flags if push is not None else None, + push_rank=push.rank if push is not None else 0, + push_world=push.world if push is not None else 0, + ) + _entry.k3_small_batch_moe(**kwargs) + return kwargs + + +_E2M1 = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() not in (100, 103, 107), + reason="the kernel targets the SM100 family (dispatched on SM107)", +) + +SHAPES = { + # hidden, intermediate per rank, experts, local experts, local offset, top-k + "k3_ep16": (3584, 3072, 896, 56, 392, 16), + "k3_tp4ep4": (3584, 768, 896, 224, 224, 16), + "small": (1024, 256, 64, 16, 0, 8), +} + + +def _fp4_to_f64(packed: torch.Tensor) -> torch.Tensor: + lut = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + p = packed.view(torch.uint8).long() + return torch.stack([lut[p & 0xF], lut[p >> 4]], dim=-1).flatten(-2) + + +def _e4m3_to_f64(sf: torch.Tensor) -> torch.Tensor: + return sf.view(torch.uint8).view(torch.float8_e4m3fn).to(torch.float64) + + +def _situ(up, gate): + beta, linear_beta = SMALL_BATCH_MOE_SITU_BETA, SMALL_BATCH_MOE_SITU_LINEAR_BETA + g = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + u = linear_beta * torch.tanh(up / linear_beta) + return g * u + + +def _setup(shape, num_tokens, seed, local_tokens=True): + hidden, inter_size, num_experts, num_local, offset, top_k = SHAPES[shape] + gen = torch.Generator(device="cuda").manual_seed(seed) + logits = torch.randn(num_tokens, num_experts, device="cuda", generator=gen) + scales, ids = logits.topk(top_k, dim=-1) + ids = ids.to(torch.int32) + scales = scales.softmax(dim=-1).to(torch.float32) + if local_tokens: + ids[0, 0] = offset # at least one token on this rank + else: + ids.fill_((offset + num_local) % num_experts) # every token on another rank + return small_batch_moe_operands( + ids, scales, hidden, inter_size, num_experts, num_local, offset, gen + ) + + +def _counters(): + return _op.new_small_batch_moe_counters(torch.device("cuda")) + + +def _rows(d): + """(tile, permuted row) of every expanded slot served on this rank.""" + num_tiles = int(d["nnet"].item()) + limits = d["t2lim"][:num_tiles].tolist() + return [(t, r) for t in range(num_tiles) for r in range(t * 128, limits[t])] + + +def _dequant_w(w, sf, e, rows, cols): + lin = unswizzle_sf(sf[e].contiguous().flatten(), rows, cols, SF_VEC).view(rows, cols // SF_VEC) + return _fp4_to_f64(w[e]) * _e4m3_to_f64(lin).repeat_interleave(SF_VEC, dim=1) + + +def _reference(d, num_tokens): + """float64 FC1 + SiTU per served slot, and the routed output of those intermediates.""" + hidden, inter_size, top_k = d["hidden"], d["inter_size"], d["top_k"] + rows = _rows(d) + row_idx = torch.tensor([r for _, r in rows], device="cuda") + tiles = torch.tensor([t for t, _ in rows], device="cuda") + expanded = d["p2e"][row_idx].long() + tokens, slots = expanded // top_k, expanded % top_k + experts = d["t2g"][tiles].long() + x_deq = _fp4_to_f64(d["x"]) * _e4m3_to_f64(d["x_sf"]).repeat_interleave(SF_VEC, dim=1) + h_ref = torch.empty(len(rows), inter_size, dtype=torch.float64, device="cuda") + for e in experts.unique().tolist(): + sel = (experts == e).nonzero().flatten() + w1 = _dequant_w(d["w1"], d["w1_sf"], e, 2 * inter_size, hidden) + acc = (x_deq[tokens[sel]] @ w1.t()) * float(d["alpha1"][e]) + acc = acc.view(len(sel), -1, 2, 16) # [gate16, up16] row groups + h_ref[sel] = _situ(acc[:, :, 1], acc[:, :, 0]).reshape(len(sel), -1) + return rows, tokens, slots, experts, h_ref + + +def _nvfp4_roundtrip(h, global_sf): + """``h`` quantized like the kernel's FC1 epilogue (16-element blocks, E4M3 scales of + ``amax / 6 * global_sf``), dequantized in the scaled domain the FC2 alphas expect.""" + grid = torch.tensor(_E2M1[:8], dtype=torch.float64, device=h.device) + blocks = h.view(h.size(0), -1, SF_VEC) * global_sf + amax = blocks.abs().amax(-1, keepdim=True) + sf = (amax / 6.0).float().to(torch.float8_e4m3fn).double() + v = torch.where(sf > 0, blocks / sf, torch.zeros_like(blocks)).clamp(-6.0, 6.0) + q = grid[(v.abs()[..., None] - grid).abs().argmin(-1)] * v.sign() + return (q * sf).view_as(h) + + +def _routed_sum(d, num_tokens, h, tokens, slots, experts): + """float64 routed output of the intermediates ``h`` and the magnitude sum of its terms.""" + hidden, inter_size = d["hidden"], d["inter_size"] + partial = torch.empty(len(tokens), hidden, dtype=torch.float64, device="cuda") + for e in experts.unique().tolist(): + sel = (experts == e).nonzero().flatten() + w2 = _dequant_w(d["w2"], d["w2_sf"], e, hidden, inter_size) + partial[sel] = h[sel] @ w2.t() * float(d["alpha2"][e]) + terms = partial * d["scales"][tokens, slots].double()[:, None] + ref = torch.zeros(num_tokens, hidden, dtype=torch.float64, device="cuda") + mag = torch.zeros_like(ref) + ref.index_add_(0, tokens, terms) + mag.index_add_(0, tokens, terms.abs()) + return ref, mag + + +@pytest.mark.parametrize("shape", list(SHAPES)) +@pytest.mark.parametrize("num_tokens", [1, 5, 8, 13, 16, 27, 32]) +def test_small_batch_moe_matches_reference(shape, num_tokens): + """Permuted layout, slice combine: FC1 and the output against the float64 reference, schedule + independence, CUDA graph replays; the default red combine within the BF16 accumulation bound.""" + d = _setup(shape, num_tokens, 1234 + num_tokens) + hidden, inter_size = d["hidden"], d["inter_size"] + assert _op.small_batch_moe_supports(inter_size, hidden) + counters = _counters() + out = torch.full((num_tokens, hidden), float("nan"), dtype=torch.bfloat16, device="cuda") + run_small_batch_moe(d, out, counters, token_indexed=False, deterministic=True) + torch.cuda.synchronize() + assert int(counters.abs().sum().item()) == 0, "counters not reset" + assert bool(torch.isfinite(out.float()).all()), "output not fully written" + + rows, tokens, slots, experts, h_ref = _reference(d, num_tokens) + h_q = _nvfp4_roundtrip(h_ref, float(d["global_sf"][0])) + assert float((h_q != 0).float().mean()) > 0.5, "intermediate mostly zero: scales too small" + ref, mag = _routed_sum(d, num_tokens, h_q, tokens, slots, experts) + assert float((ref[tokens.unique()] != 0).float().mean()) > 0.9, "reference mostly zero" + rel = float((out.double() - ref).norm() / ref.norm()) + assert rel < 0.03, f"relative error vs the float64 reference {rel:.3e}" + + # schedule independence: other persistent CTA counts give the same bits + for ctas in (1, 37): + out2 = torch.full_like(out, float("nan")) + run_small_batch_moe( + d, out2, counters, token_indexed=False, deterministic=True, num_ctas=ctas + ) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), out2.view(torch.int16)), f"ctas={ctas}" + assert int(counters.abs().sum().item()) == 0 + + # CUDA graph replays + out_g = torch.empty_like(out) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_small_batch_moe(d, out_g, counters, token_indexed=False, deterministic=True) + for _ in range(3): + out_g.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), out_g.view(torch.int16)), "graph replay" + assert int(counters.abs().sum().item()) == 0, "counters not reset after a replay" + + # red combine: the slice-combine result within the BF16 accumulation bound of its terms + out_r = torch.full_like(out, float("nan")) + run_small_batch_moe(d, out_r, counters, token_indexed=False, deterministic=False) + torch.cuda.synchronize() + assert int(counters.abs().sum().item()) == 0 + bound = d["top_k"] * 2.0**-8 * mag + 2.0**-8 * out.double().abs() + 1e-30 + assert bool(((out_r.double() - out.double()).abs() <= bound).all()), "red combine" + + +@pytest.mark.parametrize("shape", list(SHAPES)) +@pytest.mark.parametrize("num_tokens", [1, 8, 13, 32]) +def test_small_batch_moe_token_indexed_matches_permuted(shape, num_tokens): + """The token-indexed layout computes the permuted layout's FP32 sums: with the slice combine + the bits are equal while both schedules split the tiles into the same slices (at most 16 + tiles), else within the BF16 rounding of the other slice order; for any CTA count and over + graph replays.""" + d = _setup(shape, num_tokens, 2468 + num_tokens) + counters = _counters() + ref = torch.full((num_tokens, d["hidden"]), float("nan"), dtype=torch.bfloat16, device="cuda") + run_small_batch_moe(d, ref, counters, token_indexed=False, deterministic=True) + same_slices = int(d["nnet"].item()) <= 16 + for ctas in (0, 1, 37): + out = torch.full_like(ref, float("nan")) + run_small_batch_moe(d, out, counters, token_indexed=True, deterministic=True, num_ctas=ctas) + torch.cuda.synchronize() + assert int(counters.abs().sum().item()) == 0, f"counters not reset (ctas={ctas})" + if same_slices: + assert torch.equal(ref.view(torch.int16), out.view(torch.int16)), f"ctas={ctas}" + else: + torch.testing.assert_close(out.float(), ref.float(), rtol=2.0**-7, atol=1e-6) + out_g = torch.empty_like(ref) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_small_batch_moe(d, out_g, counters, token_indexed=True, deterministic=False) + for _ in range(2): + out_g.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert int(counters.abs().sum().item()) == 0 + a, b = out_g.float(), ref.float() + assert float((a - b).norm() / b.norm().clamp_min(1e-30)) < 1e-2, "red combine replay" + + +@pytest.mark.parametrize("shape", list(SHAPES)) +@pytest.mark.parametrize("num_tokens", [1, 8, 13, 32]) +def test_small_batch_moe_published_routing_matches_tables(shape, num_tokens): + """Tables built in the kernel from the published top-k are moe_sort's: bitwise the + token-indexed launch on the tables with the slice combine, for any CTA count; the counters + (the routing word too) are zero after every launch, also over graph replays.""" + d = _setup(shape, num_tokens, 1357 + num_tokens) + counters = _counters() + ref = torch.full((num_tokens, d["hidden"]), float("nan"), dtype=torch.bfloat16, device="cuda") + run_small_batch_moe(d, ref, counters, token_indexed=True, deterministic=True) + for ctas in (0, 1, 37): + out = torch.full_like(ref, float("nan")) + run_small_batch_moe( + d, out, counters, token_indexed=True, deterministic=True, num_ctas=ctas, published=True + ) + torch.cuda.synchronize() + assert int(counters.abs().sum().item()) == 0, f"counters not reset (ctas={ctas})" + assert torch.equal(ref.view(torch.int16), out.view(torch.int16)), f"ctas={ctas}" + out_g = torch.empty_like(ref) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_small_batch_moe( + d, out_g, counters, token_indexed=True, deterministic=True, published=True + ) + for _ in range(2): + out_g.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert int(counters.abs().sum().item()) == 0 + assert torch.equal(ref.view(torch.int16), out_g.view(torch.int16)), "graph replay" + + +@pytest.mark.parametrize("token_indexed", [False, True]) +@pytest.mark.parametrize("deterministic", [False, True]) +def test_small_batch_moe_no_local_expert(token_indexed, deterministic): + """No token on this rank: the kernel zeroes the output.""" + d = _setup("k3_ep16", 8, 7, local_tokens=False) + assert int(d["nnet"].item()) == 0 + counters = _counters() + out = torch.full((8, d["hidden"]), 3.0, dtype=torch.bfloat16, device="cuda") + run_small_batch_moe(d, out, counters, token_indexed=token_indexed, deterministic=deterministic) + torch.cuda.synchronize() + assert int(out.abs().sum().item()) == 0 + assert int(counters.abs().sum().item()) == 0 + + +@pytest.mark.parametrize( + "cell", + _entry.k3_small_batch_moe.CELLS, + ids=[c.why[:48] for c in _entry.k3_small_batch_moe.CELLS], +) +def test_small_batch_moe_cells(cell): + """``k3_small_batch_moe.CELLS`` (K3: 896 experts, top 16, latent 3584) under ``validating``, + against the entry's FP64 ``reference``: the relative error bound of + :func:`test_small_batch_moe_matches_reference` (a few elements of the intermediate round to the + neighbouring FP4 value in the requantization), every element written, counters zero.""" + s = cell.spec + m, num_local, offset, inter = s["m"], s["local"], s["offset"], s["inter"] + gen = torch.Generator(device="cuda").manual_seed(4321 + m + offset) + logits = torch.randn(m, 896, device="cuda", generator=gen) + scales, ids = logits.topk(16, dim=-1) + ids = ids.to(torch.int32) + ids[0, 0] = offset # at least one row on this rank + scales = scales.softmax(dim=-1).to(torch.float32) + d = small_batch_moe_operands(ids, scales, 3584, inter, 896, num_local, offset, gen) + assert _op.small_batch_moe_supports(inter, 3584) + counters = _counters() + out = torch.full((m, 3584), float("nan"), dtype=torch.bfloat16, device="cuda") + with validating(_entry.k3_small_batch_moe): + kwargs = run_small_batch_moe( + d, + out, + counters, + token_indexed=s["layout"] == "token", + deterministic=True, + published=s.get("published", False), + ) + torch.cuda.synchronize() + assert int(counters.abs().sum().item()) == 0, "counters not reset" + assert bool(torch.isfinite(out.float()).all()), "output not fully written" + ref = _entry.k3_small_batch_moe.reference(**kwargs) + assert float((ref != 0).float().mean()) > 0.5, "reference mostly zero" + rel = float((out.double() - ref).norm() / ref.norm()) + assert rel < 0.03, f"relative error vs the float64 reference {rel:.3e}" diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_spec_decode_head.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_spec_decode_head.py new file mode 100644 index 000000000000..a069b5b1ea8b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3_spec_decode_head.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The private one-kernel speculative decode step head writes what main's native head writes. + +``test_head_matches_native_preprocess_inputs`` runs main's ``DecoderInputPreparer.preprocess_inputs`` +on the same step, so it also pins the preparer and metadata names the head reads. +""" + +import functools +import zlib +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.spec_decode_head import ( # noqa: E501 + SpecDecodeHead, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec.k3_spec_decode_head import ( # noqa: E501 + k3_spec_decode_head, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA GPU") + + +def _helix_metadata(kv_lens, helix_pos, tokens_per_block, cp_size, cp_rank, max_tokens, num_gen): + def buffer(dtype, fill): + return torch.full((max_tokens,), fill, dtype=dtype, device="cuda") + + return SimpleNamespace( + num_contexts=0, + num_generations=num_gen, + kv_lens_cuda=kv_lens, + helix_position_offsets=helix_pos, + helix_local_slots=buffer(torch.int32, -7), + helix_kv_bounds=buffer(torch.int32, -7), + kv_cache_manager=SimpleNamespace(tokens_per_block=tokens_per_block), + mapping=SimpleNamespace(cp_size=cp_size, cp_rank=cp_rank), + ) + + +@requires_cuda +@pytest.mark.parametrize("overlap", [True, False]) +@pytest.mark.parametrize( + "num_seqs, tokens_per_seq, tokens_per_block, cp_size", + [ + (1, 7, 32, 16), + (2, 7, 32, 16), + (4, 8, 64, 16), + (3, 7, 32, 4), + (16, 7, 32, 16), + (96, 8, 32, 8), + ], +) +def test_fused_head_matches_native(overlap, num_seqs, tokens_per_seq, tokens_per_block, cp_size): + gen = torch.Generator(device="cuda").manual_seed(num_seqs * 131 + tokens_per_block + overlap) + max_tokens = 1024 + n = num_seqs * tokens_per_seq + for cp_rank in sorted({0, 1, cp_size // 2, cp_size - 1}): + # Groups start near page boundaries, so a group straddles owners. + base = torch.randint(0, 1 << 18, (num_seqs,), device="cuda", generator=gen) + base = base // tokens_per_block * tokens_per_block + base += 2 + torch.randint(0, 3, (num_seqs,), device="cuda", generator=gen) + offsets = torch.arange(tokens_per_seq, device="cuda") + positions = (base.view(-1, 1) + offsets).view(-1).to(torch.int32) + pos_offsets = torch.randint(-2, 1, (max_tokens,), device="cuda", generator=gen) + pos_offsets = pos_offsets.to(torch.int32) + kv_offsets = torch.randint(-2, 1, (128,), device="cuda", generator=gen).to(torch.int32) + kv_lens = torch.randint(0, 1 << 16, (128,), device="cuda", generator=gen).to(torch.int32) + position_ids = torch.randint(0, 1 << 18, (1, n), device="cuda", generator=gen) + position_ids = position_ids.to(torch.int32) + helix_pos = torch.zeros(max_tokens, dtype=torch.int32, device="cuda") + helix_pos[:n] = positions + + # Native: DecoderInputPreparer.preprocess_inputs + recompute_helix_spec_buffers. + ref_pos_ids = position_ids.clone() + ref = _helix_metadata( + kv_lens.clone(), + helix_pos.clone(), + tokens_per_block, + cp_size, + cp_rank, + max_tokens, + num_seqs, + ) + if overlap: + ref_pos_ids[0, :] += pos_offsets[:n] + ref.kv_lens_cuda[:num_seqs] += kv_offsets[:num_seqs] + ref.helix_position_offsets[:n] += pos_offsets[:n] + TrtllmAttentionMetadata.recompute_helix_spec_buffers(ref, n, tokens_per_seq) + + out_pos_ids = position_ids.clone() + out = _helix_metadata( + kv_lens.clone(), + helix_pos.clone(), + tokens_per_block, + cp_size, + cp_rank, + max_tokens, + num_seqs, + ) + zero_kv = torch.ones(max_tokens, dtype=torch.bool, device="cuda") + k3_spec_decode_head( + out_pos_ids[0, :], + pos_offsets, + n, + out.kv_lens_cuda, + kv_offsets, + 0, + num_seqs, + has_overlap=overlap, + helix_positions=out.helix_position_offsets, + helix_slots=out.helix_local_slots, + helix_bounds=out.helix_kv_bounds, + helix_zero_kv=zero_kv, + num_helix=n, + tokens_per_block=tokens_per_block, + cp_size=cp_size, + cp_rank=cp_rank, + kv_bound_start=0, + tokens_per_seq=tokens_per_seq, + ) + torch.cuda.synchronize() + assert torch.equal(out_pos_ids, ref_pos_ids) + assert torch.equal(out.kv_lens_cuda, ref.kv_lens_cuda) + for name in ("helix_position_offsets", "helix_local_slots", "helix_kv_bounds"): + assert torch.equal(getattr(out, name)[:n], getattr(ref, name)[:n]), name + assert torch.equal(zero_kv[:n], ref.helix_kv_bounds[:n] == 0) + + +@requires_cuda +@pytest.mark.parametrize("num_gen_tokens", [28, 1500]) +def test_fused_head_overlap_only_chunked_range(num_gen_tokens): + """Without Helix buffers only the overlap correction runs, on any KV range.""" + gen = torch.Generator(device="cuda").manual_seed(num_gen_tokens) + position_ids = torch.randint(0, 1000, (1, 12 + num_gen_tokens), device="cuda", generator=gen) + position_ids = position_ids.to(torch.int32) + pos_offsets = torch.randint(-3, 4, (2048,), device="cuda", generator=gen).to(torch.int32) + kv_lens = torch.randint(0, 1000, (16,), device="cuda", generator=gen).to(torch.int32) + kv_offsets = torch.randint(-3, 4, (16,), device="cuda", generator=gen).to(torch.int32) + ref_pos, ref_kv = position_ids.clone(), kv_lens.clone() + ref_pos[0, 12:] += pos_offsets[:num_gen_tokens] + ref_kv[3:5] += kv_offsets[:2] + k3_spec_decode_head( + position_ids[0, 12:], + pos_offsets, + num_gen_tokens, + kv_lens, + kv_offsets, + 3, + 2, + has_overlap=True, + ) + torch.cuda.synchronize() + assert torch.equal(position_ids, ref_pos) + assert torch.equal(kv_lens, ref_kv) + + +def _step_metadata(num_ctx, num_ctx_tokens, num_gen, num_chunked, helix, gen_positions, gen): + """TRT-LLM attention metadata of a mixed step: ``num_ctx`` context rows (the last + ``num_chunked`` of them extend a context) and ``num_gen`` verify groups.""" + max_tokens = 256 + md = Mock(spec=TrtllmAttentionMetadata) + md.num_contexts, md.num_generations = num_ctx, num_gen + md.num_seqs = num_ctx + num_gen + md.num_ctx_tokens, md.num_chunked_ctx_requests = num_ctx_tokens, num_chunked + md.kv_lens_cuda = torch.randint(0, 1 << 12, (num_ctx + num_gen,), device="cuda", generator=gen) + md.kv_lens_cuda = md.kv_lens_cuda.to(torch.int32) + md.kv_cache_manager = SimpleNamespace(tokens_per_block=32) + md.mapping = SimpleNamespace(cp_size=4, cp_rank=1) + md._helix_spec_tokens_valid = helix + md.helix_position_offsets = torch.zeros(max_tokens, dtype=torch.int32, device="cuda") + md.helix_position_offsets[: gen_positions.numel()] = gen_positions + md.helix_local_slots = torch.full((max_tokens,), -7, dtype=torch.int32, device="cuda") + md.helix_kv_bounds = torch.full((max_tokens,), -7, dtype=torch.int32, device="cuda") + md.recompute_helix_spec_buffers = ( + lambda num_tokens, width: TrtllmAttentionMetadata.recompute_helix_spec_buffers( + md, num_tokens, width + ) + ) + return md + + +def _preparer(mapping, pos_offsets, kv_offsets): + """The fields of ``DecoderInputPreparer`` that ``preprocess_inputs`` and the head read.""" + preparer = SimpleNamespace( + decode_prologue=None, + guided_decoder=None, + runtime_hooks=SimpleNamespace(spec_offsets_op=None), + _config=SimpleNamespace(disable_overlap_scheduler=False), + _buffers=SimpleNamespace( + previous_pos_id_offsets_cuda=pos_offsets, previous_kv_lens_offsets_cuda=kv_offsets + ), + mapping=mapping, + get_runtime_tokens_per_gen_step=lambda d: d + 1, + ) + preparer._kv_lens_corrected_in_preprocess = functools.partial( + DecoderInputPreparer._kv_lens_corrected_in_preprocess, preparer + ) + return preparer + + +@requires_cuda +@pytest.mark.parametrize( + "num_ctx, num_gen, num_chunked, helix", + [(2, 3, 0, False), (2, 3, 0, True), (3, 0, 2, False), (0, 4, 0, True)], + ids=["mixed", "mixed_helix", "chunked", "gen_helix"], +) +def test_head_matches_native_preprocess_inputs(num_ctx, num_gen, num_chunked, helix): + """The head through its RuntimeHooks entry writes what main's + DecoderInputPreparer.preprocess_inputs writes: the KV-length range of mixed and + extend-context batches, and the helix buffers past the context rows.""" + draft_len = 3 + width = draft_len + 1 + num_ctx_tokens = 37 * num_ctx + gen_tokens = num_gen * width if num_gen else num_chunked * width + gen = torch.Generator(device="cuda").manual_seed(num_ctx * 7 + num_gen + num_chunked + helix) + base = torch.randint(0, 1 << 12, (max(num_gen, 1),), device="cuda", generator=gen) + gen_positions = (base.view(-1, 1) * 32 + 30 + torch.arange(width, device="cuda")).view(-1) + gen_positions = gen_positions[: num_gen * width].to(torch.int32) + pos_offsets = torch.randint(-2, 1, (256,), device="cuda", generator=gen).to(torch.int32) + kv_offsets = torch.randint(-2, 1, (16,), device="cuda", generator=gen).to(torch.int32) + position_ids = torch.randint( + 0, 1 << 16, (1, num_ctx_tokens + gen_tokens), device="cuda", generator=gen + ).to(torch.int32) + mapping = SimpleNamespace(has_cp_helix=lambda: helix) + + def step(seed): + g = torch.Generator(device="cuda").manual_seed(seed) + md = _step_metadata(num_ctx, num_ctx_tokens, num_gen, num_chunked, helix, gen_positions, g) + return { + "attn_metadata": md, + "position_ids": position_ids.clone(), + "input_ids": torch.zeros(num_ctx_tokens + gen_tokens, device="cuda"), + } + + native = _preparer(mapping, pos_offsets, kv_offsets) + ref = DecoderInputPreparer.preprocess_inputs( + native, step(11), enable_spec_decode=True, runtime_draft_len=draft_len + ) + head = SpecDecodeHead() + out = step(11) + assert head( + _preparer(mapping, pos_offsets, kv_offsets), + out, + enable_spec_decode=True, + runtime_draft_len=draft_len, + ) + torch.cuda.synchronize() + assert torch.equal(out["position_ids"], ref["position_ids"]) + out_md, ref_md = out["attn_metadata"], ref["attn_metadata"] + assert torch.equal(out_md.kv_lens_cuda, ref_md.kv_lens_cuda) + if helix: + n = gen_tokens + for name in ("helix_position_offsets", "helix_local_slots", "helix_kv_bounds"): + assert torch.equal(getattr(out_md, name)[:n], getattr(ref_md, name)[:n]), name + assert torch.equal(head._zero_kv[:n], ref_md.helix_kv_bounds[:n] == 0) + + +def test_head_keeps_the_native_prologue_outside_its_cases(): + preparer = SimpleNamespace( + _config=SimpleNamespace(disable_overlap_scheduler=False), + mapping=SimpleNamespace(has_cp_helix=lambda: False), + ) + head = SpecDecodeHead() + other_metadata = {"attn_metadata": SimpleNamespace(kv_cache_manager=object())} + assert not head(preparer, other_metadata, enable_spec_decode=True, runtime_draft_len=3) + assert not head(preparer, {}, enable_spec_decode=True, runtime_draft_len=3) + assert not head(preparer, other_metadata, enable_spec_decode=False, runtime_draft_len=3) + + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402, F401 — path above; fixture by name + + +def _on_archs(op) -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in op.ARCHS} + + +@pytest.mark.skipif(not _on_archs(k3_spec_decode_head), reason="not an ARCHS architecture") +@pytest.mark.parametrize( + "cell", k3_spec_decode_head.CELLS, ids=[c.why[:48] for c in k3_spec_decode_head.CELLS] +) +def test_certified_cells(cell, certifying): # noqa: F811 — fixture by name + """Every tensor the head writes equals the entry's reference, and nothing past the ranges.""" + certifying(k3_spec_decode_head) + spec = cell.spec + num_seqs, tokens_per_seq, tokens_per_block = spec["num_seqs"], spec["tokens_per_seq"], 128 + n = num_seqs * tokens_per_seq + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + + def ints(low, high, size): + return torch.randint(low, high, (size,), device="cuda", generator=gen).to(torch.int32) + + # Groups start a few tokens past a page boundary, so a group straddles owners. + base = ints(0, 1 << 11, num_seqs).long() * tokens_per_block - ints(0, tokens_per_seq, num_seqs) + base = base.clamp(min=0) + positions = (base.view(-1, 1) + torch.arange(tokens_per_seq, device="cuda")).view(-1) + capacity = n + 16 + helix_positions = torch.zeros(capacity, dtype=torch.int32, device="cuda") + helix_positions[:n] = positions.to(torch.int32) + args = ( + ints(0, 1 << 18, capacity), # position_ids + ints(-2, 1, capacity), # pos_offsets + n, + ints(0, 1 << 16, num_seqs + 8), # kv_lens + ints(-2, 1, num_seqs + 8), # kv_offsets + 0, + num_seqs, + ) + kwargs = dict( + has_overlap=spec["overlap"], + helix_positions=helix_positions, + helix_slots=torch.full((capacity,), -7, dtype=torch.int32, device="cuda"), + helix_bounds=torch.full((capacity,), -7, dtype=torch.int32, device="cuda"), + helix_zero_kv=torch.ones(capacity, dtype=torch.bool, device="cuda"), + num_helix=n, + tokens_per_block=tokens_per_block, + cp_size=spec["cp_size"], + cp_rank=spec["cp_rank"], + kv_bound_start=0, + tokens_per_seq=tokens_per_seq, + ) + expected = k3_spec_decode_head.reference(*args, **kwargs) + pos_offsets, kv_offsets = args[1].clone(), args[4].clone() + k3_spec_decode_head(*args, **kwargs) + torch.cuda.synchronize() + written = dict(position_ids=args[0], kv_lens=args[3], **kwargs) + for name, value in expected.items(): + assert torch.equal(written[name], value), name + assert torch.equal(args[1], pos_offsets) and torch.equal(args[4], kv_offsets) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3r_context_fmha.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3r_context_fmha.py new file mode 100644 index 000000000000..f987b66fedec --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3r_context_fmha.py @@ -0,0 +1,471 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 Rubin FP8 split-K context FMHA (k3_rubin/kernels/k3_context_fmha*.py). + +The old-O TMA stage regression of the kernel (with an artificially delayed correction warp), +and the ``k3_context_fmha`` entry through ``run_k3_split_k_context_partition`` against an +FP32 reference for both cumulative-KV-length dtypes, every merge op and per-rank head counts.""" + +import importlib.util +import math +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = ["../../../.."] + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="K3 FP8 context FMHA requires SM107", +) + +_PKG = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) + + +@pytest.fixture(scope="module", params=[None, 1, 3], ids=["normal", "delay-warp1", "delay-warp3"]) +def run_fmha(request, tmp_path_factory): + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_context_fmha, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.k3_context_fmha import ( # noqa: E501 + RubinFusedMultiHeadAttentionForward, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.k3_context_fmha_helpers import ( # noqa: E501 + MaskEnum, + ) + + kernel_type = RubinFusedMultiHeadAttentionForward + if request.param is not None and not hasattr(cute.arch, "nanosleep"): + # Without cute.arch.nanosleep the instrumented warp would spin on + # clock64 inside a divergent branch. + pytest.skip("the delayed-warp instrumentation needs cute.arch.nanosleep") + if request.param is not None: + # Delay a non-issuer warp immediately before its O0 wait, after all + # preceding pipeline releases. O1 can then finish before that wait. + # Instrument a temporary copy so production needs no testing hook. + source = Path(k3_context_fmha.__file__).read_text() + wait = ( + " cute.arch.mbarrier_wait(old_o_tma_args[3], phase=old_o_tma_phase)" + ) + assert source.count(wait) == 1 + delay = ( + f" if correction_warp_idx == {request.param}:\n" + " if (row_idx - thread_idx) % self.cta_tiler[0] == 0:\n" + " delay_start = cute.arch.clock64()\n" + " while cute.arch.clock64() - delay_start < 2_000_000:\n" + " cute.arch.nanosleep(sleep_time=1000)\n" + ) + source_path = tmp_path_factory.mktemp("old_o_delay") / "delayed_fmha.py" + source_path.write_text(source.replace(wait, delay + wait)) + module_name = f"{k3_context_fmha.__package__}._old_o_delay_{request.param}" + spec = importlib.util.spec_from_file_location(module_name, source_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + + class DelayedFmha(RubinFusedMultiHeadAttentionForward): + correction_epilog = module.RubinFusedMultiHeadAttentionForward.correction_epilog + + kernel_type = DelayedFmha + + fmha = kernel_type( + cutlass.Float32, + cutlass.Float32, + (128, 128), + (192, 128), + False, + MaskEnum.RESIDUAL_MASK, + False, + False, + use_tma_store=False, + split_k_dims=(128, 64), + ) + compiled = None + + def run(q, k, k_pe, v, output, stats, op): + nonlocal compiled + nq, heads, _ = q.shape + nk = k.shape[0] + cu_q = torch.tensor([0, nq], dtype=torch.int32, device="cuda") + cu_k = torch.tensor([0, nk], dtype=torch.int32, device="cuda") + scale = 1.0 / math.sqrt(192) + args = ( + q.view(1, nq, heads, 1, 192), + k.view(1, nk, heads, 1, 128), + k_pe.view(1, nk, 1, 1, 64), + v.view(1, nk, heads, 1, 128), + output.view(1, nq, heads, 1, 128), + tuple(cutlass.Int32(x) for x in (1, nq, nq, nk, heads, heads, 192, 128)), + cu_q, + cu_k, + None, + None, + stats, + op, + cutlass.Float32(scale * math.log2(math.e)), + cutlass.Float32(scale), + cutlass.Float32(1.0), + None, + None, + None, + None, + None, + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + False, + ) + if compiled is None: + compile_args = tuple( + from_dlpack( + x, assumed_align=x.element_size() if x.ndim == 1 else 16 + ).mark_layout_dynamic(leading_dim=x.ndim - 1) + if isinstance(x, torch.Tensor) + else x + for x in args + ) + compiled = cute.compile( + fmha, *compile_args, options="--opt-level=2 --ptxas-options '--uumn'" + ) + print(f"Compiled old-O regression kernel: delay_warp={request.param}", flush=True) + compiled(*args) + torch.cuda.synchronize() + + return run + + +@pytest.mark.parametrize("q_len", [256, 257, 512, 8192]) +def test_k3_context_old_o_barrier(run_fmha, q_len: int) -> None: + torch.manual_seed(q_len) + heads, kv_len = 96, 256 + + def fp8(*shape): + return (torch.randn(*shape, device="cuda") * 0.125).to(torch.float8_e4m3fn) + + q = fp8(q_len, heads, 192) + k = fp8(kv_len, heads, 128) + k_pe = fp8(kv_len, 1, 64) + v = fp8(kv_len, heads, 128) + full_k = torch.cat((k.float(), k_pe.float().expand(-1, heads, -1)), dim=-1) + scores = torch.einsum("qhd,khd->hqk", q.float(), full_k) / math.sqrt(192) + current_max = scores.amax(dim=-1).transpose(0, 1) + weights = (scores - current_max.transpose(0, 1).unsqueeze(-1)).exp() + current_sum = weights.sum(dim=-1).transpose(0, 1) + current_output = torch.einsum("hqk,khd->qhd", weights, v.float()) / current_sum[..., None] + del scores, weights, full_k + + output = torch.full((q_len, heads, 128), torch.nan, dtype=torch.bfloat16, device="cuda") + stats = torch.full((q_len, heads, 2), torch.nan, device="cuda") + op = torch.tensor([2], dtype=torch.int64, device="cuda") + run_fmha(q, k, k_pe, v, output, stats, op) + torch.testing.assert_close(output.float(), current_output, atol=2e-3, rtol=2e-2) + + op.fill_(1) + for _ in range(3): + # Give the old result substantial weight, so a wrong old-O load cannot + # hide behind a tiny merge coefficient. Each launch starts fresh. + old_output = (torch.randn_like(output.float()) * 0.125).to(torch.bfloat16) + old_max = torch.empty_like(current_max).uniform_(-0.5, 0.5) + old_sum = torch.empty_like(current_sum).uniform_(128, 384) + output.copy_(old_output) + stats[..., 0].copy_(old_max) + stats[..., 1].copy_(old_sum) + merged_max = torch.maximum(old_max, current_max) + old_weight = old_sum * (old_max - merged_max).exp() + current_weight = current_sum * (current_max - merged_max).exp() + merged_sum = old_weight + current_weight + expected = ( + old_output.float() * old_weight[..., None] + current_output * current_weight[..., None] + ) / merged_sum[..., None] + + run_fmha(q, k, k_pe, v, output, stats, op) + torch.testing.assert_close(output.float(), expected, atol=2e-3, rtol=2e-2) + torch.testing.assert_close( + stats[..., 0] + stats[..., 1].log(), + merged_max + merged_sum.log(), + atol=3e-4, + rtol=1e-5, + ) + + op.zero_() + saved_output, saved_stats = output.clone(), stats.clone() + run_fmha(q, k, k_pe, v, output, stats, op) + torch.testing.assert_close(output, saved_output, atol=0, rtol=0) + torch.testing.assert_close(stats, saved_stats, atol=0, rtol=0) + + +def _attention(heads: int) -> SimpleNamespace: + return SimpleNamespace( + is_mla_enable=True, + kv_lora_rank=512, + num_heads=heads, + num_kv_heads=heads, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_scaling=1.0, + ) + + +@pytest.mark.parametrize( + "q_len,is_causal", [(256, False), (257, False), (512, False), (8192, False), (257, True)] +) +@pytest.mark.parametrize("heads", [96, 48, 24, 12, 6, 3]) +def test_k3_context_fmha_cu_kv_dtypes(q_len: int, is_causal: bool, heads: int) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.k3_context_fmha_op import ( # noqa: E501 + CuteDSLK3ContextFmhaRunner, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.mla import ( # noqa: E501 + run_k3_split_k_context_partition, + supports_k3_fp8_split_k_context, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.k3_context_fmha import ( # noqa: E501 + k3_context_fmha, + ) + + torch.manual_seed(q_len) + q_lens = [q_len, 17] + kv_lens = [384, 128] + nq, nk = sum(q_lens), sum(kv_lens) + + attn = _attention(heads) + assert supports_k3_fp8_split_k_context(attn) + + def fp8(*shape): + return (torch.randn(*shape, device="cuda") * 0.125).to(torch.float8_e4m3fn) + + q, k_pe = fp8(nq, heads, 192), fp8(nk, 1, 64) + # Match the noncompact K/V views produced by the combined KV-B projection. + kv = fp8(nk, 2 * heads * 128) + k, v = (part.view(nk, heads, 128) for part in kv.chunk(2, dim=-1)) + cu_q = torch.tensor([0, q_len, nq], dtype=torch.int32, device="cuda") + cu_k = torch.tensor([0, kv_lens[0], nk], dtype=torch.int64, device="cuda") + output = torch.empty((nq, heads, 128), dtype=torch.bfloat16, device="cuda") + stats = torch.empty((nq, heads, 2), device="cuda") + old_output = torch.randn_like(output) * 0.125 + old_max = torch.zeros_like(stats[..., 0]) + old_sum = torch.full_like(old_max, 256) + expected, expected_lse = [], [] + expected_current, expected_current_lse = [], [] + q_start = k_start = 0 + for q_size, k_size in zip(q_lens, kv_lens): + full_k = torch.cat( + ( + k[k_start : k_start + k_size].float(), + k_pe[k_start : k_start + k_size].float().expand(-1, heads, -1), + ), + -1, + ) + scores = torch.einsum( + "qhd,khd->hqk", q[q_start : q_start + q_size].float(), full_k + ) / math.sqrt(192) + if is_causal: + # The op uses WINDOW_MASK (top-left aligned), not + # WINDOW_MASK_INFERENCE (bottom-right aligned). + mask = ( + torch.arange(k_size, device="cuda")[None, :] + > torch.arange(q_size, device="cuda")[:, None] + ) + scores.masked_fill_(mask[None], -torch.inf) + # Include the old partition directly in a stable reference merge. + cur_max = scores.amax(-1).transpose(0, 1) + merged_max = torch.maximum(cur_max, old_max[q_start : q_start + q_size]) + weights = (scores - merged_max.T[..., None]).exp() + old_weight = old_sum[q_start : q_start + q_size] * (-merged_max).exp() + merged_sum = weights.sum(-1).T + old_weight + merged = ( + torch.einsum("hqk,khd->qhd", weights, v[k_start : k_start + k_size].float()) + + old_output[q_start : q_start + q_size].float() * old_weight[..., None] + ) / merged_sum[..., None] + expected.append(merged) + expected_lse.append(merged_max + merged_sum.log()) + current_sum = weights.sum(-1).T + current_output = ( + torch.einsum("hqk,khd->qhd", weights, v[k_start : k_start + k_size].float()) + / current_sum[..., None] + ) + if is_causal: + # Short causal prefixes fit in one 128-key tile. Model the kernel's + # E4M3 P rounding (prescale=256, unrounded denominator) explicitly: + # the FP32-only reference misses this error even with 96 heads. + prefix = min(q_size, 128) + p = (scores[:, :prefix] - cur_max[:prefix].T[..., None]).exp() + p_quantized = (p * 256).to(torch.float8_e4m3fn).float() / 256 + current_output[:prefix] = ( + torch.einsum("hqk,khd->qhd", p_quantized, v[k_start : k_start + k_size].float()) + / p.sum(-1).T[..., None] + ) + expected_current.append(current_output) + expected_current_lse.append(merged_max + current_sum.log()) + q_start += q_size + k_start += k_size + expected, expected_lse = torch.cat(expected), torch.cat(expected_lse) + + def run(op, cu_kv): + run_k3_split_k_context_partition( + attn, + q.flatten(1), + k.flatten(1), + k_pe, + v.flatten(1), + output.flatten(1), + stats, + op, + cu_q, + cu_kv, + max_q_len=max(q_lens), + max_kv_len=max(kv_lens), + is_causal=is_causal, + ) + torch.cuda.synchronize() + + op = torch.ones(2, dtype=torch.int64, device="cuda") + cache_entries = {} + for dtype in (torch.int32, torch.int64, torch.int32): + output.copy_(old_output) + stats[..., 0].copy_(old_max) + stats[..., 1].copy_(old_sum) + run(op, cu_k.to(dtype)) + torch.testing.assert_close(output.float(), expected, atol=2e-3, rtol=2e-2) + torch.testing.assert_close( + stats[..., 0] + stats[..., 1].log(), expected_lse, atol=3e-4, rtol=1e-5 + ) + key = (q.dtype, heads, is_causal, dtype) + compiled = CuteDSLK3ContextFmhaRunner.kernel_cache[key] + if dtype in cache_entries: + assert compiled is cache_entries[dtype] + cache_entries[dtype] = compiled + assert cache_entries[torch.int32] is not cache_entries[torch.int64] + + # Exercise initialization and skip as well as the merge path above. + output.fill_(torch.nan) + stats.fill_(torch.nan) + op.fill_(2) + run(op, cu_k) + torch.testing.assert_close(output.float(), torch.cat(expected_current), atol=2e-3, rtol=2e-2) + torch.testing.assert_close( + stats[..., 0] + stats[..., 1].log(), + torch.cat(expected_current_lse), + atol=3e-4, + rtol=1e-5, + ) + if is_causal: + # Same logical heads through the original 96-head specialization must + # agree exactly, independently of the FP8-vs-FP32 reference tolerance. + repeats = 96 // heads + baseline_output = torch.empty((nq, 96, 128), dtype=output.dtype, device="cuda") + baseline_stats = torch.empty((nq, 96, 2), device="cuda") + k3_context_fmha( + torch.cat([q] * repeats, dim=1), + torch.cat([k] * repeats, dim=1), + k_pe, + torch.cat([v] * repeats, dim=1), + baseline_output, + cu_q, + cu_k, + baseline_stats, + op, + max(q_lens), + max(kv_lens), + 1.0 / math.sqrt(192), + True, + ) + torch.cuda.synchronize() + torch.testing.assert_close(output, baseline_output[:, :heads], atol=0, rtol=0) + torch.testing.assert_close(stats, baseline_stats[:, :heads], atol=0, rtol=0) + saved_output, saved_stats = output.clone(), stats.clone() + op.zero_() + run(op, cu_k) + torch.testing.assert_close(output, saved_output, atol=0, rtol=0) + torch.testing.assert_close(stats, saved_stats, atol=0, rtol=0) + + +def _entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.k3_context_fmha import ( # noqa: E501 + k3_context_fmha, + ) + + return k3_context_fmha + + +@pytest.mark.parametrize("cell", _entry().CELLS, ids=[c.why[:48] for c in _entry().CELLS]) +def test_k3_context_fmha_cells(cell) -> None: + """The ``k3_context_fmha`` entry's ``.CELLS`` under ``validating`` against its FP32 reference: + output within the FP8 FMHA band (wider on causal rows over few keys, where the kernel's E4M3 + rounding of P is not averaged out), max + log(sum) of the stats, and rows of op-0 requests + untouched.""" + import zlib + + from _validating import validating + + op = _entry() + spec = cell.spec + heads = 96 + torch.manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + q_lens, kv_lens = spec["q_lens"], spec["kv_lens"] + nq, nk = sum(q_lens), sum(kv_lens) + + def fp8(*shape): + return (torch.randn(*shape, device="cuda") * 0.125).to(torch.float8_e4m3fn) + + q, k_pe = fp8(nq, heads, 192), fp8(nk, 1, 64) + kv = fp8(nk, 2 * heads * 128) + k_nope, v = (part.view(nk, heads, 128) for part in kv.chunk(2, dim=-1)) + cu_q = torch.tensor([0, *q_lens], device="cuda").cumsum(0).to(torch.int32) + kv_dtype = torch.int64 if spec["cu_kv"] == "int64" else torch.int32 + cu_kv = torch.tensor([0, *kv_lens], device="cuda").cumsum(0).to(kv_dtype) + merge_op = torch.tensor(spec["merge_ops"], dtype=torch.int64, device="cuda") + output = (torch.randn((nq, heads, 128), device="cuda") * 0.125).to(torch.bfloat16) + stats = torch.empty((nq, heads, 2), device="cuda") + stats[..., 0].uniform_(-0.5, 0.5) + stats[..., 1].uniform_(128, 384) + args = ( + q, + k_nope, + k_pe, + v, + output, + cu_q, + cu_kv, + stats, + merge_op, + max(q_lens), + max(kv_lens), + 1.0 / math.sqrt(192), + spec["causal"], + ) + expected_out, expected_stats = op.reference(*args) + with validating(op): + op(*args) + torch.cuda.synchronize() + q0 = 0 + for q_len, merge in zip(q_lens, spec["merge_ops"]): + rows = slice(q0, q0 + q_len) + q0 += q_len + if merge == 0: + assert torch.equal(output[rows], expected_out[rows]), cell.why + assert torch.equal(stats[rows], expected_stats[rows]), cell.why + continue + got, ref = output[rows].float(), expected_out[rows].float() + if spec["causal"]: + # E4M3 P (3 mantissa bits): each probability within 2^-4 relative, so a row over few + # keys moves by up to 2^-4 of max |V|; rows past one 128-key tile average it out. + few = min(q_len, 128) + band = 2.0**-4 * float(v.float().abs().max()) + torch.testing.assert_close(got[:few], ref[:few], atol=band, rtol=0) + got, ref = got[few:], ref[few:] + torch.testing.assert_close(got, ref, atol=2e-3, rtol=2e-2) + torch.testing.assert_close( + stats[rows, :, 0] + stats[rows, :, 1].log(), + expected_stats[rows, :, 0] + expected_stats[rows, :, 1].log(), + atol=3e-4, + rtol=1e-5, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3r_dsl_hints.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3r_dsl_hints.py new file mode 100644 index 000000000000..1fe3c713357d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/kernels/test_k3r_dsl_hints.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 Rubin CuTe DSL scheduling hints (k3_rubin/kernels/dsl_hints.py). + +Each wrapper emits its ``cute.nvgpu`` hint only when the DSL provides all four hints, and nothing +otherwise. Both cases execute the module against a fake ``cutlass.cute``, so they need neither a +GPU nor the CuTe DSL; where the DSL is installed, the real module is checked against it.""" + +import importlib +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +_KERNELS = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels" +_HINTS = ("cfence", "warp_switch", "sched_res_busy_xu64", "reset_sched_res_busy_xu64") + + +def _dsl_hints_source() -> Path: + # Locates the package without importing it. + package_dir = importlib.util.find_spec("tensorrt_llm").submodule_search_locations[0] + return Path(package_dir, *_KERNELS.split(".")[1:], "dsl_hints.py") + + +def _load_dsl_hints(monkeypatch, nvgpu_hints): + """Execute a private copy of dsl_hints against a fake ``cutlass.cute`` whose ``nvgpu`` has + exactly ``nvgpu_hints``. Returns the module and the names of the hints it emitted.""" + emitted = [] + nvgpu = types.SimpleNamespace( + **{name: (lambda name=name: emitted.append(name)) for name in nvgpu_hints} + ) + cute = types.ModuleType("cutlass.cute") + cute.nvgpu = nvgpu + cutlass = types.ModuleType("cutlass") + cutlass.cute = cute + monkeypatch.setitem(sys.modules, "cutlass", cutlass) + monkeypatch.setitem(sys.modules, "cutlass.cute", cute) + spec = importlib.util.spec_from_file_location("_k3r_dsl_hints_under_test", _dsl_hints_source()) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module, emitted + + +@pytest.mark.parametrize( + "nvgpu_hints", [(), ("cfence",), _HINTS[:3]], ids=["public", "cfence-only", "three-of-four"] +) +def test_wrappers_emit_nothing_without_every_hint(monkeypatch, nvgpu_hints): + module, emitted = _load_dsl_hints(monkeypatch, nvgpu_hints) + + assert module.HAS_SCHEDULING_HINTS is False + for name in _HINTS: + getattr(module, name)() + assert emitted == [] + + +def test_wrappers_emit_their_hint_with_every_hint(monkeypatch): + module, emitted = _load_dsl_hints(monkeypatch, _HINTS) + + assert module.HAS_SCHEDULING_HINTS is True + for name in _HINTS: + getattr(module, name)() + assert emitted == list(_HINTS) + module.cfence() + module.cfence() + assert emitted[len(_HINTS) :] == ["cfence", "cfence"] + + +def test_installed_dsl_hints_are_shared_with_the_context_fmha_helpers(): + cute = pytest.importorskip("cutlass.cute") + dsl_hints = importlib.import_module(f"{_KERNELS}.dsl_hints") + helpers = importlib.import_module(f"{_KERNELS}.k3_context_fmha_helpers") + + assert dsl_hints.HAS_SCHEDULING_HINTS == all(hasattr(cute.nvgpu, name) for name in _HINTS) + assert helpers.HAS_SCHEDULING_HINTS is dsl_hints.HAS_SCHEDULING_HINTS + for name in _HINTS: + assert getattr(helpers, name) is getattr(dsl_hints, name) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_fp4_mla_decode.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_fp4_mla_decode.py new file mode 100644 index 000000000000..fd58945c0b5d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_fp4_mla_decode.py @@ -0,0 +1,378 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""K3's FP4 MLA generation decode (``ops/attention/_fp4_mla_decode``). + +Host side, with the plan and decode kernels stubbed: the step's plan geometry from the target's +features, and the plan key and metadata-launch count over a step (the lengths fused with the plan +by the step's first caller, the plan alone when the backend wrote the lengths, nothing after; the +backend's own plan untouched). On SM107: the decode against the FP4 MLA backend's decode of the +same cache and query rows under each target's settings, and the reduction's programmatic launch +bitwise. + +The FP4 cache helpers are those of ``test_k3_mla_q_fused.py`` (same directory). +""" + +from types import SimpleNamespace + +import pytest +import torch + +PLAN_WS = "_k3_fp4_mla_packed_plan" + + +def _decode(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention import ( # noqa: E501 + _fp4_mla_decode as fp4_mla_decode, + ) + + return fp4_mla_decode + + +class _Kernels: + """Stand-ins for K3's decode kernel module and plan / reduction module.""" + + PAGE_PLAN_PROFILE_KV = 160 * 1024 + + def __init__(self): + self.calls = [] + self.kernel = SimpleNamespace( + SMEM_P4_PAGE_PLAN_PROFILE_KV=self.PAGE_PLAN_PROFILE_KV, + # The next power of two: a compiled profile bucket at or above the request. + _select_runtime_kv_profile=lambda kv: 1 << (kv - 1).bit_length(), + ) + self.reduce = SimpleNamespace( + packed_plan_length=lambda requests, tasks: 3 + requests + 3 * tasks, + launch_packed_plan=lambda *args, **kwargs: self.calls.append(("plan", kwargs)), + ) + + +def _stub_kernels(monkeypatch): + dec = _decode() + kernels = _Kernels() + monkeypatch.setattr(dec, "_kernel_module", lambda: kernels.kernel) + monkeypatch.setattr(dec, "_reduce_module", lambda: kernels.reduce) + monkeypatch.setattr(dec, "_get_sm_count", lambda device: 212) + # The packed layout: whole 128-row tiles of a request's (token, head) rows, up to 1024. + monkeypatch.setattr( + dec, + "_fp4_mla_packed_rows_per_seq", + lambda heads, query_len: 0 + if heads * query_len > 1024 + else -(-heads * query_len // 128) * 128, + ) + monkeypatch.setattr( + dec, + "populate_fp4_mla_generation_lengths", + lambda *args, **kwargs: kernels.calls.append(("lengths", kwargs)), + ) + monkeypatch.setattr( + dec, + "materialize_fp4_mla_device_page_table", + lambda *args, **kwargs: kernels.calls.append(("page_table", kwargs)), + ) + return dec, kernels + + +def _geometry_metadata(device, num_requests=4): + state = SimpleNamespace( + device_page_table=True, + page_table_stride=3904, + generation_kv_lens=torch.zeros(num_requests, dtype=torch.int32, device=device), + # Preallocated: the workspace helper only allocates outside CUDA graph capture. + workspaces={PLAN_WS: torch.empty(4096, dtype=torch.int32, device=device)}, + ) + return SimpleNamespace(page_size=128, is_cuda_graph=True, fp4_mla_state=state) + + +@pytest.mark.cpu_only +def test_plan_geometry_follows_the_features(monkeypatch, k3_features): + dec, kernels = _stub_kernels(monkeypatch) + device = torch.device("cpu") + md = _geometry_metadata(device) + + def geometry(**kwargs): + args = dict(num_heads=96, query_len=8, num_requests=4, helix_kv_bounds=None) + args.update(kwargs) + return dec.packed_plan_geometry(md, device=device, **args) + + base = geometry() + # Helix8 MTP7, 96 heads: 6 row tiles, 17 tasks per wave of 106 clusters, 8 slots per request, + # the makespan and wave plans. + assert (base.splits, base.wave_tasks, base.task_capacity) == (8, 17, 32) + assert (base.lpt_m_tiles, base.lpt_wave_clusters, base.wave_plan) == (6, 106, True) + # 3904 pages of 128 tokens: tasks of ceil(1952 / 7) tiles, the 128K profile. + assert (base.split_kv_len, base.profile_tiles) == (279 * 256, 512) + # The plan lives in K3's workspace; the backend's plan workspace is not created. + assert base.plan.data_ptr() == md.fp4_mla_state.workspaces[PLAN_WS].data_ptr() + assert base.plan.numel() == kernels.reduce.packed_plan_length(4, 32) + assert set(md.fp4_mla_state.workspaces) == {PLAN_WS} + assert geometry().key == base.key + bounds = torch.zeros(32, dtype=torch.int32, device=device) + assert geometry(helix_kv_bounds=bounds).key != base.key + # features.fp4_mla_page_plan off (dep16): a profile past the page-plan one. + k3_features(fp4_mla_page_plan=False) + no_page_plan = geometry() + assert no_page_plan.split_kv_len == _Kernels.PAGE_PLAN_PROFILE_KV + 1 + assert no_page_plan.profile_tiles == 1024 and no_page_plan.key != base.key + k3_features(fp4_mla_page_plan=True) + # A batch that fills the machine with one task per request keeps the full profile. + wide = geometry(num_requests=16) + assert (wide.wave_tasks, wide.split_kv_len, wide.profile_tiles) == (17, 1952 * 256, 2048) + # Too many rows for the packed layout, or no generation pages: no geometry. + assert geometry(num_heads=128, query_len=9) is None + md.fp4_mla_state.page_table_stride = 0 + assert geometry() is None + + +def _lengths_metadata(device): + state = SimpleNamespace( + generation_kv_lens=torch.zeros(2, dtype=torch.int32, device=device), + generation_append_lens=torch.zeros(2, dtype=torch.int32, device=device), + generation_lengths_num_tokens=-1, + generation_lengths_num_seqs=-1, + generation_lengths_num_contexts=-1, + generation_lengths_capture_recorded=False, + packed_plan_key=None, + device_page_table=True, + device_page_table_valid=False, + ) + return SimpleNamespace( + num_contexts=0, + num_seqs=2, + num_tokens=8, + num_ctx_tokens=0, + kv_lens_cuda_runtime=torch.tensor([17, 25], dtype=torch.int32, device=device), + prompt_lens_cuda_runtime=torch.tensor([1, 1], dtype=torch.int32, device=device), + is_cuda_graph=False, + fp4_mla_state=state, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_step_writes_lengths_and_plan_once(monkeypatch): + dec, kernels = _stub_kernels(monkeypatch) + calls = kernels.calls + md = _lengths_metadata(torch.device("cuda")) + state = md.fp4_mla_state + plan = torch.zeros(16, dtype=torch.int32, device="cuda") + geometry = dec.PackedPlanGeometry(2, 4, 8, 17, 64, 16, 1024, None, plan, ("key",), 6, 106, True) + uniform = dec.uniform_generation_lengths + step = dec._step(md) + + # The step's first caller: one fused launch writes the lengths and the plan. + uniform(md, 8, 2, plan_geometry=geometry) + assert [name for name, _ in calls] == ["plan"] + kwargs = calls[0][1] + assert kwargs["generation_lengths"][3] == 4 + assert (kwargs["lpt_m_tiles"], kwargs["lpt_wave_clusters"], kwargs["wave_plan"]) == ( + 6, + 106, + True, + ) + assert step.plan_key == ("key",) + # The backend's plan key is neither read nor written. + assert state.packed_plan_key is None + # The other layers of the step launch nothing. + uniform(md, 8, 2, plan_geometry=geometry) + uniform(md, 8, 2) + assert len(calls) == 1 + # A decode planning on other bounds rewrites the plan alone. + uniform(md, 8, 2, plan_geometry=geometry._replace(key=("other",))) + assert calls[1][0] == "plan" and calls[1][1]["generation_lengths"] is None + # The next step, whose lengths the backend's cache update wrote first: the plan alone, even + # with the previous step's key. + dec.begin_step(md) + uniform(md, 8, 2, plan_geometry=geometry._replace(key=("other",))) + assert calls[2][0] == "plan" and calls[2][1]["generation_lengths"] is None + # Lengths invalidated by the backend (prepare): a new step's plan even without begin_step. A + # first caller without a plan writes the lengths alone; the decode then plans alone. + state.generation_lengths_num_tokens = -1 + uniform(md, 8, 2) + uniform(md, 8, 2, plan_geometry=geometry._replace(key=("other",))) + assert [name for name, _ in calls[3:]] == ["lengths", "plan"] + # The device page table once per forward. + dec.materialize_page_table(md) + state.device_page_table_valid = True + dec.materialize_page_table(md) + assert [name for name, _ in calls[5:]] == ["page_table"] + # Each metadata kernel counts once; the query mark records the count. + assert step.metadata_launches == 6 + dec.mark_query_launch(md) + assert step.query_launch_mark == 6 + with pytest.raises(RuntimeError): + uniform(md, 12, 3, plan_geometry=geometry) + with pytest.raises(RuntimeError): + uniform(md, 7, 2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_prelaunch_writes_the_step_metadata(monkeypatch): + dec, kernels = _stub_kernels(monkeypatch) + device = torch.device("cuda", torch.cuda.current_device()) + md = _lengths_metadata(device) + md.page_size = 128 + md.fp4_mla_state.page_table_stride = 64 + md.fp4_mla_state.workspaces = {} + # Only a generation-only step with the device page table. + md.num_contexts = 1 + assert not dec.prelaunch_generation_metadata(md, num_heads=96, use_helix=False, device=device) + md.num_contexts = 0 + assert dec.prelaunch_generation_metadata( + md, num_heads=96, use_helix=False, device=torch.device("cuda") + ) + assert [name for name, _ in kernels.calls] == ["plan", "page_table"] + assert kernels.calls[0][1]["generation_lengths"] is not None + # An index-less device reuses the plan workspace allocated on the indexed one. + assert md.fp4_mla_state.workspaces[PLAN_WS].device == device + step = dec._step(md) + assert step.plan_key is not None and step.metadata_launches == 2 + # Later calls of the step find the metadata current. + md.fp4_mla_state.device_page_table_valid = True + assert dec.prelaunch_generation_metadata(md, num_heads=96, use_helix=False, device=device) + assert len(kernels.calls) == 2 + + +# --------------------------------------------------------------------------- +# SM107: the decode itself. +# --------------------------------------------------------------------------- + +# Each target's decode settings (k3_rubin/features.py profiles). +_TARGETS = { + "helix8": dict(fp4_mla_reduce_pdl=True), + "ull": dict(fp4_mla_reduce_pdl=True, fp4_mla_f16_partials=True), + "dep16": dict(fp4_mla_page_plan=False), +} + +_SEQ_LENS = { + "n1": [3100], + "n4": [3100, 2900, 3000, 3300], + "n16": [3000 + 37 * i for i in range(16)], + # Long requests take several tasks, launched before the short tails. + "skew": [131, 20000, 600, 45000], +} + + +def _decode_case(seq_lens, seed, helix): + """Real FP4 cache with this step's query rows written by the cache update; returns + (manager, metadata, qf module).""" + import test_k3_mla_q_fused as qf + + if qf.k3_decode_unavailable(): + pytest.skip("needs an SM107 GPU and a CuTe DSL that builds the FP4 MLA decode kernels") + n = len(seq_lens) + mgr, md = qf._fp4_case(n, seq_lens, seed) + num_tokens = n * qf.Q_LEN + rot = qf._rotary_table(md.max_seq_len, "identity") + gen = torch.Generator(device="cuda").manual_seed(seed) + latent = ( + torch.randn(num_tokens, qf.LORA + qf.ROPE, device="cuda", generator=gen) * 0.5 + ).bfloat16() + q_input = torch.randn( + num_tokens, qf.H, qf.LORA + qf.ROPE, device="cuda", generator=gen + ).bfloat16() + q_pe = q_input[..., qf.LORA :].clone() + qf._scatter(md, latent, rot, q_pe, q_input[..., qf.LORA :], q_input) + torch.cuda.synchronize() + if not helix: + # The cache update of a multi-token step needs the per-token positions, so the case + # drops them only for the decodes. + md._helix_spec_tokens_valid = False + return mgr, md, qf + + +def _run(md, qf, decode, helix): + """One decode of every query row on fresh outputs; returns (output, stats or None).""" + state = md.fp4_mla_state + num_tokens = md.num_tokens + q = torch.empty(num_tokens, qf.H, qf.LORA + qf.ROPE, device="cuda", dtype=torch.bfloat16) + out = torch.full((num_tokens, qf.H, qf.LORA), float("nan"), device="cuda", dtype=torch.bfloat16) + stats = None + if helix: + stats = torch.full((num_tokens, qf.H, 2), float("nan"), device="cuda") + decode( + md, + 0, + 0, + q, + out, + sm_scale=0.1, + kv_lora_rank=qf.LORA, + qk_rope_head_dim=qf.ROPE, + prequantized_q=state.prequantized_q, + prequantized_q_sf=state.prequantized_q_sf, + q_batch_capacity=state.q_batch_capacity, + softmax_stats_tensor=stats, + ) + torch.cuda.synchronize() + return out.clone(), None if stats is None else stats.clone() + + +def _assert_close_rows(got, ref): + """Per (token, head) row: relative L2 error and cosine against the reference.""" + got_rows = got.float().reshape(-1, got.shape[-1]) + ref_rows = ref.float().reshape(-1, ref.shape[-1]) + ref_norms = torch.linalg.vector_norm(ref_rows, dim=1) + assert bool((ref_norms > 1e-3).all()) + relative = torch.linalg.vector_norm(got_rows - ref_rows, dim=1) / ref_norms + cosine = (got_rows * ref_rows).sum(1) / (torch.linalg.vector_norm(got_rows, dim=1) * ref_norms) + assert float(relative.max()) < 0.05, float(relative.max()) + assert float(cosine.min()) > 0.998, float(cosine.min()) + + +@pytest.mark.parametrize("helix", [True, False], ids=["helix", "no-helix"]) +@pytest.mark.parametrize("seq_lens", list(_SEQ_LENS.values()), ids=list(_SEQ_LENS)) +@pytest.mark.parametrize("target", list(_TARGETS)) +def test_decode_matches_backend_decode(target, seq_lens, helix, k3_features): + """K3's decode against the FP4 MLA backend's decode (its own kernel, plan and reduction) of + the same cache and query rows, under each target's settings.""" + from tensorrt_llm._torch.attention.backends import fp4_mla + + mgr, md, qf = _decode_case(seq_lens, seed=11 + len(seq_lens), helix=helix) + try: + k3_features(**_TARGETS[target]) + ref_out, ref_stats = _run(md, qf, fp4_mla.run_fp4_mla_attention_decode, helix) + dec = _decode() + dec.begin_step(md) + out, stats = _run(md, qf, dec.run_decode, helix) + assert not torch.isnan(out).any() + _assert_close_rows(out, ref_out) + if helix: + torch.testing.assert_close(stats[..., 0], ref_stats[..., 0], atol=2e-2, rtol=2e-2) + torch.testing.assert_close(stats[..., 1], ref_stats[..., 1], atol=5e-2, rtol=2e-2) + # The backend's step plan is not K3's. + assert dec._step(md).plan_key is not None + finally: + mgr.shutdown() + + +@pytest.mark.parametrize("seq_lens", list(_SEQ_LENS.values()), ids=list(_SEQ_LENS)) +@pytest.mark.parametrize("f16_partials", [False, True], ids=["fp32", "fp16"]) +def test_reduce_pdl_is_bitwise(seq_lens, f16_partials, k3_features): + """``features.fp4_mla_reduce_pdl`` only schedules the split reduction: output and Helix + statistics bitwise, over repeated steps.""" + mgr, md, qf = _decode_case(seq_lens, seed=23 + len(seq_lens), helix=True) + dec = _decode() + try: + k3_features(fp4_mla_f16_partials=f16_partials, fp4_mla_reduce_pdl=False) + dec.begin_step(md) + ref = _run(md, qf, dec.run_decode, True) + k3_features(fp4_mla_f16_partials=f16_partials, fp4_mla_reduce_pdl=True) + for _ in range(2): + dec.begin_step(md) + got = _run(md, qf, dec.run_decode, True) + assert torch.equal(got[0].view(torch.int16), ref[0].view(torch.int16)) + assert torch.equal(got[1].view(torch.int32), ref[1].view(torch.int32)) + finally: + mgr.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_fp8_weight_read_slice.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_fp8_weight_read_slice.py new file mode 100644 index 000000000000..5b71f8d3e52f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_fp8_weight_read_slice.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Output slices of K3's FP8 weight-read projection (the Helix context heads of the shared MLA).""" + +import pytest +import torch +from utils.util import getSMVersion + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.fp8_weight_read import ( # noqa: E501 + _Fp8BlockScaleWeightReadLinear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.layers import ( + Fp8BlockScaleLinear, +) +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + +IN_FEATURES, OUT_FEATURES = 512, 1152 + + +def _projection() -> _Fp8BlockScaleWeightReadLinear: + """A checkpoint FP8 block-scale projection swapped as K3 swaps attention projections.""" + torch.manual_seed(42) + linear = Fp8BlockScaleLinear(IN_FEATURES, OUT_FEATURES, device="cuda") + weight = torch.randn(OUT_FEATURES, IN_FEATURES, device="cuda") * 0.02 + blocks = weight.view(OUT_FEATURES // 128, 128, IN_FEATURES // 128, 128) + scale = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-8) / 448.0 + expanded = scale.repeat_interleave(128, 0).repeat_interleave(128, 1) + linear.load((weight / expanded).to(torch.float8_e4m3fn), scale) + projection = _Fp8BlockScaleWeightReadLinear.from_linear(linear) + projection.post_load_weights() + return projection + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103, 107) + or (getSMVersion() == 107 and not IS_CUTLASS_DSL_RUBIN_AVAILABLE), + reason="needs a supported block-scaled FP8 GPU and toolchain", +) +@pytest.mark.parametrize("start,end", [(0, 576), (576, 1152), (513, 1107)]) +@torch.inference_mode() +def test_forward_slice_matches_the_full_projection(start, end): + """``forward_slice`` is ``forward`` cropped to the slice; the prepared scale is reused and the + projection's own weight and scale are untouched.""" + projection = _projection() + weight = projection.weight.clone() + scale = projection.weight_scale.clone() + inputs = torch.randn(2, 17, IN_FEATURES, dtype=torch.bfloat16, device="cuda") + expected = projection(inputs)[..., start:end] + + projection.prepare_output_slice(start, end) + cached = projection._output_slice_cache.get(projection.weight, start, end)[1] + actual = projection.forward_slice(inputs, start, end) + + assert actual.shape == expected.shape and actual.is_contiguous() + torch.testing.assert_close(actual, expected, atol=0.005, rtol=0.03) + assert projection._output_slice_cache.get(projection.weight, start, end)[1] is cached + assert torch.equal(projection.weight.view(torch.uint8), weight.view(torch.uint8)) + assert torch.equal(projection.weight_scale, scale) + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103, 107) + or (getSMVersion() == 107 and not IS_CUTLASS_DSL_RUBIN_AVAILABLE), + reason="needs a supported block-scaled FP8 GPU and toolchain", +) +@torch.inference_mode() +def test_prepare_output_slice_rejects_an_out_of_range_slice(): + projection = _projection() + with pytest.raises(ValueError, match="Invalid output feature range"): + projection.prepare_output_slice(0, OUT_FEATURES + 1) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_helix_tail.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_helix_tail.py new file mode 100644 index 000000000000..53935a724642 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_helix_tail.py @@ -0,0 +1,529 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi-K3 fused Helix tail (``k3_rubin/kernels/k3_helix_tail``) on one GPU (``P = 1``, a local +receive buffer). + +The fold is bitwise the packed split reduction of K3's FP4 MLA decode (latent rows and statistics, +FP32 or normalized FP16 split partials, any task count, under the rule, wave and makespan task +plans), the V up-projection is within one BF16 ulp of an FP64 reference (and of the production +BMM), and the combine is bitwise ``helix_post_process_native`` of what the senders produced, eager +and replayed in a CUDA graph whose launches alternate the receive buffer halves. +""" + +import math + +import pytest +import torch + +__extra_import_path__ = ["../../../.."] + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="requires an SM100-or-newer CUDA device", +) + +NUM_HEADS = 96 +SM_SCALE = 0.1 +OUTPUT_SCALE = 0.25 + + +def _modules(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + fp4_mla_packed_reduce as reduce_mod, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_helix_tail as tail_mod, + ) + + return reduce_mod, tail_mod + + +def _plan_kwargs(plan_kind, physical_m): + """``launch_packed_plan`` options of a task plan: ``rule`` (tasks in request order), ``wave`` + (simulated waves, full tasks first; the default with K3's decode kernel module) or ``lpt`` + (the makespan plan K3's decode runs, with the geometry it derives).""" + if plan_kind == "rule": + return {} + if plan_kind == "wave": + return {"wave_plan": True} + assert plan_kind == "lpt" + return { + "lpt_m_tiles": physical_m // 128, + "lpt_wave_clusters": max(1, torch.cuda.get_device_properties(0).multi_processor_count // 2), + } + + +def _split_case( + reduce_mod, + kv_tiles, + splits, + wave_tasks, + query_len, + seed, + partial_dtype=torch.float32, + plan_kind="wave", +): + """Random decode split state of ``len(kv_tiles)`` requests (the reduction test's recipe): + each request's first task is empty, unplanned slots hold garbage, rows of single-task + requests are stored directly in ``rows``; FP16 partials are normalized split outputs.""" + gen = torch.Generator(device="cuda") + gen.manual_seed(seed) + requests = len(kv_tiles) + valid_m = NUM_HEADS * query_len + physical_m = (valid_m + 127) // 128 * 128 + tasks = splits * requests + kv_lens = torch.tensor([t * 256 - 3 for t in kv_tiles], dtype=torch.int32, device="cuda") + plan = torch.empty( + reduce_mod.packed_plan_length(requests, tasks), dtype=torch.int32, device="cuda" + ) + reduce_mod.launch_packed_plan( + kv_lens, + plan, + num_requests=requests, + num_splits=splits, + wave_tasks=wave_tasks, + max_tiles_per_task=max(kv_tiles), + query_len=query_len, + **_plan_kwargs(plan_kind, physical_m), + ) + begins = plan[2 : 3 + requests].tolist() + partial = torch.randn((tasks, physical_m, 512), device="cuda", generator=gen) + maxima = torch.randn((tasks, physical_m), device="cuda", generator=gen) * 12 + sums = torch.rand((tasks, physical_m), device="cuda", generator=gen) * 16 + 0.5 + for begin, end in zip(begins, begins[1:]): + if end - begin > 1: + sums[begin] = 0 + maxima.masked_fill_(sums == 0, -torch.inf) + partial.masked_fill_((sums == 0).unsqueeze(-1), torch.nan) + sums[begins[-1] :] = 3.0 + maxima[begins[-1] :] = 1e9 + partial[begins[-1] :] = torch.nan + partial = partial.to(partial_dtype) + rows = torch.randn((requests * valid_m, 512), device="cuda", generator=gen).to(torch.bfloat16) + return dict( + partial=partial, + maxima=maxima, + sums=sums, + plan=plan, + q_scale=torch.tensor(1.25, device="cuda"), + kv_scale=torch.tensor(0.75, device="cuda"), + rows=rows, + requests=requests, + splits=splits, + valid_m=valid_m, + physical_m=physical_m, + ) + + +def _reduce(reduce_mod, case): + """The packed split reduction of ``case``: (latent rows, statistics).""" + rows = case["rows"].clone() + stats = torch.empty((rows.shape[0], 2), device="cuda") + reduce_mod.reduce_packed_splits( + case["partial"], + case["maxima"], + case["sums"], + rows, + case["q_scale"], + case["kv_scale"], + num_requests=case["requests"], + num_splits=case["splits"], + valid_m=case["valid_m"], + physical_m=case["physical_m"], + sm_scale=SM_SCALE, + output_scale=OUTPUT_SCALE, + plan=case["plan"], + stats=stats, + sm_count=torch.cuda.get_device_properties(0).multi_processor_count, + ) + return rows, stats + + +def _fold_inputs(tail_mod, case): + return tail_mod.FoldInputs( + case["partial"], + case["maxima"], + case["sums"], + case["plan"], + case["q_scale"], + case["kv_scale"], + num_requests=case["requests"], + valid_m=case["valid_m"], + physical_m=case["physical_m"], + sm_scale=SM_SCALE, + output_scale=OUTPUT_SCALE, + ) + + +class _Workspace: + """One rank's receive buffer (two halves) and launch counter.""" + + def __init__(self, half_bytes: int): + self.half_bytes = half_bytes + self.buffer = torch.full((2 * half_bytes // 4,), -1, dtype=torch.int32, device="cuda") + self.state = torch.zeros(8, dtype=torch.int64, device="cuda") + + def launch(self, tail_mod, rows, stats, w_uv, out, fold, debug, mx_out=None, ull2=False): + tail_mod.k3_helix_tail( + rows, + stats, + w_uv, + out, + fold=fold, + zero_kv_mask=None, + cp_rank=0, + cp_size=1, + ws_addr=self.buffer.data_ptr(), + rank_stride=self.buffer.numel() * 4, + half_bytes=self.half_bytes, + state=self.state, + max_ctas=torch.cuda.get_device_properties(0).multi_processor_count, + use_pdl=True, + mx_out=mx_out, + debug=debug, + ull2=ull2, + ) + + +def _debug_buffers(num_rows): + return ( + torch.empty((num_rows, 512), dtype=torch.bfloat16, device="cuda"), + torch.empty((num_rows, 128), dtype=torch.bfloat16, device="cuda"), + torch.empty((num_rows, 2), dtype=torch.float32, device="cuda"), + ) + + +def _check(tail_mod, latent_ref, stats_ref, w_uv, out, debug): + num_tokens = out.shape[0] + latent, values, stats = debug + # Fold: bitwise the reduction's rows and statistics. + assert torch.equal(latent.view(torch.int16), latent_ref.view(torch.int16)) + assert torch.equal(stats.view(torch.int32), stats_ref.view(torch.int32)) + # V up-projection: the FP64 product of the same BF16 rows, rounded once to BF16, up to the + # FP32 summation error (relative to the sum of |terms|, which dominates near zero). + ref = tail_mod.reference_v_up(latent_ref, w_uv) + lat = latent_ref.view(-1, NUM_HEADS, 512).double() + exact = torch.einsum("tjk,jnk->tjn", lat, w_uv.double()).reshape(-1, 128) + terms = torch.einsum("tjk,jnk->tjn", lat.abs(), w_uv.double().abs()).reshape(-1, 128) + err = (values.double() - exact).abs() + tol = tail_mod.bf16_ulp(ref).double() + terms * 2.0**-20 + assert bool((err <= tol).all()), float((err - tol).max()) + # And within two ulp of the production BMM's BF16 result (cuBLAS sums in another order). + bmm = torch.empty((num_tokens, NUM_HEADS, 128), dtype=torch.bfloat16, device="cuda") + torch.ops.trtllm.bmm_out( + latent_ref.view(num_tokens, NUM_HEADS, 512).transpose(0, 1), + w_uv.transpose(1, 2), + bmm.transpose(0, 1), + ) + diff = (values.double() - bmm.view(-1, 128).double()).abs() + tol = 2 * tail_mod.bf16_ulp(bmm.view(-1, 128)).double() + terms * 2.0**-19 + assert bool((diff <= tol).all()) + # Combine of the single source: bitwise helix_post_process_native of the sent values. + combined = torch.ops.trtllm.helix_post_process_native( + values.view(num_tokens, 1, NUM_HEADS, 128), + stats.view(num_tokens, 1, NUM_HEADS, 2), + 1.0, + 1, + ) + assert torch.equal(out.view(torch.int16), combined.reshape(num_tokens, -1).view(torch.int16)) + + +CASES = [ + # (kv tiles per request, splits, wave tasks, query length) + ([86], 8, 106, 8), # n1, 21.9k tokens: 8 tasks + ([1], 8, 106, 8), # one tile: a single task, rows stored by the decode + ([86, 60], 8, 106, 8), # n2 + ([1, 79, 3, 176], 6, 17, 8), # skewed, single-task requests between multi-task ones + ([40] * 8, 8, 106, 8), # n8 + ([20] * 16, 2, 106, 8), # n16 + ([5, 9, 3], 4, 17, 5), # tokens not a multiple of the 16-token tile + ([30, 2, 12], 4, 17, 1), # one token per request + ([86], 16, 106, 8), # n1, split 16: more tasks than one load round + ([90, 10], 16, 106, 8), # n2, split 16 +] + + +@pytest.mark.parametrize("plan_kind", ["wave", "rule", "lpt"]) +@pytest.mark.parametrize( + "partial_dtype,ull2", + [(torch.float32, False), (torch.float32, True), (torch.float16, True)], + ids=["fp32", "fp32-ull2", "fp16-ull2"], +) +@pytest.mark.parametrize("kv_tiles,splits,wave_tasks,query_len", CASES) +def test_helix_tail_fold_matches_reduction( + kv_tiles, splits, wave_tasks, query_len, partial_dtype, ull2, plan_kind +): + reduce_mod, tail_mod = _modules() + case = _split_case( + reduce_mod, + kv_tiles, + splits, + wave_tasks, + query_len, + seed=11, + partial_dtype=partial_dtype, + plan_kind=plan_kind, + ) + tasks = case["plan"][2 : 3 + case["requests"]].diff() + if splits == 16 and plan_kind == "wave": + assert int(tasks.max()) > 8, "the case must exceed one fold load round" + latent_ref, stats_ref = _reduce(reduce_mod, case) + gen = torch.Generator(device="cuda") + gen.manual_seed(5) + w_uv = (torch.randn((NUM_HEADS, 128, 512), device="cuda", generator=gen) * 0.05).to( + torch.bfloat16 + ) + num_tokens = case["requests"] * query_len + ws = _Workspace(half_bytes=num_tokens * tail_mod.record_bytes(NUM_HEADS, 128) + 4096) + out = torch.empty((num_tokens, NUM_HEADS * 128), dtype=torch.bfloat16, device="cuda") + fold = _fold_inputs(tail_mod, case) + # Two launches use both buffer halves; each re-arms what it read. + for _ in range(2): + debug = _debug_buffers(num_tokens * NUM_HEADS) + out.fill_(torch.nan) + ws.launch(tail_mod, case["rows"], None, w_uv, out, fold, debug, ull2=ull2) + torch.cuda.synchronize() + _check(tail_mod, latent_ref, stats_ref, w_uv, out, debug) + assert bool((ws.buffer == -1).all()), "the receive buffer was not re-armed" + + +@pytest.mark.parametrize("ull2", [False, True], ids=["default", "ull2"]) +def test_helix_tail_direct_rows_and_graph_replay(ull2): + """Rows stored directly (no split partials), several token counts in one CUDA graph.""" + _, tail_mod = _modules() + gen = torch.Generator(device="cuda") + gen.manual_seed(3) + w_uv = (torch.randn((NUM_HEADS, 128, 512), device="cuda", generator=gen) * 0.05).to( + torch.bfloat16 + ) + token_counts = (8, 16, 37, 64) + ws = _Workspace(half_bytes=max(token_counts) * tail_mod.record_bytes(NUM_HEADS, 128) + 4096) + + def inputs(num_tokens, seed): + g = torch.Generator(device="cuda") + g.manual_seed(seed) + rows = torch.randn((num_tokens * NUM_HEADS, 512), device="cuda", generator=g).to( + torch.bfloat16 + ) + stats = torch.empty((num_tokens * NUM_HEADS, 2), device="cuda") + stats[:, 0] = torch.randn(num_tokens * NUM_HEADS, device="cuda", generator=g) * 3 + stats[:, 1] = torch.rand(num_tokens * NUM_HEADS, device="cuda", generator=g) * 50 + 1 + return rows, stats + + cases = [inputs(n, 100 + n) for n in token_counts] + outs = [ + torch.empty((n, NUM_HEADS * 128), dtype=torch.bfloat16, device="cuda") for n in token_counts + ] + debugs = [_debug_buffers(n * NUM_HEADS) for n in token_counts] + + def launch_all(): + for (rows, stats), out, debug in zip(cases, outs, debugs): + ws.launch(tail_mod, rows, stats, w_uv, out, None, debug, ull2=ull2) + + launch_all() + torch.cuda.synchronize() + for (rows, stats), out, debug in zip(cases, outs, debugs): + _check(tail_mod, rows, stats, w_uv, out, debug) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch_all() + for rep in range(3): + for i, (rows, stats) in enumerate(cases): + new_rows, new_stats = inputs(token_counts[i], 1000 * rep + i) + rows.copy_(new_rows) + stats.copy_(new_stats) + for out in outs: + out.fill_(torch.nan) + graph.replay() + torch.cuda.synchronize() + for (rows, stats), out, debug in zip(cases, outs, debugs): + _check(tail_mod, rows, stats, w_uv, out, debug) + assert bool((ws.buffer == -1).all()) + assert math.isfinite(float(outs[0].float().abs().max())) + + +@pytest.mark.parametrize("num_tokens", [8, 37, 64]) +def test_helix_tail_mxfp8_operand(num_tokens): + """o_proj's operand: the receivers' MXFP8 form of ``bf16(out * gate)`` is bitwise the + gated low-M GEMM's quantization of it, with a row-strided gate.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_fp8w_lowm_mxfp8 as mx, + ) + + _, tail_mod = _modules() + gen = torch.Generator(device="cuda") + gen.manual_seed(num_tokens) + w_uv = (torch.randn((NUM_HEADS, 128, 512), device="cuda", generator=gen) * 0.05).to( + torch.bfloat16 + ) + rows = torch.randn((num_tokens * NUM_HEADS, 512), device="cuda", generator=gen).to( + torch.bfloat16 + ) + stats = torch.empty((num_tokens * NUM_HEADS, 2), device="cuda") + stats[:, 0] = torch.randn(num_tokens * NUM_HEADS, device="cuda", generator=gen) * 3 + stats[:, 1] = torch.rand(num_tokens * NUM_HEADS, device="cuda", generator=gen) * 50 + 1 + k = NUM_HEADS * 128 + gate_storage = torch.sigmoid( + torch.randn((num_tokens, k + 64), device="cuda", generator=gen) + ).to(torch.bfloat16) + gate = gate_storage[:, :k] + xq = torch.empty((num_tokens, k), dtype=torch.float8_e4m3fn, device="cuda") + xs = torch.empty((num_tokens, k // 32), dtype=torch.uint8, device="cuda") + ws = _Workspace(half_bytes=num_tokens * tail_mod.record_bytes(NUM_HEADS, 128) + 4096) + out = torch.empty((num_tokens, k), dtype=torch.bfloat16, device="cuda") + debug = _debug_buffers(num_tokens * NUM_HEADS) + ws.launch(tail_mod, rows, stats, w_uv, out, None, debug, mx_out=(gate, xq, xs)) + torch.cuda.synchronize() + _check(tail_mod, rows, stats, w_uv, out, debug) + ref_q, ref_s = mx.mxfp8_quantize_reference(out * gate) + assert torch.equal(xq.view(torch.uint8), ref_q.view(torch.uint8)) + assert torch.equal(xs, ref_s) + + +def test_tail_output_buffer_is_the_generation_output(): + """The fused tail writes ``forward_absorption_generation``'s output in place: it is given that + output viewed as ``[T, H / P * v_head_dim]`` (the same storage, so the copy into ``output`` is + skipped), and nothing it could not write as is (no output, another dtype, size or layout).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.mla import ( # noqa: E501 + tail_output_buffer, + ) + + num_tokens, cols = 90, 12 * 128 + assert tail_output_buffer(None, num_tokens, cols) is None + + storage = torch.empty((num_tokens + 6, cols), dtype=torch.bfloat16, device="cuda") + output = storage[:num_tokens] + buf = tail_output_buffer(output, num_tokens, cols) + assert buf is not None and buf.shape == (num_tokens, cols) + assert buf.data_ptr() == output.data_ptr() + + assert tail_output_buffer(output, num_tokens - 6, cols) is None + assert tail_output_buffer(output.float(), num_tokens, cols) is None + wide = torch.empty((num_tokens, cols + 64), dtype=torch.bfloat16, device="cuda") + assert tail_output_buffer(wide[:, :cols], num_tokens, cols) is None + + +def _entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.k3_helix_tail import ( # noqa: E501 + k3_helix_tail, + ) + + return k3_helix_tail + + +def _mx_dequant(q: torch.Tensor, s: torch.Tensor) -> torch.Tensor: + scale = torch.exp2(s.float() - 127).repeat_interleave(32, dim=-1) + return q.float() * scale + + +@pytest.mark.parametrize("cell", _entry().CELLS, ids=[c.why[:48] for c in _entry().CELLS]) +def test_k3_helix_tail_cells(cell): + """The ``k3_helix_tail`` entry's ``.CELLS`` under ``validating``: ``out`` within BF16 rounding + of the FP64 reference (fold, V up-projection, combine over the P sources of the looped-back + group), and o_proj's MXFP8 operand within one E4M3 step of the reference's.""" + import zlib + + from _validating import assert_within_ulp, validating + + op = _entry() + spec = cell.spec + reduce_mod, tail_mod = _modules() + cp_size = spec["cp_size"] + num_tokens = len(spec["kv_tiles"]) * spec["query_len"] + seed = zlib.crc32(cell.why.encode()) % 2**31 + gen = torch.Generator(device="cuda") + gen.manual_seed(seed) + if spec["partial"] is None: + rows = torch.randn((num_tokens * NUM_HEADS, 512), device="cuda", generator=gen).to( + torch.bfloat16 + ) + stats = torch.empty((num_tokens * NUM_HEADS, 2), device="cuda") + stats[:, 0] = torch.randn(num_tokens * NUM_HEADS, device="cuda", generator=gen) * 3 + stats[:, 1] = torch.rand(num_tokens * NUM_HEADS, device="cuda", generator=gen) * 50 + 1 + fold = None + else: + dtype = torch.float16 if spec["partial"] == "fp16" else torch.float32 + case = _split_case( + reduce_mod, + spec["kv_tiles"], + spec["splits"], + spec["wave_tasks"], + spec["query_len"], + seed=seed, + partial_dtype=dtype, + plan_kind="lpt", + ) + if dtype == torch.float16: + # Normalized split outputs are bounded by the values they average. + case["partial"] = (case["partial"].float() * 0.25).to(dtype) + rows, stats = case["rows"], torch.empty((num_tokens * NUM_HEADS, 2), device="cuda") + fold = _fold_inputs(tail_mod, case) + w_uv = (torch.randn((NUM_HEADS, 128, 512), device="cuda", generator=gen) * 0.05).to( + torch.bfloat16 + ) + cols = NUM_HEADS // cp_size * 128 + out = torch.full((num_tokens, cols), torch.nan, dtype=torch.bfloat16, device="cuda") + mx_out = None + if spec["mx_out"]: + gate = torch.sigmoid(torch.randn((num_tokens, cols), device="cuda", generator=gen)).to( + torch.bfloat16 + ) + mx_out = ( + gate, + torch.empty((num_tokens, cols), dtype=torch.float8_e4m3fn, device="cuda"), + torch.empty((num_tokens, cols // 32), dtype=torch.uint8, device="cuda"), + ) + half_bytes = num_tokens * cp_size * tail_mod.record_bytes(NUM_HEADS // cp_size, 128) + 4096 + buffer = torch.full((2 * half_bytes // 4,), -1, dtype=torch.int32, device="cuda") + args = ( + rows, + stats, + w_uv, + out, + fold, + None, + 0, + cp_size, + buffer.data_ptr(), + buffer.numel() * 4, + half_bytes, + torch.zeros(8, dtype=torch.int64, device="cuda"), + torch.cuda.get_device_properties(0).multi_processor_count, + True, + mx_out, + None, + cp_size > 1, + False, + spec["ull2"], + ) + # Two launches use both receive halves; each must re-arm what it read. + for _ in range(2): + out.fill_(torch.nan) + with validating(op): + op(*args) + torch.cuda.synchronize() + expected = op.reference(*args) + ref_out = expected[0] if mx_out is not None else expected + # The V up-projection rounds each source row once to BF16 (half an ulp of the row's + # scale), the fold's latent rows another half ulp; the combine adds one rounding. + assert_within_ulp(out, ref_out, element_ulp=4, rms_ulp=1) + if mx_out is not None: + got = _mx_dequant(mx_out[1], mx_out[2]) + ref = _mx_dequant(expected[1], expected[2]) + # One E4M3 step (2^-3 of the group's power-of-two scale) where out differs by an ulp. + group = ref.abs().view(num_tokens, -1, 32).amax(-1, keepdim=True) + step = (group * 2.0**-3).expand(-1, -1, 32).reshape(num_tokens, -1) + assert bool(((got - ref).abs() <= step + 1e-30).all()), cell.why + assert bool((buffer == -1).all()), "the receive buffer was not re-armed" diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_ctx_fp8.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_ctx_fp8.py new file mode 100644 index 000000000000..dc01797a7c28 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_ctx_fp8.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MLA context: the Triton static FP8 quantization the split-K context path uses +(``ctx_split_k``) against the C++ one, and the ``static_quantize_e4m3`` entry on its cells.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../.."] + +NOPE, ROPE, V, LORA = 128, 64, 128, 512 + +requires_fp8_cvt = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9), + reason="E4M3 conversion instructions require SM89+", +) + + +# --------------------------------------------------------------------------------------------- +# GPU: Triton FP8 kernels against the C++ static FP8 quantization +# --------------------------------------------------------------------------------------------- + + +def _cpp_quant(x, scale): + return torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x.contiguous(), scale)[0] + + +def _quant_test_input(rows, cols, seed): + g = torch.Generator(device="cuda") + g.manual_seed(seed) + x = torch.randn(rows, cols, generator=g, device="cuda") * 64 + flat = x.view(-1) + special = [ + float("inf"), + float("-inf"), + float("nan"), + 1e30, + -1e30, + 1e-30, + -0.0, + 0.0, + 447.0, + 448.0, + 456.0, + 464.0, + 465.0, + 480.0, + -464.0, + -465.0, + 2.0**-9, + 2.0**-10, + ] + flat[: len(special)] = torch.tensor(special, device="cuda") + n = min(flat.numel() - len(special), 8192) + flat[len(special) : len(special) + n] = torch.logspace(-12, 12, n, device="cuda") * torch.sign( + torch.randn(n, generator=g, device="cuda") + ) + return x.to(torch.bfloat16) + + +@requires_fp8_cvt +@pytest.mark.parametrize("shape", [(1, 18432), (257, 18432), (4096, 192), (3, 6)]) +@pytest.mark.parametrize("scale", [1.0, 0.5, 3.0, 0.37]) +def test_triton_static_quant_matches_cpp(shape, scale): + import tensorrt_llm # noqa: F401 + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + fp8_static_quant, + ) + + x = _quant_test_input(*shape, seed=shape[0]) + s = torch.tensor([scale], dtype=torch.float32, device="cuda") + y = fp8_static_quant.static_quantize_e4m3(x, s) + assert y.shape == x.shape and y.dtype == torch.float8_e4m3fn + assert torch.equal(y.view(torch.uint8), _cpp_quant(x, s).view(torch.uint8)) + + +@requires_fp8_cvt +@pytest.mark.parametrize("rows", [1, 63, 64, 1000, 16384]) +@pytest.mark.parametrize("scale", [1.0, 0.37]) +def test_triton_static_quant_rows_matches_cpp(rows, scale): + """Row-strided views (the K-pe columns of latent_cache) quantize to the C++ bytes.""" + import tensorrt_llm # noqa: F401 + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + fp8_static_quant, + ) + + latent = _quant_test_input(rows, LORA + ROPE, seed=rows) + s = torch.tensor([scale], dtype=torch.float32, device="cuda") + for view in (latent[:, LORA:], latent[:, LORA:].contiguous(), latent[:, :ROPE]): + assert fp8_static_quant.supports_rows(view) + y = fp8_static_quant.static_quantize_e4m3_rows(view, s) + assert y.is_contiguous() and y.shape == view.shape + assert torch.equal(y.view(torch.uint8), _cpp_quant(view, s).view(torch.uint8)) + # Odd storage offsets, non-power-of-two widths and non-BF16 inputs are not supported. + for view in (latent[:, 511:575], latent[:, LORA : LORA + 48], latent[:, LORA:].float()): + assert not fp8_static_quant.supports_rows(view) + with pytest.raises(ValueError): + fp8_static_quant.static_quantize_e4m3_rows(view, s) + + +# --------------------------------------------------------------------------------------------- +# GPU: the static_quantize_e4m3 entry on its cells +# --------------------------------------------------------------------------------------------- + + +def _entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.static_quantize_e4m3 import ( # noqa: E501 + static_quantize_e4m3, + ) + + return static_quantize_e4m3 + + +@requires_fp8_cvt +@pytest.mark.parametrize("cell", _entry().CELLS, ids=[c.why[:48] for c in _entry().CELLS]) +def test_static_quantize_e4m3_cells(cell): + """The entry's ``.CELLS`` under ``validating``: bitwise its reference (one IEEE FP32 + reciprocal, saturating round-to-nearest-even E4M3) and the C++ quantizer.""" + from _validating import validating + + op = _entry() + spec = cell.spec + tokens = spec["tokens"] + width = {"q": 96 * (NOPE + ROPE), "latent_k_pe": LORA + ROPE, "k_pe": ROPE}[spec["source"]] + x = _quant_test_input(tokens, width, seed=tokens + width) + if spec["source"] == "latent_k_pe": + x = x[:, LORA:] + scale = torch.ones(1, dtype=torch.float32, device="cuda") + with validating(op): + y = op(x, scale, rows=spec["rows"]) + assert y.dtype == torch.float8_e4m3fn and y.shape == x.shape and y.is_contiguous() + expected = op.reference(x, scale, rows=spec["rows"]) + assert torch.equal(y.view(torch.uint8), expected.view(torch.uint8)), cell.why + assert torch.equal(y.view(torch.uint8), _cpp_quant(x, scale).view(torch.uint8)), cell.why diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_ctx_split_k.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_ctx_split_k.py new file mode 100644 index 000000000000..d39872587d86 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_ctx_split_k.py @@ -0,0 +1,700 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MLA cached-prefix context on the Rubin FP8 split-K context FMHA +(k3_rubin/mla/ctx_split_k.py). + +CPU: which layers and calls take the split-K path and when kv_b_proj takes the MXFP8 cache +gather. GPU (SM89+): the partition schedule handed to the FMHA. SM107: the MXFP8 gather against +the dense gather plus the C++ FP8 quantizers, and the whole chunked prefill on a real FP4 cache +against an FP32 reference.""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = ["../../../.."] + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + mla as context, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + mla as ctx_split_k, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_context_fmha_op, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + +H, NOPE, ROPE, V, LORA = 96, 128, 64, 128, 512 +N_KVB = H * (NOPE + V) +PAGE = 128 + +requires_fp8_cvt = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9), + reason="E4M3 conversion instructions require SM89+", +) +requires_sm107 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="The K3 split-K context FMHA, the MXFP8 gather and the MXFP8 GEMM require SM107", +) + + +# --------------------------------------------------------------------------------------------- +# CPU: selection +# --------------------------------------------------------------------------------------------- + + +def _backend(heads=H, kv_heads=None, nope=NOPE, rope=ROPE, v=V): + return SimpleNamespace( + num_heads=heads, + num_kv_heads=heads if kv_heads is None else kv_heads, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope, + v_head_dim=v, + ) + + +@pytest.fixture +def info_once(monkeypatch): + calls = [] + monkeypatch.setattr(ctx_split_k.logger, "info_once", lambda msg, key: calls.append((msg, key))) + return calls + + +@pytest.mark.cpu_only +def test_split_k_support_is_shape_gated(monkeypatch, info_once): + monkeypatch.setattr(ctx_split_k, "get_sm_version", lambda: 107) + monkeypatch.setattr(k3_context_fmha_op, "IS_K3_CONTEXT_FMHA_AVAILABLE", True) + monkeypatch.setattr(k3_context_fmha_op, "K3_CONTEXT_FMHA_IMPORT_ERROR", None) + + for heads in (96, 48, 24, 12, 6, 3): + assert ctx_split_k.supports_k3_fp8_split_k_context(_backend(heads)) + assert not info_once + + cases = [ + (_backend(12, 6), "matching positive Q/KV head counts per rank, got 12/6"), + (_backend(0), "matching positive Q/KV head counts per rank, got 0/0"), + (_backend(nope=64), "requires head dimensions"), + (_backend(rope=128), "requires head dimensions"), + ] + for attn, reason in cases: + assert not ctx_split_k.supports_k3_fp8_split_k_context(attn) + assert reason in info_once[-1][0] + # One line per reason. + assert ctx_split_k.split_k_context_unavailable_reason(_backend(12, 6)) in info_once[0][1] + + monkeypatch.setattr(ctx_split_k, "get_sm_version", lambda: 103) + assert not ctx_split_k.supports_k3_fp8_split_k_context(_backend()) + assert "requires SM107, got SM103" in info_once[-1][0] + + monkeypatch.setattr(ctx_split_k, "get_sm_version", lambda: 107) + monkeypatch.setattr(k3_context_fmha_op, "IS_K3_CONTEXT_FMHA_AVAILABLE", False) + assert not ctx_split_k.supports_k3_fp8_split_k_context(_backend()) + assert "lacks split-K epilogue merge support" in info_once[-1][0] + monkeypatch.setattr( + k3_context_fmha_op, "K3_CONTEXT_FMHA_IMPORT_ERROR", ImportError("missing CuTe runtime") + ) + assert not ctx_split_k.supports_k3_fp8_split_k_context(_backend()) + assert "import failed" in info_once[-1][0] and "missing CuTe runtime" in info_once[-1][0] + assert "cannot attend a cached prefix" in info_once[-1][0] + + +def _trtllm_backend(*, fp4=True, sparse=None): + attention = object.__new__(TrtllmAttention) + attention.is_mla_enable = True + attention.has_fp4_kv_cache = fp4 + attention.has_fp8_kv_cache = not fp4 + attention.sparse_params = sparse + return attention + + +def _layer(mha, kv_b_proj=None): + """The attributes of the MLA layer the split-K context reads.""" + return SimpleNamespace( + mha=mha, + kv_b_proj=kv_b_proj, + out_scale=None, + kv_lora_rank=LORA, + qk_rope_head_dim=ROPE, + qk_nope_head_dim=NOPE, + v_head_dim=V, + num_heads_tp=H, + ) + + +@pytest.mark.cpu_only +def test_split_k_serves_bf16_fp4_cache_layers_only(monkeypatch): + supported = {"value": True} + monkeypatch.setattr( + ctx_split_k, "supports_k3_fp8_split_k_context", lambda attn: supported["value"] + ) + bf16 = torch.empty(1, dtype=torch.bfloat16) + + def uses(layer, q=bf16, output=bf16): + return ctx_split_k.uses_k3_split_k_context(layer, q, output) + + layer = _layer(_trtllm_backend()) + assert uses(layer) + assert not uses(layer, q=bf16.half()) + assert not uses(layer, output=bf16.float()) + layer.out_scale = torch.ones(1) + assert not uses(layer) + layer.out_scale = None + # Only NVFP4 latent caches (not FP8 / sparse layers) and only a TrtllmAttention backend. + assert not uses(_layer(_trtllm_backend(fp4=False))) + assert not uses(_layer(_trtllm_backend(sparse=object()))) + assert not uses(_layer(SimpleNamespace(uses_fp4_mla_attention=True))) + supported["value"] = False + assert not uses(layer) + + +def _kv_b_proj(**overrides): + fields = dict( + weight=torch.empty(2, 2, dtype=torch.float8_e4m3fn), + weight_scale=torch.empty(2, dtype=torch.uint8), + supports_prequantized_input=True, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +@pytest.mark.cpu_only +def test_kv_b_takes_mxfp8_where_the_weight_is_in_the_cute_mxfp8_layout(monkeypatch): + monkeypatch.setattr(ctx_split_k, "get_sm_version", lambda: 107) + monkeypatch.setattr(ctx_split_k, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", True) + mha = _trtllm_backend() + + def takes(**overrides): + return ctx_split_k.kv_b_takes_mxfp8(_layer(mha, _kv_b_proj(**overrides))) + + assert takes() + assert not takes(supports_prequantized_input=False) + assert not takes(bias=torch.zeros(2)) + assert not takes(weight=torch.empty(2, 2, dtype=torch.bfloat16)) + assert not takes(weight_scale=torch.empty(2, dtype=torch.int32)) + assert not takes(weight_scale=None) + layer = _layer(mha, _kv_b_proj()) + layer.kv_lora_rank = 256 + assert not ctx_split_k.kv_b_takes_mxfp8(layer) + monkeypatch.setattr(ctx_split_k, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", False) + assert not takes() + monkeypatch.setattr(ctx_split_k, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", True) + monkeypatch.setattr(ctx_split_k, "get_sm_version", lambda: 103) + assert not takes() + + +@pytest.mark.cpu_only +def test_fp4_cache_chunked_prefill_runs_split_k_or_raises(monkeypatch): + """The context dispatch hands an FP4-cache chunked prefill to the split-K context FMHA, and + raises where it does not apply (there is no other FP4 partition loop).""" + monkeypatch.setattr(context, "get_sm_version", lambda: 107) + applies = {"value": True} + monkeypatch.setattr( + ctx_split_k, "uses_k3_split_k_context", lambda layer, q, output: applies["value"] + ) + monkeypatch.setattr( + ctx_split_k, "forward_context_split_k_fp4", lambda layer, q, ckv, latent, md, out: "split_k" + ) + mha = SimpleNamespace( + uses_fp4_mla_attention=True, + has_cached_kv_for_mla_context_warmup=lambda md: False, + is_chunked_prefill_for_mla_context=lambda md: True, + has_cached_kv_for_mla_context=lambda md: True, + ) + layer = _layer(mha) + md = SimpleNamespace(helix_context_local_slots=None) + run = lambda: context.forward_context(layer, None, None, md, None, None) # noqa: E731 + assert run() == "split_k" + applies["value"] = False + with pytest.raises(NotImplementedError, match="split-K context FMHA"): + run() + mha.is_chunked_prefill_for_mla_context = lambda md: False + with pytest.raises(RuntimeError, match="enable chunked prefill"): + run() + + +# --------------------------------------------------------------------------------------------- +# Shared GPU fixtures: a real dense FP4 MLA cache and chunked-prefill metadata +# --------------------------------------------------------------------------------------------- + + +def _e4m3_byte(value): + return int(torch.tensor([value]).to(torch.float8_e4m3fn).view(torch.uint8).item()) + + +def _build_prefix_case(prefix, new, chunk_size, buffer_batch, seed=0, global_scale=1.0): + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + + nctx = len(prefix) + g = torch.Generator(device="cuda") + g.manual_seed(seed) + pages = [math.ceil((p + q) / PAGE) for p, q in zip(prefix, new)] + stride = max(pages) + total_pages = sum(pages) + 2 + kv = torch.randint( + 0, 256, (total_pages, 1, PAGE, 1, 320), dtype=torch.uint8, device="cuda", generator=g + ) + sf = torch.randint( + _e4m3_byte(0.0625), + _e4m3_byte(0.5), + (total_pages, PAGE, 40), + dtype=torch.uint8, + device="cuda", + generator=g, + ) + perm = torch.randperm(total_pages, generator=torch.Generator().manual_seed(seed)).tolist() + rows, base = [], 0 + for p in pages: + rows += perm[base : base + p] + [0] * (stride - p) + base += p + loops = math.ceil(sum(prefix) / (chunk_size * buffer_batch)) + shim = SimpleNamespace( + num_contexts=nctx, + runtime_features=SimpleNamespace( + chunk_size=chunk_size, chunked_prefill_buffer_batch_size=buffer_batch + ), + kv_cache_params=SimpleNamespace(num_cached_tokens_per_seq=list(prefix)), + ) + seq = torch.zeros((loops, nctx), dtype=torch.int) + cu = torch.zeros((loops, nctx + 1), dtype=torch.int64) + offsets = torch.zeros((loops + 1, nctx), dtype=torch.int64) + merge = torch.empty((loops + 1, nctx), dtype=torch.int64) + max_len = [] + TrtllmAttentionMetadata.pre_process_for_chunked_prefill( + shim, + chunked_seq_len=seq, + chunked_global_offset=offsets, + cu_chunked_seq_len=cu, + merge_op_tensor=merge, + max_chunk_len_per_loop=max_len, + chunked_loop_num=loops, + ) + T = sum(new) + cu_q = torch.tensor( + [0] + torch.tensor(new).cumsum(0).tolist(), dtype=torch.int32, device="cuda" + ) + meta = SimpleNamespace( + num_contexts=nctx, + num_ctx_tokens=T, + chunked_loop_num=loops, + chunked_seq_len=seq.cuda(), + host_chunked_seq_len=seq, + cu_chunked_seq_len=cu.cuda(), + host_cu_chunked_seq_len=cu, + chunked_global_offset=offsets.cuda(), + max_chunk_len_per_loop=max_len, + merge_op_tensor=merge.cuda(), + max_ctx_seq_len=max(new), + mla_prepare_ctx_cu_seqlens=lambda: cu_q, + fp4_mla_state=SimpleNamespace( + device_page_table=True, + device_page_table_valid=True, + page_table_stride=stride, + _paged_kv_indices=torch.tensor(rows, dtype=torch.int32, device="cuda"), + kv_global_scale=torch.full((1,), global_scale, dtype=torch.float32, device="cuda"), + ), + page_size=PAGE, + kv_cache_manager=SimpleNamespace(get_fp4_mla_cache_buffers=lambda _layer: (kv, sf)), + ) + q = torch.randn(T, H * (NOPE + ROPE), generator=g, device="cuda").to(torch.bfloat16) + latent = (torch.randn(T, LORA + ROPE, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + return SimpleNamespace( + prefix=prefix, + new=new, + meta=meta, + q=q, + latent=latent, + rows=rows, + stride=stride, + T=T, + loops=loops, + cu_q=cu_q, + ) + + +def _dense_gather(meta, *, tokens, cu, offsets, max_len): + from tensorrt_llm._torch.attention.backends.fp4_mla import load_fp4_mla_chunked_kv_cache + + return load_fp4_mla_chunked_kv_cache( + meta, + 0, + num_ctx_cached_tokens=tokens, + cu_chunked_seq_len=cu, + chunked_global_offset=offsets, + chunked_max_seq_len=max_len, + out_dtype=torch.bfloat16, + kv_lora_rank=LORA, + qk_rope_head_dim=ROPE, + ) + + +class _KvBProj: + """kv_b_proj of the reference: a BF16 GEMM and, with ``mxfp8`` (SM100 family), its 1x128 + MXFP8 quantization in the layout ``cute_dsl_mxfp8_gemm_rubin`` reads.""" + + bias = None + + def __init__(self, mxfp8=False, seed=1): + g = torch.Generator(device="cuda") + g.manual_seed(seed) + self.weight_bf16 = (torch.randn(N_KVB, LORA, generator=g, device="cuda") / 16).to( + torch.bfloat16 + ) + self.weight = self.weight_scale = None + if mxfp8: + self.weight, self.weight_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + self.weight_bf16 + ) + + def __call__(self, x): + return x @ self.weight_bf16.T + + +class _Mha: + """The attention backend calls of the split-K path, over the case's FP4 cache.""" + + layer_idx = 0 + num_heads = num_kv_heads = H + qk_nope_head_dim, qk_rope_head_dim, v_head_dim = NOPE, ROPE, V + q_scaling = 1.0 + + def __init__(self, loader=None): + self.kv_cache_scaling_factor = torch.ones(1, dtype=torch.float32, device="cuda") + self._loader = loader + + def append_fp4_mla_context_cache(self, q, latent_cache, metadata): + pass + + def load_chunked_kv_cache_for_mla( + self, + metadata, + num_ctx_cached_tokens, + cu_chunked_seq_len, + chunked_global_offset, + chunked_max_seq_len, + out_dtype, + ): + assert out_dtype == torch.bfloat16 + if self._loader is not None: + return self._loader(num_ctx_cached_tokens) + return _dense_gather( + metadata, + tokens=num_ctx_cached_tokens, + cu=cu_chunked_seq_len, + offsets=chunked_global_offset, + max_len=chunked_max_seq_len, + ) + + +def _split_k_forward(monkeypatch, mha, kv_b_proj, mxfp8, case): + """``ctx_split_k.forward_context_split_k_fp4`` of one case, with kv_b on the MXFP8 gather or + not.""" + monkeypatch.setattr(ctx_split_k, "kv_b_takes_mxfp8", lambda layer: mxfp8) + out = torch.empty(case.T, H * V, dtype=torch.bfloat16, device="cuda") + ctx_split_k.forward_context_split_k_fp4( + _layer(mha, kv_b_proj), + case.q, + case.latent[:, :LORA].contiguous(), + case.latent, + case.meta, + out, + ) + torch.cuda.synchronize() + return out + + +def _unit_quant(x): + unit = torch.ones(1, dtype=torch.float32, device=x.device) + return torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x.contiguous(), unit)[0] + + +_CASES = [ + ([5000], [300], 1024, 2), # slices 2048, 2048, 904 of one request + ([1500, 3000, 700], [128, 64, 200], 1024, 2), # slices span requests + ([66000], [256], 16384, 4), # production slice size: 65536 + 464 +] + + +# --------------------------------------------------------------------------------------------- +# GPU: the partition schedule +# --------------------------------------------------------------------------------------------- + + +@requires_fp8_cvt +def test_split_k_partitions_follow_the_chunk_schedule(monkeypatch): + """Every cached-prefix slice and then the current chunk reach the FMHA once, with their KV + lengths, merge ops and mask, unconcatenated FP8 K-nope / K-pe / V, and one FP8 Q.""" + import tensorrt_llm # noqa: F401 + + case = _build_prefix_case([1500, 3000, 700], [128, 64, 200], 1024, 2) + loaded = [] + + def loader(tokens): + latent = (torch.randn(tokens, LORA + ROPE, device="cuda") * 0.5).to(torch.bfloat16) + loaded.append(latent) + return latent[:, :LORA].contiguous(), latent[:, LORA:].contiguous() + + calls = [] + + def record(attn, q, k_nope, k_pe, v, output, stats, merge_op, cu_q, cu_kv, **kwargs): + calls.append( + dict( + q=q, + k_nope=k_nope, + k_pe=k_pe, + v=v, + output=output, + stats=stats, + merge_op=merge_op, + cu_q=cu_q, + cu_kv=cu_kv, + **kwargs, + ) + ) + + monkeypatch.setattr(ctx_split_k, "run_k3_split_k_context_partition", record) + kv_b_proj = _KvBProj() + out = _split_k_forward(monkeypatch, _Mha(loader), kv_b_proj, False, case) + + meta = case.meta + assert len(calls) == case.loops + 1 == len(loaded) + 1 + q_fp8 = calls[0]["q"] + assert q_fp8.dtype == torch.float8_e4m3fn + assert torch.equal(q_fp8.view(torch.uint8), _unit_quant(case.q).view(torch.uint8)) + stats = calls[0]["stats"] + assert stats.shape == (case.T, H, 2) and stats.dtype == torch.float32 + for loop_idx, call in enumerate(calls): + last = loop_idx == case.loops + latent = case.latent if last else loaded[loop_idx] + assert call["q"] is q_fp8 and call["stats"] is stats and call["cu_q"] is case.cu_q + assert call["output"].data_ptr() == out.data_ptr() + assert call["is_causal"] == last + assert call["max_q_len"] == meta.max_ctx_seq_len + assert call["merge_op"].data_ptr() == meta.merge_op_tensor[loop_idx].data_ptr() + if last: + assert call["cu_kv"] is case.cu_q and call["max_kv_len"] == meta.max_ctx_seq_len + else: + assert call["cu_kv"].data_ptr() == meta.cu_chunked_seq_len[loop_idx].data_ptr() + assert call["max_kv_len"] == meta.max_chunk_len_per_loop[loop_idx] + kv = _unit_quant(kv_b_proj(latent[:, :LORA].contiguous())) + assert torch.equal(call["k_nope"].view(torch.uint8), kv[:, : H * NOPE].view(torch.uint8)) + assert torch.equal(call["v"].view(torch.uint8), kv[:, H * NOPE :].view(torch.uint8)) + assert call["k_nope"].stride(-1) == 1 and call["v"].stride(-1) == 1 + assert call["k_pe"].shape == (latent.shape[0], ROPE) + assert torch.equal( + call["k_pe"].view(torch.uint8), _unit_quant(latent[:, LORA:]).view(torch.uint8) + ) + + +# --------------------------------------------------------------------------------------------- +# SM107: the MXFP8 gather and the whole chunked prefill +# --------------------------------------------------------------------------------------------- + + +@requires_sm107 +@pytest.mark.parametrize("global_scale", [1.0, 32.0]) +@pytest.mark.parametrize("prefix,new,chunk_size,buffer_batch", _CASES[:2]) +def test_mxfp8_gather_matches_the_quantized_dense_gather( + prefix, new, chunk_size, buffer_batch, global_scale +): + """Per partition, the MXFP8 latent (every scale byte, padding included) and the FP8 K-pe are + the C++ MXFP8 / unit FP8 quantization of the dense BF16 gather (power-of-two global scales + keep the division exact in both).""" + import tensorrt_llm # noqa: F401 + + case = _build_prefix_case(prefix, new, chunk_size, buffer_batch, global_scale=global_scale) + meta = case.meta + for loop_idx in range(case.loops): + tokens = int(meta.host_cu_chunked_seq_len[loop_idx, meta.num_contexts]) + args = dict( + cu=meta.cu_chunked_seq_len[loop_idx], + offsets=meta.chunked_global_offset[loop_idx], + max_len=int(meta.max_chunk_len_per_loop[loop_idx]), + ) + (latent, latent_sf), k_pe = ctx_split_k.load_fp4_mla_chunked_kv_cache_mxfp8( + meta, + 0, + num_ctx_cached_tokens=tokens, + cu_chunked_seq_len=args["cu"], + chunked_global_offset=args["offsets"], + chunked_max_seq_len=args["max_len"], + kv_lora_rank=LORA, + qk_rope_head_dim=ROPE, + ) + dense, dense_pe = _dense_gather(meta, tokens=tokens, **args) + expected, expected_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(dense) + assert torch.equal(latent.view(torch.uint8), expected.view(torch.uint8)) + assert torch.equal(latent_sf.view(torch.uint8), expected_sf.view(torch.uint8)) + assert torch.equal(k_pe.view(torch.uint8), _unit_quant(dense_pe).view(torch.uint8)) + + +@requires_sm107 +def test_mxfp8_gather_validates_its_inputs(): + import tensorrt_llm # noqa: F401 + + case = _build_prefix_case([300], [10], 1024, 1) + meta = case.meta + kwargs = dict( + num_ctx_cached_tokens=300, + cu_chunked_seq_len=meta.cu_chunked_seq_len[0], + chunked_global_offset=meta.chunked_global_offset[0], + chunked_max_seq_len=300, + kv_lora_rank=LORA, + qk_rope_head_dim=ROPE, + ) + with pytest.raises(ValueError, match="kv_lora_rank=512"): + ctx_split_k.load_fp4_mla_chunked_kv_cache_mxfp8(meta, 0, **{**kwargs, "kv_lora_rank": 256}) + with pytest.raises(ValueError, match="int64 cumulative-length"): + ctx_split_k.load_fp4_mla_chunked_kv_cache_mxfp8( + meta, 0, **{**kwargs, "cu_chunked_seq_len": meta.cu_chunked_seq_len[0].int()} + ) + (latent, latent_sf), k_pe = ctx_split_k.load_fp4_mla_chunked_kv_cache_mxfp8( + meta, 0, **{**kwargs, "num_ctx_cached_tokens": 0} + ) + assert latent.shape == (0, LORA) and latent_sf.numel() == 0 and k_pe.shape == (0, ROPE) + + +def _reference(case, kv_b_proj, mxfp8): + """FP32 attention over the FP8 Q / K / V the FMHA consumes, the whole prefix in one piece.""" + q8 = _unit_quant(case.q).float().view(case.T, H, NOPE + ROPE) + + def expand(latent): + if mxfp8: + a, a_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(latent) + return _unit_quant( + torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + a, kv_b_proj.weight, a_sf, kv_b_proj.weight_scale + ) + ).float() + return _unit_quant(kv_b_proj(latent)).float() + + cur_kv = expand(case.latent[:, :LORA].contiguous()) + cur_pe = _unit_quant(case.latent[:, LORA:]).float() + out = torch.empty(case.T, H, V, device="cuda") + q0 = 0 + for r, (p, n) in enumerate(zip(case.prefix, case.new)): + meta = SimpleNamespace(**vars(case.meta)) + meta.num_contexts = 1 + meta.fp4_mla_state = SimpleNamespace(**vars(case.meta.fp4_mla_state)) + meta.fp4_mla_state._paged_kv_indices = torch.tensor( + case.rows[r * case.stride : (r + 1) * case.stride], dtype=torch.int32, device="cuda" + ) + dense, dense_pe = _dense_gather( + meta, + tokens=p, + cu=torch.tensor([0, p], dtype=torch.int64, device="cuda"), + offsets=torch.zeros(1, dtype=torch.int64, device="cuda"), + max_len=p, + ) + kv = torch.cat((expand(dense), cur_kv[q0 : q0 + n])) + pe = torch.cat((_unit_quant(dense_pe).float(), cur_pe[q0 : q0 + n])) + k_nope = kv[:, : H * NOPE].view(-1, H, NOPE) + v = kv[:, H * NOPE :].view(-1, H, V) + causal = torch.arange(n, device="cuda")[None, :] > torch.arange(n, device="cuda")[:, None] + for h0 in range(0, H, 16): + heads = slice(h0, h0 + 16) + scores = torch.einsum( + "qhd,khd->hqk", q8[q0 : q0 + n, heads, :NOPE], k_nope[:, heads] + ) + torch.einsum("qhd,kd->hqk", q8[q0 : q0 + n, heads, NOPE:], pe) + scores /= math.sqrt(NOPE + ROPE) + scores[:, :, p:].masked_fill_(causal[None], -torch.inf) + out[q0 : q0 + n, heads] = torch.einsum("hqk,khd->qhd", scores.softmax(-1), v[:, heads]) + q0 += n + return out.view(case.T, H * V) + + +@requires_sm107 +@pytest.mark.parametrize("mxfp8", [True, False], ids=["mxfp8_gather", "bf16_gather"]) +@pytest.mark.parametrize("prefix,new,chunk_size,buffer_batch", _CASES) +def test_split_k_chunked_prefill_matches_reference( + prefix, new, chunk_size, buffer_batch, mxfp8, monkeypatch +): + """Output against an FP32 reference over the C++ FP8 quantization of Q / K / V (the Triton + quantizers' bytes are those, see ``test_k3_mla_ctx_fp8.py``).""" + import tensorrt_llm # noqa: F401 + + reason = ctx_split_k.split_k_context_unavailable_reason(_Mha()) + if reason is not None: + pytest.skip(f"K3 split-K context FMHA unavailable: {reason}") + kv_b_proj = _KvBProj(mxfp8) + case = _build_prefix_case(prefix, new, chunk_size, buffer_batch) + out = _split_k_forward(monkeypatch, _Mha(), kv_b_proj, mxfp8, case) + torch.testing.assert_close( + out.float(), _reference(case, kv_b_proj, mxfp8), atol=5e-2, rtol=5e-2 + ) + + +# --------------------------------------------------------------------------------------------- +# SM107: the fp4_mla_gather_mxfp8 entry on its cells +# --------------------------------------------------------------------------------------------- + + +def _gather_entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.fp4_mla_gather_mxfp8 import ( # noqa: E501 + fp4_mla_gather_mxfp8, + ) + + return fp4_mla_gather_mxfp8 + + +@requires_sm107 +@pytest.mark.parametrize( + "cell", _gather_entry().CELLS, ids=[c.why[:48] for c in _gather_entry().CELLS] +) +def test_fp4_mla_gather_mxfp8_cells(cell): + """Every partition of the cell's chunk schedule under ``validating``: K-pe and the MXFP8 + latent within one E4M3 step of the reference (the kernel's approximate division by the + global scale and reciprocal of 448 can move a BF16 value by an ulp, or a block's exponent by + one at a power-of-two boundary).""" + from _validating import validating + + import tensorrt_llm # noqa: F401 + from tensorrt_llm._torch.attention.backends.fp4_mla import FP4_MLA_KV_GLOBAL_SCALE + + op = _gather_entry() + spec = cell.spec + case = _build_prefix_case( + spec["prefix"], + spec["new"], + spec["chunk_size"], + spec["buffer_batch"], + global_scale=FP4_MLA_KV_GLOBAL_SCALE, + ) + meta = case.meta + kv, sf = meta.kv_cache_manager.get_fp4_mla_cache_buffers(0) + state = meta.fp4_mla_state + for loop_idx in range(case.loops): + tokens = int(meta.host_cu_chunked_seq_len[loop_idx, meta.num_contexts]) + a = torch.empty((tokens, LORA), dtype=torch.float8_e4m3fn, device="cuda") + asc = torch.full(((tokens + 127) // 128 * 128 * 16,), 0xFF, dtype=torch.uint8, + device="cuda") # fmt: skip + pe = torch.empty((tokens, ROPE), dtype=torch.float8_e4m3fn, device="cuda") + args = (a, asc, pe, kv, sf.view(torch.float8_e4m3fn), state._paged_kv_indices, + meta.cu_chunked_seq_len[loop_idx], meta.chunked_global_offset[loop_idx], + state.kv_global_scale, int(meta.max_chunk_len_per_loop[loop_idx]), PAGE, + state.page_table_stride, meta.num_contexts) # fmt: skip + ref_a, ref_asc, ref_pe = op.reference(*args) + with validating(op): + op(*args) + torch.cuda.synchronize() + # The kernel divides by the global scale with div.full (2 ulp); a BF16 rounding can then + # move a K-pe value by one E4M3 step (2^-3 relative), nothing more. + pe_err = (pe.float() - ref_pe.float()).abs() + assert bool((pe_err <= ref_pe.float().abs() * 2.0**-3 + 2.0**-9).all()), cell.why + assert int((asc.int() - ref_asc.int()).abs().max()) <= 1, cell.why + + def dequant(codes, scales): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.k3_mla_q_fused import ( # noqa: E501 + swizzled_sf_offset, + ) + + rows = torch.arange(tokens, device="cuda")[:, None] + offsets = swizzled_sf_offset(rows, torch.arange(16, device="cuda")[None, :], 16) + exponent = scales[offsets].float().repeat_interleave(32, -1) + return codes.float() * torch.exp2(exponent - 127) + + got, ref = dequant(a, asc), dequant(ref_a, ref_asc) + step = ref.abs().view(tokens, 4, 128).amax(-1, keepdim=True) * 2.0**-3 + assert bool(((got - ref).abs().view(tokens, 4, 128) <= step).all()), cell.why diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_kv_fp8.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_kv_fp8.py new file mode 100644 index 000000000000..993212d819cf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_kv_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The target's ``mla_kv_to_fp8`` (``kernels/mla/mla_kv_fp8.py``) writes the attention op's FP8 context K / V +bytes.""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers the C++ ops) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.mla.mla_kv_fp8 import ( # noqa: E501 + mla_kv_to_fp8, +) + +NOPE, ROPE, V, LORA = 128, 64, 128, 512 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9), + reason="needs FP8 conversion instructions (SM89+)", +) + + +def _input(rows: int, cols: int, seed: int) -> torch.Tensor: + """BF16 values spanning the E4M3 range, with saturating values and signed zeros.""" + generator = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(rows, cols, device="cuda", generator=generator) * 64 + x.view(-1)[::97] = 1000.0 + x.view(-1)[1::97] = -1000.0 + x.view(-1)[2::97] = -0.0 + return x.to(torch.bfloat16) + + +def _cpp_unit_scale_fp8(x: torch.Tensor) -> torch.Tensor: + unit = torch.ones(1, dtype=torch.float32, device="cuda") + return torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x.contiguous(), unit)[0] + + +@pytest.mark.parametrize("tokens,heads", [(1, 96), (31, 96), (1000, 96), (4099, 12), (65553, 8)]) +@pytest.mark.parametrize("strided_pe", [False, True], ids=["contiguous_pe", "latent_cache_pe"]) +def test_mla_kv_to_fp8_matches_cat_and_cpp_quant(tokens: int, heads: int, strided_pe: bool): + """K / V are the unit-scale C++ FP8 bytes of the BF16 [T, H, nope + rope] concat of k_nope + and k_pe and of the V columns.""" + kv = _input(tokens, heads * (NOPE + V), seed=tokens) + latent = _input(tokens, LORA + ROPE, seed=tokens + 1) + k_pe = latent[:, LORA:] if strided_pe else latent[:, LORA:].contiguous() + k8, v8 = mla_kv_to_fp8(kv, k_pe, heads, NOPE, V) + k_bf16 = torch.cat( + (kv[:, : heads * NOPE].view(-1, heads, NOPE), k_pe[:, None].expand(-1, heads, -1)), -1 + ).view(tokens, -1) + for x8, x in ((k8, k_bf16), (v8, kv[:, heads * NOPE :])): + assert x8.dtype == torch.float8_e4m3fn and x8.is_contiguous() and x8.shape == x.shape + assert torch.equal(x8.view(torch.uint8), _cpp_unit_scale_fp8(x).view(torch.uint8)) + + +def test_mla_kv_to_fp8_rejects_bad_layouts(): + heads, tokens = 4, 16 + kv = torch.randn(tokens, heads * (NOPE + V), device="cuda").to(torch.bfloat16) + latent = torch.randn(tokens, LORA + ROPE, device="cuda").to(torch.bfloat16) + k, v = mla_kv_to_fp8(kv[:0], latent[:0, LORA:], heads, NOPE, V) + assert k.shape == (0, heads * (NOPE + ROPE)) and v.shape == (0, heads * V) + bad_calls = ( + (kv[:, 1:], latent[:, LORA:]), # kv width + (kv.float(), latent[:, LORA:]), # kv dtype + (kv, latent[:, LORA - 1 : -1]), # odd k_pe offset + (kv, latent[:, LORA : LORA + 48]), # k_pe width not a power of two + (kv[:-1], latent[:, LORA:]), # token counts differ + ) + for kv_arg, pe_arg in bad_calls: + with pytest.raises(ValueError): + mla_kv_to_fp8(kv_arg, pe_arg, heads, NOPE, V) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_q_fused.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_q_fused.py new file mode 100644 index 000000000000..80080dd5fddf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_mla_q_fused.py @@ -0,0 +1,709 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi-K3 fused FP4 MLA query (``k3_mla_q_fused``) against the default query chain. + +The FP4 cache, HP ring and Helix per-token metadata are built through the real +``Fp4MlaKVCacheManagerV2`` (context scatter, then one generation step of ``q_len`` tokens per +request), with the per-batch FP4 MLA state on ``metadata.fp4_mla_state``; the decode is K3's own +(``ops/attention/_fp4_mla_decode``). The default chain is the production one: the q_b low-M GEMM, +the BF16 W_UK BMM and the fused FP4 generation update with its query rows. Checks: + +* The KV append without query rows (``k3_rubin/kernels/fp4_mla_kv_append``) writes the cache + bytes of the backend's full update (bitwise). +* The fused kernel's query quantization equals the Triton kernel's on the same BF16 values + (bitwise), including the RoPE rotation with a non-identity table. +* Its q_b and W_UK products agree with the production GEMMs up to BF16 rounding of a different + FP32 summation order. +* With the q_a RMSNorm inside the kernel the query rows are those after FlashInfer's separate + norm (bitwise), and ``k3_mla_latent_row`` is bitwise FlashInfer's norm plus the ``k_pe`` copy. +* ``generation.forward_gen_q_fused`` (multi-stream, KV half on the aux stream) gives + the default chain's cache bytes and, through K3's FP4 decode, its attention output within the + FP4 query quantization. +* The ``k3_mla_q_fused`` entry on its ``.CELLS`` is within one FP4 step of its plain-torch + reference. + +``test_k3_ull2_mla.py`` / ``test_k3_ull4_mla.py`` / ``test_k3_ull3_mla.py`` (same directory) build +on the FP4 cache helpers here. +""" + +import functools +import math +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = ["../../../.."] + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="needs an SM107 GPU", +) + +H, QL, LORA, ROPE, PAGE = 96, 1536, 512, 64, 128 +Q_LEN = 8 +EPS = 1e-6 + + +def k3_decode_unavailable() -> bool: + """Whether the packed CuTe DSL FP4 MLA decode cannot run here (SM107 with a CuTe DSL that + builds the FP4 MLA kernels).""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + return True + from tensorrt_llm._torch.attention.backends.fp4_mla.config import _cutedsl_backend_available + + return not _cutedsl_backend_available() + + +# Tests that run K3's FP4 MLA decode (``ops/attention/_fp4_mla_decode``). +needs_k3_decode = pytest.mark.skipif( + k3_decode_unavailable(), + reason="needs an SM107 GPU and a CuTe DSL that builds the FP4 MLA decode kernels", +) + + +def k3_decode(): + """K3's FP4 MLA decode module.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention import ( # noqa: E501 + _fp4_mla_decode as fp4_mla_decode, + ) + + return fp4_mla_decode + + +def _pow2_scales(rows, cols, gen, lo, hi): + grid = (math.ceil(rows / 128), math.ceil(cols / 128)) + return torch.exp2(torch.randint(lo, hi, grid, generator=gen).float()) + + +@functools.lru_cache(maxsize=1) +def _weights(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_fp8w_lowm_gemm as lowm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_mla_q_fused as fused, + ) + + dev = torch.device("cuda") + gen = torch.Generator().manual_seed(7) + w_qb = (torch.randn(H * 192, QL, generator=gen) * 2).to(torch.float8_e4m3fn) + s_qb = _pow2_scales(H * 192, QL, gen, -8, -6) + w_qb, packed = lowm.prepare_k3_fp8w_lowm_weight(w_qb.to(dev), s_qb.to(dev)) + return SimpleNamespace( + w_qb=w_qb, + s_qb_f32=s_qb.to(dev), + s_qb_packed=packed, + qb_exp=fused.block_exponents(s_qb.to(dev), H * 192, QL), + w_uk=(torch.randn(H, LORA, 128, generator=gen) * 0.06).to(torch.bfloat16).to(dev), + q_norm=(1 + 0.1 * torch.randn(QL, generator=gen)).bfloat16().to(dev), + kv_norm=(1 + 0.1 * torch.randn(LORA, generator=gen)).bfloat16().to(dev), + ) + + +def _rotary_table(max_pos, kind): + dev = torch.device("cuda") + rot = torch.zeros((max_pos, ROPE, 2), dtype=torch.float32, device=dev) + if kind == "identity": + rot[..., 0] = 1.0 + else: + gen = torch.Generator(device=dev).manual_seed(5) + angle = torch.rand((max_pos, ROPE), device=dev, generator=gen) * 6.28 + rot[..., 0] = torch.cos(angle) + rot[..., 1] = torch.sin(angle) + return rot.view(max_pos, ROPE * 2).contiguous() + + +def _fp4_case(n, kv, seed): + """Real FP4 cache of ``n`` requests with ``kv`` tokens each (or the ``n`` lengths in ``kv``; + the last ``Q_LEN`` of them this step's), Helix speculative per-token metadata; returns + (manager, metadata).""" + import tensorrt_llm + from tensorrt_llm._torch.attention.backends.fp4_mla import ( + FP4_MLA_KV_GLOBAL_SCALE, + FP4_MLA_Q_GLOBAL_SCALE, + scatter_fp4_mla_kv_cache, + ) + from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 + from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState + from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + from tensorrt_llm.mapping import Mapping + + torch.manual_seed(seed) + dev = torch.device("cuda") + seq_lens = list(kv) if isinstance(kv, (list, tuple)) else [kv] * n + assert len(seq_lens) == n + num_blocks = [(s + PAGE - 1) // PAGE for s in seq_lens] + max_seq_len = max(seq_lens) + 32 * PAGE + max_tokens = sum(num_blocks) * PAGE + max_tokens += max(max_tokens // 2, n * 64 * PAGE) + ctx_lens = [s - Q_LEN for s in seq_lens] + mgr = Fp4MlaKVCacheManagerV2( + LlmKvCacheConfig( + max_tokens=max_tokens, dtype="nvfp4", enable_block_reuse=False, host_cache_size=0 + ), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=LORA + ROPE, + tokens_per_block=PAGE, + max_seq_len=max_seq_len, + max_batch_size=n, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=tensorrt_llm.bindings.DataType.NVFP4, + spec_config=MTPDecodingConfig(max_draft_len=Q_LEN - 1), + max_num_tokens=max_tokens, + pretrained_config=SimpleNamespace(kv_lora_rank=LORA), + ) + mgr.add_dummy_requests(list(range(n)), seq_lens) + kv_cache, sf_cache = mgr.get_fp4_mla_cache_buffers(0) + kv_cache.zero_() + sf_cache.zero_() + request_ids = list(range(n)) + block_ids = mgr.get_batch_cache_indices(request_ids, layer_idx=0) + spec = mgr.get_fp4_mla_page_table_spec(0) + hp_block_ids = mgr._get_batch_cache_indices_by_pool_id( + request_ids, pool_id=spec.hp_pool_id, is_kv_aggregate=False + ) + stride = max(num_blocks) + rows, hp_rows = [], [] + for i in range(n): + b = block_ids[i][: num_blocks[i]] + rows.extend(b + [0] * (stride - len(b))) + hb = hp_block_ids[i][: num_blocks[i]] + hp_rows.extend(hb + [0] * (stride - len(hb))) + idx = torch.tensor(rows, dtype=torch.int32, device=dev) + hp_idx = torch.tensor(hp_rows, dtype=torch.int32, device=dev) + indptr = torch.arange(n + 1, dtype=torch.int32, device=dev) * stride + state = Fp4MlaState( + hp_pool=mgr.get_fp4_mla_hp_pool(), + hp_page_indices=hp_idx, + v_scale_pool=mgr.get_mla_v_scale_pool(), + q_global_scale=torch.tensor([FP4_MLA_Q_GLOBAL_SCALE], dtype=torch.float32, device=dev), + kv_global_scale=torch.tensor([FP4_MLA_KV_GLOBAL_SCALE], dtype=torch.float32, device=dev), + generation_kv_lens=torch.empty(n, dtype=torch.int32, device=dev), + generation_append_lens=torch.empty(n, dtype=torch.int32, device=dev), + generation_lengths_num_tokens=-1, + generation_lengths_num_seqs=-1, + generation_lengths_num_contexts=-1, + generation_lengths_capture_recorded=False, + device_page_table=True, + device_page_table_valid=True, + page_table_stride=stride, + context_repack_max_touched_pages=stride, + _paged_kv_indices=idx, + _paged_kv_indptr=indptr, + paged_kv_indptr_decode=indptr.clone(), + num_sequences=n, + num_blocks=None, + num_context_blocks=n * stride, + num_generation_blocks=0, + ) + md = SimpleNamespace( + kv_cache_manager=mgr, + page_size=PAGE, + num_contexts=n, + num_seqs=n, + beam_width=1, + request_ids=request_ids, + runtime_features=SimpleNamespace(has_speculative_draft_tokens=False), + is_cuda_graph=False, + fp4_mla_state=state, + ) + v_packed = mgr.get_mla_v_packed_pool_base() + if v_packed is not None: + v_packed.zero_() + state.v_scale_pool.zero_() + md.kv_lens_cuda_runtime = torch.tensor(ctx_lens, dtype=torch.int32, device=dev) + md.prompt_lens_cuda_runtime = md.kv_lens_cuda_runtime.clone() + md.prompt_lens_cpu_runtime = torch.tensor(ctx_lens, dtype=torch.int32) + state.batch_indices = torch.cat( + [torch.full((s,), i, dtype=torch.int32, device=dev) for i, s in enumerate(ctx_lens)] + ) + state.positions = torch.cat([torch.arange(s, dtype=torch.int32, device=dev) for s in ctx_lens]) + latent = (torch.randn(sum(ctx_lens), LORA + ROPE, device=dev) * 0.5).clamp_(-3, 3).bfloat16() + scatter_fp4_mla_kv_cache( + md, latent, layer_idx=0, token_offset=0, phase="context", local_layer=0, v_head_dim=LORA + ) + md.num_contexts = 0 + state.num_context_blocks = 0 + state.num_generation_blocks = n * stride + md.kv_lens_cuda_runtime = torch.tensor(seq_lens, dtype=torch.int32, device=dev) + md.prompt_lens_cuda_runtime = torch.full((n,), Q_LEN, dtype=torch.int32, device=dev) + md.prompt_lens_cpu_runtime = torch.full((n,), Q_LEN, dtype=torch.int32) + state.batch_indices = torch.arange(n, dtype=torch.int32, device=dev).repeat_interleave(Q_LEN) + state.positions = torch.cat( + [torch.arange(c, s, dtype=torch.int32, device=dev) for c, s in zip(ctx_lens, seq_lens)] + ) + md.num_tokens = n * Q_LEN + md.num_ctx_tokens = 0 + md.num_generations = n + md.max_num_tokens = n * Q_LEN + md.max_num_sequences = n + md.max_total_draft_tokens = Q_LEN - 1 + md.max_seq_len = max_seq_len + md.runtime_features.has_speculative_draft_tokens = True + md.helix_position_offsets = state.positions.clone() + md.helix_local_slots = state.positions.clone() + md.helix_kv_bounds = state.positions + 1 + md.helix_is_inactive_rank = torch.zeros(n, dtype=torch.bool, device=dev) + md._helix_spec_tokens_valid = True + return mgr, md + + +def _pools(mgr): + kv, sf = mgr.get_fp4_mla_cache_buffers(0) + pools = [kv, sf, mgr.get_fp4_mla_hp_pool(), mgr.get_mla_v_scale_pool()] + v_packed = mgr.get_mla_v_packed_pool_base() + return pools + ([v_packed] if v_packed is not None else []) + + +def _rows_per_seq(md): + """Query rows one request owns in the decode's Q buffers (packed or one row per head).""" + return int(md.fp4_mla_state.prequantized_q_rows_per_seq or 0) or Q_LEN * H + + +def _q_bytes(md, num_tokens): + from tensorrt_llm._torch.attention.backends.fp4_mla import _get_fp4_mla_swizzled_scale_size + + state = md.fp4_mla_state + rows = _rows_per_seq(md) * (num_tokens // Q_LEN) + sf_bytes = _get_fp4_mla_swizzled_scale_size(rows, 640) + return ( + state.prequantized_q[:rows].clone(), + state.prequantized_q_sf.view(torch.uint8)[:sf_bytes].clone(), + ) + + +def _scatter(md, latent, rot, q_pe, q_rope_out, q_input): + from tensorrt_llm._torch.attention.backends.fp4_mla import scatter_fp4_mla_kv_cache + + assert scatter_fp4_mla_kv_cache( + md, + latent, + 0, + token_offset=0, + phase="generation", + local_layer=0, + v_head_dim=LORA, + rotary_cos_sin=rot, + q_pe=q_pe, + q_rope_out=q_rope_out, + q_quant_input=q_input, + helix_position_offsets=md.helix_position_offsets, + helix_is_inactive_rank=md.helix_is_inactive_rank, + ) + + +def _attention(rot=None): + """The attributes of the FP4 MLA layer's ``TrtllmAttention`` that the KV-only generation + update and K3's generation decode use.""" + return SimpleNamespace( + num_heads=H, + kv_lora_rank=LORA, + qk_rope_head_dim=ROPE, + qk_nope_head_dim=128, + q_scaling=1.0, + uses_fp4_mla_attention=True, + rotary_cos_sin=rot, + layer_idx=0, + _ensure_rope_table_size=lambda _: None, + get_fp4_mla_local_layer_idx=lambda _: 0, + ) + + +def _external_scatter(md, latent, rot): + """The KV-only generation update: cache rows and the packed-Q buffers, no query rows.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + fp4_mla_kv_append, + ) + + fp4_mla_kv_append.fp4_mla_kv_append_generation(_attention(rot), latent, md, H) + + +def _run_fused(md, w, x, rot, debug=None, norm=None): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_mla_q_fused as fused, + ) + + state = md.fp4_mla_state + fused.k3_mla_q_fused( + x, + w.w_qb, + w.qb_exp, + w.w_uk, + state.q_global_scale, + rot, + md.helix_position_offsets, + state.prequantized_q, + state.prequantized_q_sf.view(torch.uint8), + gen_len=Q_LEN, + rows_per_seq=_rows_per_seq(md), + debug=debug, + norm=norm, + ) + + +def _bf16_ulps(a, b): + """``|a - b|`` in units of the BF16 spacing at ``b``.""" + a, b = a.float(), b.float() + ulp = torch.exp2(torch.floor(torch.log2(b.abs().clamp_min(2.0**-126))) - 7) + return (a - b).abs() / ulp + + +def _inputs(num_tokens, seed): + gen = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(num_tokens, QL, device="cuda", generator=gen).bfloat16() + latent = (torch.randn(num_tokens, LORA + ROPE, device="cuda", generator=gen) * 0.5).bfloat16() + return x, latent + + +@pytest.mark.parametrize("n", [1, 2, 4]) +def test_kv_append_without_query_rows_is_bitwise(n): + mgr, md = _fp4_case(n, 1000, seed=n) + try: + num_tokens = n * Q_LEN + rot = _rotary_table(md.max_seq_len, "identity") + _, latent = _inputs(num_tokens, 3) + before = [p.clone() for p in _pools(mgr)] + q_input = torch.randn(num_tokens, H, LORA + ROPE, device="cuda").bfloat16() + q_pe = torch.randn(num_tokens, H, ROPE, device="cuda").bfloat16() + _scatter(md, latent, rot, q_pe, q_input[..., LORA:], q_input) + full = [p.clone() for p in _pools(mgr)] + for p, p0 in zip(_pools(mgr), before): + p.copy_(p0) + _external_scatter(md, latent, rot) + torch.cuda.synchronize() + for p, ref in zip(_pools(mgr), full): + assert torch.equal(p, ref) + assert md.fp4_mla_state.prequantized_q is not None + assert md.fp4_mla_state.prequantized_q_sf is not None + finally: + mgr.shutdown() + + +@pytest.mark.parametrize("rope", ["identity", "random"]) +@pytest.mark.parametrize("n", [1, 2]) +def test_query_quantization_matches_cache_update(n, rope): + """The fused kernel's FP4 query rows are the generation update's on the kernel's own BF16 + absorbed / q_b values.""" + w = _weights() + mgr, md = _fp4_case(n, 1000, seed=10 + n) + try: + num_tokens = n * Q_LEN + rot = _rotary_table(md.max_seq_len, rope) + x, latent = _inputs(num_tokens, 4) + _external_scatter(md, latent, rot) + dbg_qb = torch.zeros(num_tokens, H * 192, device="cuda", dtype=torch.bfloat16) + dbg_q = torch.zeros(num_tokens, H, LORA + ROPE, device="cuda", dtype=torch.bfloat16) + _run_fused(md, w, x, rot, debug=(dbg_qb, dbg_q)) + torch.cuda.synchronize() + q_fused, sf_fused = _q_bytes(md, num_tokens) + md.fp4_mla_state.prequantized_q.fill_(0) + q_input = dbg_q.clone() + q_rope_out = torch.empty(num_tokens, H, ROPE, device="cuda", dtype=torch.bfloat16) + q_pe = dbg_qb.view(num_tokens, H, 192)[..., 128:] + _scatter(md, latent, rot, q_pe, q_rope_out, q_input) + torch.cuda.synchronize() + q_ref, sf_ref = _q_bytes(md, num_tokens) + assert torch.equal(q_rope_out, dbg_q[..., LORA:]) + assert torch.equal(q_fused, q_ref) + assert torch.equal(sf_fused, sf_ref) + finally: + mgr.shutdown() + + +@pytest.mark.parametrize("num_tokens", [5, 8, 16, 24]) +def test_fused_norm_matches_flashinfer_norm(num_tokens): + """The q_a RMSNorm inside the kernel gives the query rows of the kernel run after + FlashInfer's norm, bit for bit; so does the KV latent row kernel.""" + from flashinfer.norm import rmsnorm + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_mla_q_fused as fused, + ) + + w = _weights() + n = max(1, num_tokens // Q_LEN) + mgr, md = _fp4_case(n, 600, seed=40 + num_tokens) + try: + rot = _rotary_table(md.max_seq_len, "identity") + gen = torch.Generator(device="cuda").manual_seed(8) + qkv = torch.randn(num_tokens, QL + LORA + ROPE, device="cuda", generator=gen).bfloat16() + q, ckv, kpe = qkv.split([QL, LORA, ROPE], -1) + latent = fused.k3_mla_latent_row(ckv, kpe, w.kv_norm, EPS) + ref_latent = torch.cat([rmsnorm(ckv, w.kv_norm, EPS), kpe], dim=-1) + torch.cuda.synchronize() + assert torch.equal(latent, ref_latent) + if num_tokens % Q_LEN or num_tokens > fused.MAX_TOKENS: + return + _external_scatter(md, latent, rot) + _run_fused(md, w, rmsnorm(q, w.q_norm, EPS), rot) + torch.cuda.synchronize() + ref = _q_bytes(md, num_tokens) + md.fp4_mla_state.prequantized_q.fill_(0) + _run_fused(md, w, q, rot, norm=(w.q_norm, EPS)) + torch.cuda.synchronize() + got = _q_bytes(md, num_tokens) + assert torch.equal(got[0], ref[0]) and torch.equal(got[1], ref[1]) + finally: + mgr.shutdown() + + +@pytest.mark.parametrize("num_tokens", [8, 16]) +def test_projections_match_production_gemms(num_tokens): + """q_b against the low-M FP8 GEMM, W_UK against the BF16 BMM: the same products, another FP32 + summation order, so BF16 results agree to rounding.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_fp8w_lowm_gemm as lowm, + ) + + w = _weights() + n = num_tokens // Q_LEN + mgr, md = _fp4_case(n, 600, seed=20 + n) + try: + rot = _rotary_table(md.max_seq_len, "identity") + x, latent = _inputs(num_tokens, 5) + _external_scatter(md, latent, rot) + dbg_qb = torch.zeros(num_tokens, H * 192, device="cuda", dtype=torch.bfloat16) + dbg_q = torch.zeros(num_tokens, H, LORA + ROPE, device="cuda", dtype=torch.bfloat16) + _run_fused(md, w, x, rot, debug=(dbg_qb, dbg_q)) + ref_qb = torch.empty_like(dbg_qb) + lowm.run_k3_fp8w_lowm_gemm(x, w.w_qb, w.s_qb_packed, ref_qb) + q = dbg_qb.view(num_tokens, H, 192) + ref_uk = torch.bmm(q[..., :128].transpose(0, 1), w.w_uk.transpose(1, 2)).transpose(0, 1) + torch.cuda.synchronize() + for got, ref in ((dbg_qb, ref_qb), (dbg_q[..., :LORA], ref_uk)): + ulps = _bf16_ulps(got, ref) + assert float(ulps.max()) <= 2.0 + assert float((got != ref).float().mean()) < 0.02 + finally: + mgr.shutdown() + + +class _DecodeAsAttention: + """``generation.attn_forward_gen`` stand-in: K3's FP4 generation decode of every head (no + Helix exchange), as ``generation.gen_attention`` runs it.""" + + def __init__(self, md): + self.md = md + self.q_rows_external = None + + def __call__( + self, layer, fused_q, md, *, helix_v_up=False, q_rows_external=False, hooks=None, **kw + ): + assert helix_v_up and kw["latent_cache"] is not None + assert md.fp4_mla_state.generation_cache_scattered + self.q_rows_external = q_rows_external + stats = torch.empty(fused_q.shape[0], H, 2, device="cuda", dtype=torch.float32) + out = k3_decode().run_generation( + _attention(), fused_q, md, softmax_stats_tensor=stats, q_rows_external=q_rows_external + ) + assert not md.fp4_mla_state.generation_cache_scattered + return out + + +def _mla_stand_in(w, md, rot, qkv): + """The attributes of the MLA layer that ``generation.forward_gen_q_fused`` reads, with the + real kernels behind them.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.mla.q_norm_pdl import ( # noqa: E501 + rmsnorm_release_at_launch, + ) + + return SimpleNamespace( + num_heads_tp=H, + num_heads_tp_cp=H, + v_head_dim=LORA, + q_lora_rank=QL, + kv_lora_rank=LORA, + qk_rope_head_dim=ROPE, + mqa=_attention(rot), + aux_stream=torch.cuda.Stream(), + _kv_events=(torch.cuda.Event(), torch.cuda.Event()), + kv_a_layernorm=SimpleNamespace(weight=w.kv_norm, variance_epsilon=EPS), + q_a_layernorm=SimpleNamespace(weight=w.q_norm, variance_epsilon=EPS), + q_b_proj=SimpleNamespace(weight=w.w_qb, weight_scale=w.s_qb_f32), + k_b_proj_trans=w.w_uk, + k_b_proj_trans_dequant=None, + _q_fused_qb_exp=None, + project_qkv_a=lambda _: qkv, + q_a_norm=lambda q: rmsnorm_release_at_launch(q, w.q_norm, EPS), + ) + + +@needs_k3_decode +@pytest.mark.parametrize("n", [1, 2]) +def test_forward_gen_q_fused_matches_default_chain(n, monkeypatch): + from flashinfer.norm import rmsnorm + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + mla as generation, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( + k3_fp8w_lowm_gemm as lowm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.mla.q_norm_pdl import ( # noqa: E501 + rmsnorm_release_at_launch, + ) + from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream + + w = _weights() + mgr, md = _fp4_case(n, 2000, seed=30 + n) + try: + num_tokens = n * Q_LEN + rot = _rotary_table(md.max_seq_len, "identity") + gen = torch.Generator(device="cuda").manual_seed(6) + qkv = torch.randn(num_tokens, QL + LORA + ROPE, device="cuda", generator=gen).bfloat16() + before = [p.clone() for p in _pools(mgr)] + + # Default chain: q RMSNorm, q_b GEMM, W_UK BMM, generation update with the query rows. + q, ckv, kpe = qkv.split([QL, LORA, ROPE], -1) + latent = torch.cat([rmsnorm(ckv, w.kv_norm, EPS), kpe], dim=-1) + qb = torch.empty(num_tokens, H * 192, device="cuda", dtype=torch.bfloat16) + lowm.run_k3_fp8w_lowm_gemm(rmsnorm_release_at_launch(q, w.q_norm, EPS), w.w_qb, + w.s_qb_packed, qb) # fmt: skip + fused_q = torch.empty(num_tokens, H, LORA + ROPE, device="cuda", dtype=torch.bfloat16) + qv = qb.view(num_tokens, H, 192) + torch.bmm(qv[..., :128].transpose(0, 1), w.w_uk.transpose(1, 2), + out=fused_q[..., :LORA].transpose(0, 1)) # fmt: skip + _scatter(md, latent, rot, qv[..., 128:], fused_q[..., LORA:], fused_q) + assert md.fp4_mla_state.generation_cache_scattered + ref_out = _DecodeAsAttention(md)(None, fused_q.view(num_tokens, -1), md, helix_v_up=True, + latent_cache=latent) # fmt: skip + ref_pools = [p.clone() for p in _pools(mgr)] + for p, p0 in zip(_pools(mgr), before): + p.copy_(p0) + + mla = _mla_stand_in(w, md, rot, qkv) + decode = _DecodeAsAttention(md) + monkeypatch.setattr(generation, "attn_forward_gen", decode) + k3_decode().begin_step(md) + with with_multi_stream(True): + out = generation.forward_gen_q_fused(mla, qkv, md) + torch.cuda.synchronize() + # The fused query wrote the query rows and marked its launch for the decode + # (features.ull2_mla_decode_pdl). + assert decode.q_rows_external is True + assert k3_decode()._step(md).query_launch_mark is not None + for p, ref in zip(_pools(mgr), ref_pools): + assert torch.equal(p, ref) + rel = float((out.float() - ref_out.float()).norm() / ref_out.float().norm()) + assert rel < 0.05, rel + finally: + mgr.shutdown() + + +# --------------------------------------------------------------------------------------------- +# The k3_mla_q_fused entry on its cells, against its plain-torch reference +# --------------------------------------------------------------------------------------------- + + +def _q_fused_entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention import ( # noqa: E501 + k3_mla_q_fused as entry_module, + ) + + return entry_module + + +def _dequant_rows(entry_module, q_fp4, q_sf, rows, gs): + """FP32 ``[rows, 576]`` values of FP4 query rows (RoPE: main + residual) and their FP32 + ``[rows, 40]`` group scales, in units of the unscaled query.""" + cols = 40 + col = torch.arange(cols, device=q_fp4.device) + offsets = entry_module.swizzled_sf_offset(rows[:, None], col[None, :], cols) + sf = q_sf.view(torch.uint8).view(-1)[offsets.flatten()].view(-1, cols) + sf = sf.view(torch.float8_e4m3fn).float() + codes = entry_module.unpack_e2m1(q_fp4[rows]) + values = entry_module.e2m1_values(codes).view(-1, cols, 16) * sf.unsqueeze(-1) + prefix = values[:, :32].flatten(1) + rope = values[:, 32:].view(-1, 4, 2, 16) + return torch.cat((prefix, (rope[:, :, 0] + rope[:, :, 1]).flatten(1)), -1) / gs, sf / gs + + +_Q_FUSED_CELLS = _q_fused_entry().k3_mla_q_fused.CELLS + + +@pytest.mark.parametrize("cell", _Q_FUSED_CELLS, ids=[c.why[:48] for c in _Q_FUSED_CELLS]) +def test_k3_mla_q_fused_cells(cell): + """The entry's ``.CELLS`` under ``validating``: every FP4 group of the written query rows + within one FP4 step (two units of its scale) of the reference's, codes equal in nearly every + group (the q_b MMA's summation order moves a BF16 value by an ulp now and then), and every + byte outside the tokens' rows untouched.""" + import zlib + + from _validating import validating + + from tensorrt_llm._torch.attention.backends.fp4_mla import FP4_MLA_Q_GLOBAL_SCALE + from tensorrt_llm._torch.attention.backends.fp4_mla.config import ( + FP4_MLA_Q_LOGICAL_DIM, + _fp4_mla_packed_rows_per_seq, + ) + from tensorrt_llm._torch.attention.backends.fp4_mla.layout import ( + _get_fp4_mla_swizzled_scale_size, + ) + + entry_module = _q_fused_entry() + op = entry_module.k3_mla_q_fused + spec = cell.spec + w = _weights() + requests, gen_len = spec["requests"], spec["gen_len"] + tokens = requests * gen_len + seed = zlib.crc32(cell.why.encode()) % 2**31 + x, _ = _inputs(tokens, seed) + rows_per_seq = _fp4_mla_packed_rows_per_seq(H, gen_len) + num_rows = requests * rows_per_seq + q_fp4 = torch.zeros((num_rows, 320), dtype=torch.uint8, device="cuda") + q_sf = torch.zeros( + _get_fp4_mla_swizzled_scale_size(num_rows, FP4_MLA_Q_LOGICAL_DIM), + dtype=torch.uint8, + device="cuda", + ) + gs = torch.tensor([FP4_MLA_Q_GLOBAL_SCALE], dtype=torch.float32, device="cuda") + rot = _rotary_table(4096, "identity") + positions = (1000 + torch.arange(tokens, device="cuda")).to(torch.int32) + norm = (w.q_norm, EPS) if spec["norm"] else None + if norm is None: + x = _flashinfer_norm(x, w.q_norm) + args = (x, w.w_qb, w.qb_exp, w.w_uk, gs, rot, positions, q_fp4, q_sf, gen_len, rows_per_seq, + True, norm, spec["early_trigger"]) # fmt: skip + ref_q, ref_sf = op.reference(*args) + with validating(op): + op(*args) + torch.cuda.synchronize() + + t = torch.arange(tokens, device="cuda") + token_rows = t // gen_len * rows_per_seq + t % gen_len * H + rows = (token_rows[:, None] + torch.arange(H, device="cuda")[None, :]).flatten() + written = torch.zeros(num_rows, dtype=torch.bool, device="cuda") + written[rows] = True + assert torch.equal(q_fp4[~written], ref_q[~written]), "bytes outside the tokens' rows moved" + got, got_sf = _dequant_rows(entry_module, q_fp4, q_sf, rows, float(gs)) + ref, ref_sf = _dequant_rows(entry_module, ref_q, ref_sf, rows, float(gs)) + step = 2 * torch.maximum(got_sf, ref_sf) + step = torch.cat((step[:, :32].repeat_interleave(16, -1), + step[:, 32::2].repeat_interleave(16, -1)), -1) # fmt: skip + assert bool(((got - ref).abs() <= step).all()), cell.why + groups_equal = torch.eq(q_fp4[rows].view(-1, 40, 8), ref_q[rows].view(-1, 40, 8)).all(-1) + assert float(groups_equal.float().mean()) >= 0.95, cell.why + + +def _flashinfer_norm(x, weight): + from flashinfer.norm import rmsnorm + + out = torch.empty_like(x) + rmsnorm(x, weight, EPS, out=out) + return out diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull2_mla.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull2_mla.py new file mode 100644 index 000000000000..8e4043460145 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull2_mla.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi-K3 ULL MLA attention side against today's path, one GPU. + +``features.ull2_mla_decode_pdl``: the fused query (``k3_mla_q_fused``) releases its dependents at +entry and K3's FP4 packed decode (``ops/attention/_fp4_mla_decode``), launched as its programmatic +dependent, issues everything but the query load (barriers, TMEM, plan, page table, the first KV +tile) before ``griddepcontrol.wait``. The chain is the module's: the step's decode metadata written +ahead (``features.ull4_mla_plan_early``), KV-only update on a side stream joined before the +decode, fused query on the main stream (its launch marked for the decode). Checks: decode output +and Helix statistics bitwise those of the normal launch, eager and over CUDA graph replays with +new inputs; the decode falls back to a normal launch when a metadata kernel (plan, lengths) runs in +its own call or the query rows come from the cache update. + +``features.ull2_mla_tail`` launch policy: at most 32 tokens, the decode whose reduction the fused +tail folds (``DeferredPackedReduce(release_early=False)``) releases its dependents from its +epilogue instead of at entry; the recorded reduction and the decode outputs are unchanged. + +The FP4 cache helpers are those of ``test_k3_mla_q_fused.py`` (same directory, importable by name +under pytest's default ``prepend`` import mode). +""" + +import math + +import pytest +import test_k3_mla_q_fused as qf +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="needs an SM107 GPU", +) + +H, LORA, ROPE, Q_LEN = qf.H, qf.LORA, qf.ROPE, qf.Q_LEN + + +class _Recorder: + """Records the ``pdl_wait`` and ``launch_dependents`` of every packed decode launch; the keys + of ``override`` replace the decode's own arguments.""" + + def __init__(self, module): + self.module = module + self.inner = module.run_trtllm_fp4_mla_decode_packed + self.calls = [] + self.releases = [] + self.override = {} + + def __call__(self, *args, **kwargs): + kwargs.update(self.override) + self.calls.append(bool(kwargs.get("pdl_wait", False))) + self.releases.append(kwargs.get("launch_dependents", False)) + return self.inner(*args, **kwargs) + + +def _record_decode(monkeypatch): + """Wrap the packed decode of K3's kernel module in a :class:`_Recorder`.""" + module = qf.k3_decode()._kernel_module() + recorder = _Recorder(module) + monkeypatch.setattr(module, "run_trtllm_fp4_mla_decode_packed", recorder) + return recorder + + +class _Chain: + """The fused-query generation chain of one layer on a real FP4 cache: KV-only update on a + side stream, ``k3_mla_q_fused`` on the main stream, join, K3's FP4 packed decode (+ split + reduction) with Helix statistics.""" + + def __init__(self, n, kv, seed): + self.mgr, self.md = qf._fp4_case(n, kv, seed) + self.w = qf._weights() + self.T = n * Q_LEN + self.rot = qf._rotary_table(self.md.max_seq_len, "identity") + self.x, self.latent = qf._inputs(self.T, seed + 1) + self.side = torch.cuda.Stream() + self.fork, self.ready = torch.cuda.Event(), torch.cuda.Event() + self.out = torch.empty(self.T, H, LORA, device="cuda", dtype=torch.bfloat16) + self.stats = torch.empty(self.T, H, 2, device="cuda", dtype=torch.float32) + self.q_stub = torch.empty(self.T, H, LORA + ROPE, device="cuda", dtype=torch.bfloat16) + # The first update writes the Q buffers the fused query fills. + qf._external_scatter(self.md, self.latent, self.rot) + torch.cuda.synchronize() + + def decode(self, q_rows_external=True): + state = self.md.fp4_mla_state + qf.k3_decode().run_decode( + self.md, + 0, + 0, + self.q_stub, + self.out, + sm_scale=1.0 / math.sqrt(192.0), + kv_lora_rank=LORA, + qk_rope_head_dim=ROPE, + prequantized_q=state.prequantized_q, + prequantized_q_sf=state.prequantized_q_sf, + q_batch_capacity=state.q_batch_capacity, + softmax_stats_tensor=self.stats, + q_rows_external=q_rows_external, + ) + + def begin_step(self, plan_early=True): + """A new step: the previous plan is stale; with ``plan_early`` the step's decode + metadata is written ahead on the main stream (``features.ull4_mla_plan_early``).""" + dec = qf.k3_decode() + dec.begin_step(self.md) + if plan_early: + assert dec.prelaunch_generation_metadata( + self.md, num_heads=H, use_helix=True, device=self.out.device + ) + + def run(self, early_trigger, before_decode=None, plan_early=True, new_step=True): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( # noqa: E501 + k3_mla_q_fused as fused, + ) + + md = self.md + state = md.fp4_mla_state + if new_step: + self.begin_step(plan_early) + self.fork.record() + with torch.cuda.stream(self.side): + self.fork.wait() + qf._external_scatter(md, self.latent, self.rot) + self.ready.record() + fused.k3_mla_q_fused( + self.x, + self.w.w_qb, + self.w.qb_exp, + self.w.w_uk, + state.q_global_scale, + self.rot, + md.helix_position_offsets, + state.prequantized_q, + state.prequantized_q_sf.view(torch.uint8), + gen_len=Q_LEN, + rows_per_seq=qf._rows_per_seq(md), + norm=(self.w.q_norm, qf.EPS), + early_trigger=early_trigger, + ) + qf.k3_decode().mark_query_launch(md) + self.ready.wait() + if before_decode is not None: + before_decode() + self.decode() + + def result(self): + torch.cuda.synchronize() + return self.out.clone(), self.stats.clone() + + +def _assert_bitwise(got, ref): + assert torch.equal(got[0].view(torch.int16), ref[0].view(torch.int16)) + assert torch.equal(got[1].view(torch.int32), ref[1].view(torch.int32)) + + +@qf.needs_k3_decode +@pytest.mark.parametrize("n", [1, 2]) +def test_decode_pdl_matches_normal_launch(n, monkeypatch, k3_features): + recorder = _record_decode(monkeypatch) + chain = _Chain(n, 3000, seed=50 + n) + try: + k3_features(ull2_mla_decode_pdl=False) + chain.run(early_trigger=False) + ref = chain.result() + assert recorder.calls[-1] is False + k3_features(ull2_mla_decode_pdl=True) + for _ in range(2): + chain.out.fill_(torch.nan) + chain.run(early_trigger=True) + _assert_bitwise(chain.result(), ref) + assert recorder.calls[-1] is True + # A normal decode after the early-released query is also exact. + k3_features(ull2_mla_decode_pdl=False) + chain.run(early_trigger=True) + _assert_bitwise(chain.result(), ref) + assert recorder.calls[-1] is False + finally: + chain.mgr.shutdown() + + +@qf.needs_k3_decode +@pytest.mark.parametrize("n", [1, 2]) +def test_decode_pdl_graph_replay(n, k3_features): + """Captured with the programmatic edge, replayed on new queries and new KV rows (the decode's + early KV tile loads include this step's tokens: about two tiles per split task): each replay + equals the eager normal-launch chain on the same inputs.""" + k3_features(ull2_mla_decode_pdl=True) + chain = _Chain(n, 5000, seed=60 + n) + try: + chain.run(early_trigger=True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + chain.run(early_trigger=True) + torch.cuda.synchronize() + gen = torch.Generator(device="cuda").manual_seed(9) + for _ in range(3): + chain.x.copy_(torch.randn(chain.x.shape, device="cuda", generator=gen).bfloat16()) + chain.latent.copy_( + (torch.randn(chain.latent.shape, device="cuda", generator=gen) * 0.5).bfloat16() + ) + chain.out.fill_(torch.nan) + graph.replay() + got = chain.result() + k3_features(ull2_mla_decode_pdl=False) + chain.run(early_trigger=False) + ref = chain.result() + k3_features(ull2_mla_decode_pdl=True) + _assert_bitwise(got, ref) + finally: + chain.mgr.shutdown() + + +@qf.needs_k3_decode +def test_decode_pdl_falls_back(monkeypatch, k3_features): + """No programmatic launch when the decode call itself launches a metadata kernel (a stale + plan, or a step whose plan was not written ahead) or when the cache update wrote the query + rows.""" + k3_features(ull2_mla_decode_pdl=True) + recorder = _record_decode(monkeypatch) + chain = _Chain(1, 3000, seed=70) + try: + chain.run(early_trigger=True) + ref = chain.result() + assert recorder.calls[-1] is True + # A stale plan noticed by the decode call itself: the plan kernel runs right before it. + chain.run( + early_trigger=True, + before_decode=lambda: setattr(qf.k3_decode()._step(chain.md), "plan_key", None), + ) + _assert_bitwise(chain.result(), ref) + assert recorder.calls[-1] is False + # The first layer of a step without the plan written ahead. + chain.run(early_trigger=True, plan_early=False) + _assert_bitwise(chain.result(), ref) + assert recorder.calls[-1] is False + chain.decode(q_rows_external=False) + _assert_bitwise(chain.result(), ref) + assert recorder.calls[-1] is False + finally: + chain.mgr.shutdown() + + +@qf.needs_k3_decode +@pytest.mark.parametrize("release_early", [True, False]) +def test_deferred_reduce_release(release_early, monkeypatch, k3_features): + """``DeferredPackedReduce``: with ``release_early`` the decode releases its dependents at + entry (``features.fp4_mla_reduce_pdl``), without it from its epilogue (when its producer warps + finish the last KV tile); the reduction it recorded, launched afterwards, gives the outputs of + the plain decode bitwise.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.fp4_mla_packed_reduce import ( # noqa: E501 + DeferredPackedReduce, + ) + + k3_features(fp4_mla_reduce_pdl=True) + recorder = _record_decode(monkeypatch) + module = recorder.module + # The entry release stays the boolean True of the early modes. + assert module.LAUNCH_DEPENDENTS_ENTRY == True # noqa: E712 + assert module.LAUNCH_DEPENDENTS_EPILOGUE not in (0, 1) + chain = _Chain(1, 6000, seed=90) + try: + chain.run(early_trigger=False) + ref = chain.result() + assert recorder.releases[-1] == module.LAUNCH_DEPENDENTS_ENTRY + chain.out.fill_(torch.nan) + with DeferredPackedReduce(release_early=release_early) as deferred: + chain.run(early_trigger=False) + expected = ( + module.LAUNCH_DEPENDENTS_ENTRY if release_early else module.LAUNCH_DEPENDENTS_EPILOGUE + ) + assert recorder.releases[-1] == expected + assert deferred.args is not None, "the split reduction was not deferred" + deferred.run() + _assert_bitwise(chain.result(), ref) + finally: + chain.mgr.shutdown() + + +def test_late_tail_policy(k3_features): + """The ULL2 tail is released by the decode's completion up to 32 tokens.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.mla import ( # noqa: E501 + decode_releases_tail_early, + ) + + k3_features(ull2_mla_tail=False) + assert decode_releases_tail_early(8) and decode_releases_tail_early(32) + k3_features(ull2_mla_tail=True) + assert not decode_releases_tail_early(8) + assert not decode_releases_tail_early(16) + assert not decode_releases_tail_early(32) + assert decode_releases_tail_early(33) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull3_mla.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull3_mla.py new file mode 100644 index 000000000000..4811a8f9bd17 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull3_mla.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi-K3 ULL3 MLA attention-side scheduling against the late-released chain, one GPU. + +The epilogue release: K3's FP4 packed decode whose split reduction the fused Helix tail folds +(``DeferredPackedReduce(release_early=False)``) releases the tail when its producer warps finish +the last KV tile instead of at its completion. The tail, launched as its programmatic dependent, +stages ``W_UV`` and the plan during the decode's end and folds after ``griddepcontrol.wait``. +Checks: the tail output (fold, V up-projection, combine, MXFP8 o_proj operand) is bitwise that of +the chain whose decode releases nothing (the tail waits for its completion), eagerly and over CUDA +graph replays, at n1 / n2 with FP16 and FP32 split partials. + +``features.ull3_mla_gate_at``: the per-batch placement policy of the side-stream output gate +parses as documented. + +The FP4 cache and chain helpers are those of ``test_k3_ull2_mla.py`` / ``test_k3_mla_q_fused.py`` +(same directory). +""" + +import pytest +import test_k3_ull2_mla as u2 +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="needs an SM107 GPU", +) + +H, LORA, Q_LEN = u2.H, u2.LORA, u2.Q_LEN + + +class _Tail: + """The fused tail on a local receive buffer (``P = 1``): fold, V up-projection, push into its + own buffer, combine, MXFP8 operand of ``bf16(out * gate)``.""" + + def __init__(self, num_tokens, seed): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels import ( # noqa: E501 + k3_helix_tail as tail_mod, + ) + + self.mod = tail_mod + gen = torch.Generator(device="cuda") + gen.manual_seed(seed) + self.w_uv = (torch.randn((H, 128, LORA), device="cuda", generator=gen) * 0.05).to( + torch.bfloat16 + ) + k = H * 128 + self.gate = torch.sigmoid(torch.randn((num_tokens, k), device="cuda", generator=gen)).to( + torch.bfloat16 + ) + self.out = torch.empty((num_tokens, k), dtype=torch.bfloat16, device="cuda") + self.xq = torch.empty((num_tokens, k), dtype=torch.float8_e4m3fn, device="cuda") + self.xs = torch.empty((num_tokens, k // 32), dtype=torch.uint8, device="cuda") + self.half_bytes = num_tokens * tail_mod.record_bytes(H, 128) + 4096 + self.buffer = torch.full((2 * self.half_bytes // 4,), -1, dtype=torch.int32, device="cuda") + self.state = torch.zeros(8, dtype=torch.int64, device="cuda") + + def launch(self, rows, stats, fold): + self.mod.k3_helix_tail( + rows, + stats, + self.w_uv, + self.out, + fold=fold, + zero_kv_mask=None, + cp_rank=0, + cp_size=1, + ws_addr=self.buffer.data_ptr(), + rank_stride=self.buffer.numel() * 4, + half_bytes=self.half_bytes, + state=self.state, + max_ctas=torch.cuda.get_device_properties(0).multi_processor_count, + use_pdl=True, + mx_out=(self.gate, self.xq, self.xs), + ull2=True, + ) + + def result(self): + torch.cuda.synchronize() + return self.out.clone(), self.xq.view(torch.uint8).clone(), self.xs.clone() + + +def _assert_same(got, ref): + """Tail output (BF16), MXFP8 codes and exponents bit for bit.""" + assert torch.equal(got[0].view(torch.int16), ref[0].view(torch.int16)) + assert torch.equal(got[1], ref[1]) and torch.equal(got[2], ref[2]) + + +def _layer(chain, tail): + """Decode with the split reduction deferred (late release), then the fused tail on it.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.fp4_mla_packed_reduce import ( # noqa: E501 + DeferredPackedReduce, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.mla import ( # noqa: E501 + fold_inputs_of, + ) + + with DeferredPackedReduce(release_early=False) as deferred: + chain.run(early_trigger=True) + fold = fold_inputs_of(deferred, chain.out.view(chain.T, -1), chain.stats) + assert fold is not None, "the split reduction was not deferred" + tail.launch(chain.out.view(chain.T * H, LORA), chain.stats.view(chain.T * H, 2), fold) + + +@u2.qf.needs_k3_decode +@pytest.mark.parametrize("f16_partials", [True, False], ids=["fp16", "fp32"]) +@pytest.mark.parametrize("n", [1, 2]) +def test_tail_trigger_matches_late_release(n, f16_partials, monkeypatch, k3_features): + k3_features( + ull2_mla_tail=True, + ull2_mla_decode_pdl=True, + fp4_mla_reduce_pdl=True, + fp4_mla_f16_partials=f16_partials, + ) + recorder = u2._record_decode(monkeypatch) + module = recorder.module + chain = u2._Chain(n, 6000 + 1500 * n, seed=70 + n) + tail = _Tail(chain.T, seed=5 + n) + try: + # The reference: the decode releases nothing, the tail waits for its completion. + recorder.override = {"launch_dependents": False} + _layer(chain, tail) + ref = tail.result() + assert recorder.releases[-1] is False + assert torch.isfinite(ref[0].float()).all() and ref[0].abs().max() > 0 + recorder.override = {} + for _ in range(2): + tail.out.fill_(torch.nan) + _layer(chain, tail) + _assert_same(tail.result(), ref) + assert recorder.releases[-1] == module.LAUNCH_DEPENDENTS_EPILOGUE + # CUDA graph: the programmatic edges decode -> tail as the model captures them. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _layer(chain, tail) + torch.cuda.synchronize() + for _ in range(3): + tail.out.fill_(torch.nan) + graph.replay() + _assert_same(tail.result(), ref) + assert bool((tail.buffer == -1).all()), "the receive buffer was not re-armed" + finally: + chain.mgr.shutdown() + + +def test_gate_place_policy(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + mla as k3mla, + ) + + assert k3mla._parse_ull3_gate_at("decode") == ("decode", 16, "decode") + assert k3mla._parse_ull3_gate_at("exchange:16,start") == ("exchange", 16, "start") + assert k3mla._parse_ull3_gate_at("exchange") == ("exchange", 16, "exchange") + assert k3mla._parse_ull3_gate_at("exchange:8,start") == ("exchange", 8, "start") + for bad in ("early", "exchange,late", "start:x", "query", "tail"): + with pytest.raises(ValueError): + k3mla._parse_ull3_gate_at(bad) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull4_mla.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull4_mla.py new file mode 100644 index 000000000000..d43ebdb27983 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/mla/test_k3_ull4_mla.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi-K3 ULL4 MLA attention-side scheduling switch against today's path, one GPU. + +``features.ull4_mla_plan_early``: the step's FP4 MLA metadata (generation lengths fused with K3's +packed split plan, the device page table) written ahead on a side stream +(``_fp4_mla_decode.prelaunch_generation_metadata``): the layer's decode then launches no metadata +kernel and gives the same cache bytes, outputs and statistics. + +Helpers: ``test_k3_ull2_mla.py`` (FP4 chain, same directory). +""" + +import math + +import pytest +import test_k3_ull2_mla as u2 +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="needs an SM107 GPU", +) + + +def _plan(md): + """K3's packed split plan workspace of the step.""" + return md.fp4_mla_state.workspaces[u2.qf.k3_decode()._PLAN_WS] + + +@u2.qf.needs_k3_decode +@pytest.mark.parametrize("n", [1, 2]) +def test_step_metadata_prelaunch(n, k3_features): + """The step metadata written ahead on a side stream: the chain's decode launches no metadata + kernel, and the cache bytes, decode output and statistics are bitwise those of the chain whose + decode writes the metadata itself; also captured in a CUDA graph.""" + dec = u2.qf.k3_decode() + k3_features(ull2_mla_decode_pdl=True) + chain = u2._Chain(n, 5000 + 1000 * n, seed=80 + n) + md = chain.md + state = md.fp4_mla_state + try: + # What the metadata resets at every step (Fp4MlaState.invalidate_generation_lengths). + state.invalidate_generation_lengths() + chain.run(early_trigger=True, plan_early=False) + # Eager plans hold spare task slots the plan kernel leaves untouched: compare against a + # reference written over the same fill. + _plan(md).fill_(-1) + state.invalidate_generation_lengths() + chain.run(early_trigger=True, plan_early=False) + ref = chain.result() + ref_plan = _plan(md).clone() + ref_lens = state.generation_kv_lens.clone() + pools = [p.clone() for p in u2.qf._pools(chain.mgr)] + assert dec._step(md).plan_key is not None + + def prelaunched_run(): + state.invalidate_generation_lengths() + dec.begin_step(md) + side = torch.cuda.Stream() + fork, done = torch.cuda.Event(), torch.cuda.Event() + fork.record() + with torch.cuda.stream(side): + fork.wait() + launched = dec.prelaunch_generation_metadata( + md, num_heads=u2.H, use_helix=True, device=torch.device("cuda") + ) + done.record() + torch.cuda.current_stream().wait_event(done) + before = dec._step(md).metadata_launches + chain.run(early_trigger=True, new_step=False) + return launched, dec._step(md).metadata_launches - before + + chain.out.fill_(torch.nan) + _plan(md).fill_(-1) + plan_ptr = _plan(md).data_ptr() + launched, extra_launches = prelaunched_run() + assert launched and extra_launches == 0 + # An index-less device must not reallocate the plan workspace. + assert _plan(md).data_ptr() == plan_ptr + u2._assert_bitwise(chain.result(), ref) + assert torch.equal(_plan(md), ref_plan) + assert torch.equal(state.generation_kv_lens, ref_lens) + for got, want in zip(u2.qf._pools(chain.mgr), pools): + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + # Captured: the prelaunch, the join and the chain in one graph. + prelaunched_run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + prelaunched_run() + for _ in range(2): + chain.out.fill_(torch.nan) + _plan(md).fill_(-1) + graph.replay() + u2._assert_bitwise(chain.result(), ref) + assert torch.equal(_plan(md), ref_plan) + assert math.isfinite(float(ref[0].float().abs().max())) + finally: + chain.mgr.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/test_k3_mla_latent_row.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/test_k3_mla_latent_row.py new file mode 100644 index 000000000000..64b5b878b6fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/test_k3_mla_latent_row.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The ``k3_mla_latent_row`` entry (kv_a RMSNorm of the latent and the K-pe copy in one launch) +on its ``.CELLS``: the norm within one BF16 rounding of the FP32 reference, the K-pe columns +copied exactly, and the row bitwise FlashInfer's ``rmsnorm`` followed by the copy.""" + +import zlib + +import pytest +import torch + +__extra_import_path__ = ["../../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 3), (10, 7)), + reason="certified on sm_103 and sm_107", +) + +LORA, ROPE = 512, 64 + + +def _entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.k3_mla_latent_row import ( # noqa: E501 + k3_mla_latent_row, + ) + + return k3_mla_latent_row + + +@pytest.mark.parametrize("cell", _entry().CELLS, ids=[c.why[:40] for c in _entry().CELLS]) +@pytest.mark.parametrize("use_pdl", [True, False]) +def test_k3_mla_latent_row_cells(cell, use_pdl): + from flashinfer.norm import rmsnorm + + op = _entry() + spec = cell.spec + tokens, stride, eps = spec["tokens"], spec["row_stride"], spec["eps"] + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + proj = torch.randn(tokens, stride, device="cuda", generator=gen).to(torch.bfloat16) + c_kv = proj[:, stride - LORA - ROPE : stride - ROPE] + k_pe = proj[:, stride - ROPE :] + weight = (1 + 0.1 * torch.randn(LORA, device="cuda", generator=gen)).to(torch.bfloat16) + with validating(op): + out = op(c_kv, k_pe, weight, eps, use_pdl=use_pdl) + assert out.shape == (tokens, LORA + ROPE) and out.is_contiguous() + expected = op.reference(c_kv, k_pe, weight, eps, use_pdl=use_pdl) + assert_within_ulp(out[:, :LORA], expected[:, :LORA], element_ulp=1, rms_ulp=0.5) + assert torch.equal(out[:, LORA:], expected[:, LORA:]), cell.why + from flashinfer.norm import _use_cuda_norm + + if _use_cuda_norm(c_kv.device): + return # FlashInfer runs its CUDA JIT kernel here, not the CuTe DSL one repeated + normed = torch.empty((tokens, LORA), dtype=torch.bfloat16, device="cuda") + rmsnorm(c_kv, weight, eps, out=normed) + assert torch.equal(out[:, :LORA].view(torch.int16), normed.view(torch.int16)), cell.why diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/test_rmsnorm_release_at_launch.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/test_rmsnorm_release_at_launch.py new file mode 100644 index 000000000000..929946678f88 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/attention/test_rmsnorm_release_at_launch.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The ``rmsnorm_release_at_launch`` entry (the K3 q_a RMSNorm with its PDL trigger at launch) on +its ``.CELLS``: within one BF16 rounding of its FP32 reference and bitwise FlashInfer's +``rmsnorm`` (the same kernel code), on the strided q_a rows of the fused projection output.""" + +import zlib + +import pytest +import torch + +__extra_import_path__ = ["../../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 3), (10, 7)), + reason="certified on sm_103 and sm_107", +) + + +def _entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.attention.rmsnorm_release_at_launch import ( # noqa: E501 + rmsnorm_release_at_launch, + ) + + return rmsnorm_release_at_launch + + +@pytest.mark.parametrize("cell", _entry().CELLS, ids=[c.why[:40] for c in _entry().CELLS]) +def test_rmsnorm_release_at_launch_cells(cell): + from flashinfer.norm import rmsnorm + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.mla import ( # noqa: E501 + q_norm_pdl, + ) + + op = _entry() + spec = cell.spec + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + rows = torch.randn(spec["tokens"], spec["row_stride"], device="cuda", generator=gen) + x = rows.to(torch.bfloat16)[:, : spec["width"]] + weight = (1 + 0.1 * torch.randn(spec["width"], device="cuda", generator=gen)).to(torch.bfloat16) + eps = spec["eps"] + assert q_norm_pdl.can_use(x, weight), cell.why + with validating(op): + y = op(x, weight, eps) + assert y.is_contiguous() and y.shape == x.shape and y.dtype == torch.bfloat16 + # One BF16 rounding of the FP32 result, against the row's largest magnitude. + assert_within_ulp(y, op.reference(x, weight, eps), element_ulp=1, rms_ulp=0.5) + from flashinfer.norm import _use_cuda_norm + + if _use_cuda_norm(x.device): + return # FlashInfer runs its CUDA JIT kernel here, not the CuTe DSL one repeated + expected = torch.empty_like(y) + rmsnorm(x, weight, eps, out=expected, enable_pdl=True) + assert torch.equal(y.view(torch.int16), expected.view(torch.int16)), cell.why diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_cells.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_cells.py new file mode 100644 index 000000000000..e3b0070badce --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_cells.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The single-GPU FP8-weight GEMM entries of the K3 default target over their certified cells: +``k3_fp8w_lowm_gemm.CELLS``, ``k3_fp8w_mid_gemm.CELLS``, ``k3_fp8w_lowm_mxfp8_gemm.CELLS`` and +``k3_latent_down_fp4.CELLS``, each driven through ``validating`` and checked against the entry's +own ``reference``. (``k3_mxfp8_gemm.CELLS`` run in ``core/test_core_mxfp8_gemm.py``, on SM107.) + +The weights are random E4M3 codes with power-of-two 128 x 128 block scales, and the activation rows +have a different magnitude per 32-column group, so every UE8M0 exponent path is exercised. +""" + +import importlib + +import pytest +import torch + +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 + +K3 = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() not in (100, 103, 107), + reason="the low-M FP8-weight kernels run on SM100 / SM103 / SM107", +) + +lowm = importlib.import_module(f"{K3}.kernels.k3_fp8w_lowm_gemm") +mx = importlib.import_module(f"{K3}.kernels.k3_fp8w_lowm_mxfp8") +ref = importlib.import_module(f"{K3}.ops.gemm._fp8_reference") +k3_fp8w_lowm_gemm = importlib.import_module(f"{K3}.ops.gemm.k3_fp8w_lowm_gemm").k3_fp8w_lowm_gemm +k3_fp8w_mid_gemm = importlib.import_module(f"{K3}.ops.gemm.k3_fp8w_mid_gemm").k3_fp8w_mid_gemm +k3_fp8w_lowm_mxfp8_gemm = importlib.import_module( + f"{K3}.ops.gemm.k3_fp8w_lowm_mxfp8_gemm" +).k3_fp8w_lowm_mxfp8_gemm +k3_latent_down_fp4 = importlib.import_module(f"{K3}.ops.gemm.k3_latent_down_fp4").k3_latent_down_fp4 + +#: The band of every BF16 GEMM here. Kernel and reference multiply the same E4M3 operands with +#: power-of-two scales, so every product is exact; they differ only in how the products are summed +#: (FP32 split-K partials against FP64), which leaves the sums within ~1e-6 of each other, far +#: below a BF16 ulp. The single rounding to BF16 can then fall on the other side of a boundary: +#: one BF16 ulp of the element, which is at most 2 units of ``ULP`` (2^-8) of the row's largest +#: element. RMS: such flips are rare, well under 1 unit. +ELEMENT_ULP, RMS_ULP = 2, 1 + + +def _ids(cells): + return [c.why[:48] for c in cells] + + +def _weight(n: int, k: int, seed: int): + """UE8M0-resmoothed E4M3 ``[n, k]`` weight and its DeepGEMM-packed scale.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, device="cuda", generator=gen) * 0.5).to(torch.float8_e4m3fn) + blocks = (-(-n // 128), k // 128) + scale = torch.exp2(torch.randint(-8, -2, blocks, device="cuda", generator=gen).float()) + return lowm.prepare_k3_fp8w_lowm_weight(w, scale) + + +def _rows(m: int, k: int, seed: int) -> torch.Tensor: + """BF16 ``[m, k]`` with one magnitude per 32-column group (2^-6 .. 2^2).""" + gen = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(m, k, device="cuda", generator=gen) + mag = torch.randint(-6, 3, (m, k // 32), device="cuda", generator=gen).float() + return (x * torch.exp2(mag).repeat_interleave(32, dim=1)).to(torch.bfloat16) + + +@pytest.mark.parametrize("cell", k3_fp8w_lowm_gemm.CELLS, ids=_ids(k3_fp8w_lowm_gemm.CELLS)) +def test_k3_fp8w_lowm_gemm_cells(cell): + s = cell.spec + n, k, m, variant = s["n"], s["k"], s["m"], s["variant"] + w, w_scale = _weight(n, k, 1) + x = _rows(m, k, 2) + kwargs = {} + if variant in ("out", "gated"): + kwargs["out"] = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + if variant == "gated": + gen = torch.Generator(device="cuda").manual_seed(3) + kwargs["gate"] = torch.rand(m, k, device="cuda", generator=gen).to(torch.bfloat16) + if variant == "sigmoid": + kwargs["sigmoid"] = True + with validating(k3_fp8w_lowm_gemm): + got = k3_fp8w_lowm_gemm(x, w, w_scale, **kwargs) + if "out" in kwargs: + assert got.data_ptr() == kwargs["out"].data_ptr() + want = k3_fp8w_lowm_gemm.reference(x, w, w_scale, **kwargs) + assert_within_ulp(got, want, ELEMENT_ULP, RMS_ULP) + + +@pytest.mark.parametrize("cell", k3_fp8w_mid_gemm.CELLS, ids=_ids(k3_fp8w_mid_gemm.CELLS)) +def test_k3_fp8w_mid_gemm_cells(cell): + n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"] + plan = lowm.mid_prequant_plan(m, n, k) + assert plan is not None, "a mid-M cell must be above the low-M cut-over" + tactic = plan[0] + w, w_scale = _weight(n, k, 4) + xq, xs = lowm.quantize_mid_operand(_rows(m, k, 5), tactic) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + with validating(k3_fp8w_mid_gemm): + k3_fp8w_mid_gemm(xq, xs, w, w_scale, out, tactic) + want = k3_fp8w_mid_gemm.reference(xq, xs, w, w_scale, out, tactic) + assert_within_ulp(out, want, ELEMENT_ULP, RMS_ULP) + + +@pytest.mark.parametrize( + "cell", k3_fp8w_lowm_mxfp8_gemm.CELLS, ids=_ids(k3_fp8w_lowm_mxfp8_gemm.CELLS) +) +def test_k3_fp8w_lowm_mxfp8_gemm_cells(cell): + s = cell.spec + n, k, m = s["n"], s["k"], s["m"] + w, w_scale = _weight(n, k, 6) + xq, xs = mx.mxfp8_quantize_reference(_rows(m, k, 7)) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + with validating(k3_fp8w_lowm_mxfp8_gemm): + k3_fp8w_lowm_mxfp8_gemm(xq, xs, w, w_scale, out, stage_weights=s["stage"]) + want = k3_fp8w_lowm_mxfp8_gemm.reference(xq, xs, w, w_scale, out, stage_weights=s["stage"]) + assert_within_ulp(out, want, ELEMENT_ULP, RMS_ULP) + + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def _nvfp4_values(codes, scales): + """FP64 values of NVFP4 ``codes [M, N / 2]`` / linear E4M3 ``scales [M, N / 16]`` in the scaled + domain (``global_sf`` not divided out, so every value is exact) and the per-element scale.""" + lut = torch.tensor(_E2M1, dtype=torch.float64, device=codes.device) + c = codes.long() + e2m1 = torch.stack([lut[c & 0xF], lut[c >> 4]], dim=-1).flatten(-2) + s = scales.view(torch.float8_e4m3fn).to(torch.float64).repeat_interleave(16, dim=1) + return e2m1 * s, s + + +@pytest.mark.parametrize("cell", k3_latent_down_fp4.CELLS, ids=_ids(k3_latent_down_fp4.CELLS)) +def test_k3_latent_down_fp4_cells(cell): + """NVFP4 of the BF16 projection. Where the kernel's BF16 value is one ulp off the reference's + (see ``ELEMENT_ULP``), the block maximum or an element near an E2M1 midpoint can quantize the + other way, so codes and scales are compared by value: every element within one E2M1 step of + its block (``|a - b| <= 2 * max(scale_a, scale_b)``: each quantization is within half a step, + at most one unit of its scale), and at most 1% of the codes / scales differing at all.""" + m = cell.spec["m"] + n, k = 3584, 7168 + w, w_scale = _weight(n, k, 8) + x = _rows(m, k, 9) + amax = float(ref.lowm_gemm(x, w, w_scale).abs().amax()) + global_sf = torch.tensor([448.0 * 6.0 / amax], dtype=torch.float32, device="cuda") + with validating(k3_latent_down_fp4): + codes, scales = k3_latent_down_fp4(x, w, w_scale, global_sf) + want_codes, want_scales = k3_latent_down_fp4.reference(x, w, w_scale, global_sf) + assert codes.shape == want_codes.shape and scales.shape == want_scales.shape + assert codes.dtype == want_codes.dtype and scales.dtype == want_scales.dtype + got, got_s = _nvfp4_values(codes, scales) + want, want_s = _nvfp4_values(want_codes, want_scales) + bound = 2.0 * torch.maximum(got_s, want_s) + excess = float(((got - want).abs() - bound).amax()) + assert excess <= 0, f"an element is more than one E2M1 step off ({excess:.3e})" + code_diff = float((codes != want_codes).float().mean()) + scale_diff = float((scales != want_scales).float().mean()) + assert code_diff <= 0.01 and scale_diff <= 0.01, (code_diff, scale_diff) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_push_cells.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_push_cells.py new file mode 100644 index 000000000000..bfe7838d8504 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/gemm/test_k3_fp8w_gemm_push_cells.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The collective FP8-weight GEMM entries of the K3 default target over their certified cells, on a +4-rank MNNVL group: ``k3_fp8w_lowm_gemm_push.CELLS``, ``k3_fp8w_mid_gemm_push.CELLS`` and +``k3_fp8w_lowm_mxfp8_push.CELLS``. + +Every rank draws every rank's inputs from the same seeds, so each holds the group's inputs in rank +order and checks its own ``k3_lamport_ar.run_lamport_finish`` output against the entry's collective +``reference``. Each rank pushes under ``validating``; a cell that only SM107 admits (``sm_107`` in +its spec) is skipped elsewhere and reported. + +Band: each rank's BF16 GEMM output can be one BF16 ulp off the reference's (see +``test_k3_fp8w_gemm_cells.ELEMENT_ULP``) and the finish sums four such parts in FP32 before its +own rounding; four one-ulp parts plus the rounding of the sum stay within 8 units of ``ULP`` +(2^-8) of the row's largest element, and the RMS within 2. +""" + +import os +import pickle +import sys +import traceback +from pathlib import Path + +import cloudpickle +import pytest +import torch +from mpi4py import MPI +from utils.util import skip_pre_blackwell + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL) + +pytestmark = pytest.mark.threadleak(enabled=False) + +K3 = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) +#: tests/unittest/_torch/modeling_v2, where ``_validating`` lives (imported by path in the ranks). +MV2_TESTS = str(Path(__file__).resolve().parents[5]) +WORLD = 4 +#: One Lamport buffer: the largest cell (128 x 7168 two-shot over 4 ranks) needs 3.7 MB. +BUFFER_BYTES = 8 << 20 +ELEMENT_ULP, RMS_ULP = 8, 2 +ENTRIES = ("k3_fp8w_lowm_gemm_push", "k3_fp8w_mid_gemm_push", "k3_fp8w_lowm_mxfp8_push") + + +def _weight(n, k, seed): + """Random E4M3 ``[n, k]`` codes and UE8M0 block exponents ``uint8 [ceil(n/128), k/128]``.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, device="cuda", generator=gen) * 0.5).to(torch.float8_e4m3fn) + exps = torch.randint(119, 125, (-(-n // 128), k // 128), device="cuda", generator=gen) + return w, exps.to(torch.uint8) + + +def _rows(m, k, seed): + gen = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(m, k, device="cuda", generator=gen) + mag = torch.randint(-6, 3, (m, k // 32), device="cuda", generator=gen).float() + return (x * torch.exp2(mag).repeat_interleave(32, dim=1)).to(torch.bfloat16) + + +def _bf16(m, n, seed, scale=1.0): + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(m, n, device="cuda", generator=gen) * scale).to(torch.bfloat16) + + +def _run_rank(entry_name, world, mv2_tests): + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + import importlib + + sys.path.insert(0, mv2_tests) + from _validating import assert_within_ulp, validating + + from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy + from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace + from tensorrt_llm.mapping import Mapping + + lowm = importlib.import_module(f"{K3}.kernels.k3_fp8w_lowm_gemm") + mx = importlib.import_module(f"{K3}.kernels.k3_fp8w_lowm_mxfp8") + lamport = importlib.import_module(f"{K3}.kernels.k3_lamport_ar") + op = getattr(importlib.import_module(f"{K3}.ops.gemm.{entry_name}"), entry_name) + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + torch.cuda.set_device(rank) + sm107 = torch.cuda.get_device_capability() == (10, 7) + report = {"failed": [], "skipped": [], "passed": 0} + try: + mapping = Mapping(world_size=world, tp_size=world, rank=rank) + allreduce = AllReduce( + mapping=mapping, strategy=AllReduceStrategy.MNNVL, dtype=torch.bfloat16 + ) + assert allreduce.mnnvl_allreduce is not None + workspace = get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16, BUFFER_BYTES) + ws = lamport.LamportWorkspace.from_mnnvl(workspace, world, rank) + + def finish(m, n, one_shot, ag=False): + out = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + lamport.run_lamport_finish(ws, out, one_shot, ag=ag) + torch.cuda.synchronize() + return out + + def lowm_push(i, s): + n, k, m, one_shot = s["n"], s["k"], s["m"], s["one_shot"] + strided = s.get("strided", False) + k_total = s.get("k_total", k) + xs, ws_, scales, gates, addends = [], [], [], [], [] + for r in range(world): + x_full = _rows(m, k_total, 100 * i + r) + w_full, exps = _weight(n, k_total, 200 * i + r) + lo = r * k % k_total if strided else 0 + xs.append(x_full[:, lo : lo + k]) + ws_.append(w_full[:, lo : lo + k]) + scales.append( + lowm.pack_ue8m0_exponents( + exps[:, lo // 128 : (lo + k) // 128].contiguous(), n, k + ) + ) + gates.append(_bf16(m, k, 300 * i + r).sigmoid() if s.get("gate") else None) + addends.append(_bf16(m, n, 400 * i + r) if s.get("addend") else None) + kwargs = {"strided": strided} + if s.get("gate"): + kwargs["gate"] = gates[rank] + if s.get("addend"): + kwargs["addend"] = addends[rank] + with validating(op): + op(xs[rank], ws_[rank], scales[rank], ws, one_shot, **kwargs) + got = finish(m, n, one_shot) + if s.get("gate"): + kwargs["gate"] = gates + if s.get("addend"): + kwargs["addend"] = addends + return got, op.reference(xs, ws_, scales, ws, one_shot, **kwargs) + + def mid_push(i, s): + n, k, m, one_shot = s["n"], s["k"], s["m"], s["one_shot"] + tactic = lowm.mid_ar_push_tactic(m, n, k) + assert tactic is not None, "a mid-M push cell must be above the low-M cut-over" + xqs, xss, ws_, scales = [], [], [], [] + for r in range(world): + xq, xs = lowm.quantize_mid_operand(_rows(m, k, 500 * i + r), tactic) + w, exps = _weight(n, k, 600 * i + r) + xqs.append(xq) + xss.append(xs) + ws_.append(w) + scales.append(lowm.pack_ue8m0_exponents(exps, n, k)) + with validating(op): + op(xqs[rank], xss[rank], ws_[rank], scales[rank], ws, one_shot, tactic) + got = finish(m, n, one_shot) + return got, op.reference(xqs, xss, ws_, scales, ws, one_shot, tactic) + + def mxfp8_push(i, s): + variant, k, m = s["variant"], s["k"], s["m"] + if variant == "ag": + hidden = s["hidden"] + n = hidden // world + # The column-parallel tail: the same latent rows on every rank. + xq, xs = mx.mxfp8_quantize_reference(_rows(m, k, 700 * i)) + w_full, exps = _weight(hidden, k, 800 * i) + col0s = [r * n for r in range(world)] + ws_ = [w_full[c : c + n].contiguous() for c in col0s] + scales = [ + lowm.pack_ue8m0_exponents(exps[c // 128 : (c + n) // 128].contiguous(), n, k) + for c in col0s + ] + tactic = mx.default_mx_tactic(m, hidden, k) + zeros = torch.zeros(m, hidden, dtype=torch.bfloat16, device="cuda") + lamport.run_lamport_push(ws, zeros, True) + with validating(op): + op( + xq, + xs, + ws_[rank], + scales[rank], + ws, + col0=col0s[rank], + hidden=hidden, + tactic=tactic, + ) + got = finish(m, hidden, True, ag=True) + want = op.reference( + [xq] * world, + [xs] * world, + ws_, + scales, + ws, + col0=col0s, + hidden=hidden, + tactic=tactic, + ) + return got, want + n = s["n"] + one_shot = s.get("one_shot", True) + k_total = s.get("k_total", k) + xqs, xss, ws_, scales, addends, k0s = [], [], [], [], [], [] + for r in range(world): + xq, xs = mx.mxfp8_quantize_reference(_rows(m, k_total, 900 * i + r)) + w, exps = _weight(n, k_total, 1000 * i + r) + lo = r * k % k_total + xqs.append(xq) + xss.append(xs) + ws_.append(w) + k0s.append(lo) + if variant == "slice": + sl = exps[:, lo // 128 : (lo + k) // 128].contiguous() + scales.append(lowm.pack_ue8m0_exponents(sl, n, k)) + addends.append(_bf16(m, n, 1100 * i + r)) + else: + scales.append(lowm.pack_ue8m0_exponents(exps, n, k)) + if variant == "slice": + with validating(op): + op( + xqs[rank], + xss[rank], + ws_[rank], + scales[rank], + ws, + one_shot, + k0=k0s[rank], + k1=k, + addend=addends[rank], + ) + got = finish(m, n, one_shot) + want = op.reference( + xqs, xss, ws_, scales, ws, one_shot, k0=k0s, k1=k, addend=addends + ) + return got, want + with validating(op): + op(xqs[rank], xss[rank], ws_[rank], scales[rank], ws, one_shot) + got = finish(m, n, one_shot) + return got, op.reference(xqs, xss, ws_, scales, ws, one_shot) + + run = { + "k3_fp8w_lowm_gemm_push": lowm_push, + "k3_fp8w_mid_gemm_push": mid_push, + "k3_fp8w_lowm_mxfp8_push": mxfp8_push, + }[entry_name] + for i, cell in enumerate(op.CELLS): + if cell.spec.get("sm_107") and not sm107: + report["skipped"].append(cell.why) + continue + comm.Barrier() + got, want = run(i, cell.spec) + try: + assert_within_ulp(got, want, ELEMENT_ULP, RMS_ULP) + report["passed"] += 1 + except AssertionError as e: + report["failed"].append(f"{cell.why}: {e}") + except Exception: + traceback.print_exc() + raise + return report + + +@skip_pre_blackwell +@pytest.mark.parametrize("entry_name", ENTRIES) +@pytest.mark.parametrize("mpi_pool_executor", [WORLD], indirect=True) +def test_k3_fp8w_gemm_push_cells(mpi_pool_executor, entry_name): + """Drives ``k3_fp8w_lowm_gemm_push.CELLS``, ``k3_fp8w_mid_gemm_push.CELLS`` or + ``k3_fp8w_lowm_mxfp8_push.CELLS`` on every rank.""" + world = mpi_pool_executor.num_workers + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + args = [(entry_name, world, MV2_TESTS)] * world + reports = list(mpi_pool_executor.map(_run_rank, *zip(*args))) + failed = [f"rank {r}: {msg}" for r, rep in enumerate(reports) for msg in rep["failed"]] + assert not failed, "\n".join(failed) + assert all(rep["passed"] for rep in reports), reports + skipped = reports[0]["skipped"] + if skipped: + print(f"{entry_name}: skipped off SM107: {skipped}") diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/test_k3_moe_entry_cells.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/test_k3_moe_entry_cells.py new file mode 100644 index 000000000000..9aef9e3bf530 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/moe/test_k3_moe_entry_cells.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The single-GPU MoE entries of the K3 default target over their certified cells: +``k3_router_logits.CELLS``, ``k3_route_tail.CELLS`` and ``k3_moe_residual_add.CELLS``, each driven +through ``validating`` and checked against the entry's own ``reference``. +(``k3_small_batch_moe.CELLS`` run in ``kernels/test_k3_small_batch_moe.py``.) +""" + +import importlib + +import pytest +import torch + +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../../.."] +from _validating import ULP, assert_within_ulp, validating # noqa: E402 + +K3 = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() not in (100, 103, 107), + reason="the K3 routing kernels run on SM100 / SM103 / SM107", +) + +route_tail = importlib.import_module(f"{K3}.kernels.k3_route_tail") +k3_router_logits = importlib.import_module(f"{K3}.ops.moe.k3_router_logits").k3_router_logits +k3_route_tail = importlib.import_module(f"{K3}.ops.moe.k3_route_tail").k3_route_tail +k3_moe_residual_add = importlib.import_module( + f"{K3}.ops.moe.k3_moe_residual_add" +).k3_moe_residual_add + +HIDDEN, NUM_EXPERTS, TOP_K = 7168, 896, 16 +ROUTED_SCALING_FACTOR = 2.5 + + +def _ids(cells): + return [c.why[:48] for c in cells] + + +@pytest.mark.parametrize("cell", k3_router_logits.CELLS, ids=_ids(k3_router_logits.CELLS)) +def test_k3_router_logits_cells(cell): + """FP32 logits of BF16 operands: the products are exact in FP32 and only the summation order + differs from the reference's (14 slices of 512, then the slices in order, against cuBLAS's + FP32 order). A random-walk estimate of that difference over K = 7168 terms of these operands + is ~1.6e-6 of the row's largest logit, about 13 FP32 ulps (2^-23); the gates allow 64 per + element and 16 in RMS. The note's bitwise claim against ``dsv3_router_gemm_op`` is checked + too.""" + m = cell.spec["m"] + gen = torch.Generator(device="cuda").manual_seed(m) + x = torch.randn(m, HIDDEN, device="cuda", generator=gen).to(torch.bfloat16) + w = (torch.randn(NUM_EXPERTS, HIDDEN, device="cuda", generator=gen) * 0.02).to(torch.bfloat16) + with validating(k3_router_logits): + got = k3_router_logits(x, w) + want = k3_router_logits.reference(x, w) + assert_within_ulp(got, want, 64, 16, ULP[torch.float32]) + stock = torch.ops.trtllm.dsv3_router_gemm_op(x, w.t(), bias=None, out_dtype=torch.float32) + assert torch.equal(got, stock), "not bitwise dsv3_router_gemm_op" + + +@pytest.mark.parametrize("cell", k3_moe_residual_add.CELLS, ids=_ids(k3_moe_residual_add.CELLS)) +def test_k3_moe_residual_add_cells(cell): + """Two BF16 adds, each rounded once: bitwise the reference.""" + m = cell.spec["m"] + gen = torch.Generator(device="cuda").manual_seed(m) + routed, shared, residual = ( + torch.randn(m, HIDDEN, device="cuda", generator=gen).to(torch.bfloat16) for _ in range(3) + ) + with validating(k3_moe_residual_add): + got = k3_moe_residual_add(routed, shared, residual) + want = k3_moe_residual_add.reference(routed, shared, residual) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)) + + +@pytest.mark.parametrize("cell", k3_route_tail.CELLS, ids=_ids(k3_route_tail.CELLS)) +def test_k3_route_tail_cells(cell): + """The top-k ids (selection and order) exactly; the weights within 4 FP32 ulps of the row's + largest (both renormalize in FP64 from the same FP32 scores; the kernel's sigmoid is its own); + of the ``moe_sort`` tables every written entry exactly (an expert's rows in ascending token + order, so in ascending expanded index, as the reference sorts them).""" + s = cell.spec + m, offset, local, splits = s["m"], s["offset"], s["local"], s["splits"] + gen = torch.Generator(device="cuda").manual_seed(1000 + m) + shape = (splits, m, NUM_EXPERTS) if splits else (m, NUM_EXPERTS) + logits = torch.randn(*shape, device="cuda", generator=gen) + if splits: + logits = logits / splits + bias = 0.1 * torch.randn(NUM_EXPERTS, device="cuda", generator=gen) + workspace = route_tail.wide_workspace(logits.device) if m > route_tail.MAX_M else None + args = (logits, bias, ROUTED_SCALING_FACTOR, offset, local, route_tail.TILE) + with validating(k3_route_tail): + got = k3_route_tail(*args, workspace=workspace) + torch.cuda.synchronize() + want = k3_route_tail.reference(*args, workspace=workspace) + ids, weights, tile_expert, tile_limit, e2p, p2e, total, num_tiles = got + w_ids, w_weights, w_tile_expert, w_tile_limit, w_e2p, w_p2e, w_total, w_num_tiles = want + + assert torch.equal(ids, w_ids), "top-k ids" + assert_within_ulp(weights, w_weights, 4, 1, ULP[torch.float32]) + tiles = int(w_num_tiles.item()) + assert int(num_tiles.item()) == tiles and int(total.item()) == int(w_total.item()) + assert torch.equal(tile_expert[:tiles], w_tile_expert[:tiles]) + assert torch.equal(tile_limit[:tiles], w_tile_limit[:tiles]) + local_rows = w_e2p.reshape(-1) >= 0 + assert bool(local_rows.any()), "the cell routes no row to the window" + assert torch.equal(e2p.reshape(-1)[local_rows], w_e2p.reshape(-1)[local_rows]) + written = w_p2e >= 0 + assert torch.equal(p2e[: written.numel()][written], w_p2e[written]) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/test_k3_attn_res_norm.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/test_k3_attn_res_norm.py new file mode 100644 index 000000000000..bd751ed82583 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/norm/test_k3_attn_res_norm.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the k3_attn_res_norm entry (``kernels/k3_attn_res.py``, one GPU). + +Every cell of ``k3_attn_res_norm.CELLS`` runs under ``validating``: the updated prefix sum bit for bit +against the reference's FP32 adds rounded to BF16, the normed output against the reference (the family +``attn_res_fwd`` mixture and the twice-rounded trailing norm) under the cosine / relative-L2 metric of +the family attn_res tests, identical bits on a repeated call, untouched inputs, and the MXFP8 output +bit for bit the quantization rule of the low-M FP8 GEMM applied to the stored BF16 output (row-major +scales, or the mid-M GEMM's scale atoms). +""" + +import importlib +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch + +__extra_import_path__ = ["../../../../.."] +from _validating import validating # noqa: E402 + +_ENTRY = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + "kimi_k3_nvfp4__sm_107__default.ops.norm.k3_attn_res_norm" +) +k3_attn_res_norm = _ENTRY.k3_attn_res_norm + +_ARCH = {(10, 0): Arch.SM_100, (10, 3): Arch.SM_103, (10, 7): Arch.SM_107} +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or _ARCH.get(torch.cuda.get_device_capability()) not in k3_attn_res_norm.ARCHS, + reason="k3_attn_res_norm is certified on its ARCHS", +) + +# The thresholds of the family attn_res tests (main's op test metric). +MIN_COSINE = 0.9999 +MAX_RELATIVE_L2 = 5e-3 + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.uint8) + + +def _similarity(actual: torch.Tensor, expected: torch.Tensor) -> tuple[float, float]: + a, e = actual.float().flatten(), expected.float().flatten() + cosine = torch.nn.functional.cosine_similarity(a, e, dim=0).item() + return cosine, ((a - e).norm() / (e.norm() + 1e-12)).item() + + +def _inputs(spec: dict, seed: int): + gen = torch.Generator(device="cuda").manual_seed(seed) + + def rnd(*shape, scale=1.0, shift=0.0): + x = torch.randn(*shape, generator=gen, device="cuda") * scale + shift + return x.to(torch.bfloat16) + + m, n, h = spec["rows"], spec["candidates"], spec["hidden"] + layer = rnd(m, h, scale=0.5) + addends = tuple(rnd(m, h, scale=0.25) for _ in range(spec["addends"])) + block = rnd(n - 1, m, h, scale=0.5) + weights = (rnd(h, scale=0.02), rnd(h, scale=0.02, shift=1.0), rnd(h, scale=0.02, shift=1.0)) + return layer, addends, block, weights + + +@pytest.mark.parametrize( + "cell", k3_attn_res_norm.CELLS, ids=[c.why[:48] for c in k3_attn_res_norm.CELLS] +) +def test_certified_cells(cell) -> None: + spec = cell.spec + layer, addends, block, weights = _inputs(spec, zlib.crc32(cell.why.encode()) % 2**31) + flags = dict(mxfp8=spec["mxfp8"], entry_trigger=spec["entry_trigger"], mx_atom=spec["mx_atom"]) + args = (layer, addends, block, *weights, spec["rms_eps"], spec["out_eps"]) + before = [t.clone() for t in (layer, *addends, block, *weights)] + with validating(k3_attn_res_norm): + got = k3_attn_res_norm(*args, **flags) + again = k3_attn_res_norm(*args, **flags) + ref = k3_attn_res_norm.reference(*args, **flags) + if not addends: + got, again, ref = (got,), (again,), (ref,) + assert len(got) == len(ref), cell.why + atom_offsets = ( + _ENTRY.mx_atom_offsets(layer.shape[0], layer.shape[1], layer.device) + if spec["mxfp8"] and spec["mx_atom"] + else None + ) + for i, (a, b) in enumerate(zip(got, again)): + if atom_offsets is not None and i == 3: + # Atom rows past M are not stored: compare the scales the kernel writes. + a, b = a[atom_offsets], b[atom_offsets] + assert torch.equal(_bits(a), _bits(b)), f"{cell.why}: a repeated call changed bits" + index = 1 if addends else 0 + output, ref_output = got[index], ref[index] + assert output.shape == layer.shape and output.dtype == torch.bfloat16, cell.why + cosine, rel_l2 = _similarity(output, ref_output) + assert cosine > MIN_COSINE and rel_l2 < MAX_RELATIVE_L2, ( + f"{cell.why}: cosine {cosine:.6f}, relative L2 {rel_l2:.3e}" + ) + if addends: + assert torch.equal(_bits(got[0]), _bits(ref[0])), f"{cell.why}: updated prefix differs" + if spec["mxfp8"]: + xq, xs = got[2], got[3] + want_q, want_s = _ENTRY.mxfp8_rows_reference(output) + assert torch.equal(_bits(xq), _bits(want_q)), f"{cell.why}: E4M3 codes differ" + if spec["mx_atom"]: + offsets = _ENTRY.mx_atom_offsets(layer.shape[0], layer.shape[1], layer.device) + assert tuple(xs.shape) == tuple(ref[3].shape), cell.why + assert torch.equal(xs[offsets], want_s.reshape(-1)), f"{cell.why}: atom scales differ" + assert torch.equal( + ref[3][offsets], _ENTRY.mxfp8_rows_reference(ref_output)[1].reshape(-1) + ) + else: + assert torch.equal(xs, want_s), f"{cell.why}: scales differ" + for old, new in zip(before, (layer, *addends, block, *weights)): + assert torch.equal(_bits(old), _bits(new)), f"{cell.why}: an input was written" + + +def test_is_valid_refuses_strided_rows() -> None: + """A row-strided layer is read as dense rows by the kernel; the guard refuses it.""" + spec = k3_attn_res_norm.CELLS[0].spec + layer, addends, block, weights = _inputs(spec, 7) + wide = torch.cat([layer, layer], dim=1)[:, : layer.shape[1]] + with pytest.raises(AssertionError), validating(k3_attn_res_norm): + k3_attn_res_norm(wide, addends, block, *weights, 1e-6, 1e-6) + + +def test_mxfp8_without_moe_addends_raises() -> None: + spec = k3_attn_res_norm.CELLS[3].spec + layer, addends, block, weights = _inputs(spec, 8) + with pytest.raises(ValueError): + k3_attn_res_norm(layer, addends, block, *weights, 1e-6, 1e-6, mxfp8=True) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_markov_entries.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_markov_entries.py new file mode 100644 index 000000000000..21a182be2a52 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_markov_entries.py @@ -0,0 +1,321 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark greedy Markov chain entries (``ops/spec``) over their certified cells. + +The one-launch chains run as one rank of the helix group on one GPU: the other ranks' pairs +(``k3_markov_chain``) or per-CTA keys (``k3_dspark_markov_greedy``) of every block position are +written into the exchange before the launch, computed from the reference chain, so the kernel +polls the group's exchange as it does on the deployment and reduces it with its own rank's. +""" + +import zlib + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("needs a CUDA GPU", allow_module_level=True) + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec import ( # noqa: E402, E501 + _reference, + k3_dspark_markov_greedy, + k3_fused_markov_chain_logits, + k3_markov_chain, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +VOCAB = 163840 +MARKOV_RANK = 256 +_ARCH = "sm_{}{}".format(*torch.cuda.get_device_capability()) +_U32 = 0xFFFFFFFF + + +def _require_arch(op) -> None: + if _ARCH not in {a.value for a in op.ARCHS}: + pytest.skip(f"{type(op).__name__} is not certified on {_ARCH}") + + +def _gen(cell) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + + +def _weights(gen) -> tuple[torch.Tensor, torch.Tensor]: + """BF16 ``markov_w1`` / ``markov_w2`` ``[vocab, R]`` with entries of 0.25: a bias of unit + scale, of the order of the logits' spread, so that it moves the argmax.""" + w = 0.25 * torch.randn(2, VOCAB, MARKOV_RANK, device="cuda", generator=gen) + return w[0].to(torch.bfloat16), w[1].to(torch.bfloat16) + + +def _anchors(reqs: int, dtype: str, gen, out_of_vocab: bool) -> torch.Tensor: + anchors = torch.randint(0, VOCAB, (reqs,), device="cuda", generator=gen) + if out_of_vocab: + anchors[reqs - 1] = -1 + return anchors.to(getattr(torch, dtype)) + + +def _dot_slack(row: torch.Tensor, w2: torch.Tensor) -> torch.Tensor: + """How far an FP32 dot over the Markov rank may sit from the exact one, in any summation + order: ``R * 2**-24 * sum |w1 w2|`` (the standard bound for a length-R sum), doubled.""" + return (MARKOV_RANK * 2.0**-23 * (row.abs() @ w2.abs().t())).float() + + +def _bias_tol(bias: torch.Tensor, slack: torch.Tensor) -> torch.Tensor: + """How far an op's BF16 bias may sit from the reference's: the dot's slack, then one BF16 + rounding (an ulp is at most 2**-7 of the value) on either side.""" + return 2.0**-7 * bias.abs() + slack + + +def _chain_scores(block, anchors, markov_w1, markov_w2, tokens): + """``(scores, bias, slack)`` FP32 ``[B, K, vocab]`` of the reference chain that produced + ``tokens``: position ``i``'s logits plus the BF16 bias of ``prev_i``, and the slack of the + FP32 dot behind each bias (``_dot_slack``).""" + w2 = markov_w2.double() + prev = anchors.to(torch.int64) + scores, biases, slacks = [], [], [] + for i in range(block.shape[1]): + valid = (prev >= 0) & (prev < VOCAB) + row = markov_w1[prev.clamp(0, VOCAB - 1)].double() * valid[:, None] + bias = (row @ w2.t()).to(torch.bfloat16).float() + scores.append(block[:, i].float() + bias) + biases.append(bias) + slacks.append(_dot_slack(row, w2)) + prev = tokens[:, i].to(torch.int64) + return torch.stack(scores, 1), torch.stack(biases, 1), torch.stack(slacks, 1) + + +def _assert_chain(tokens, ref_tokens, scores, bias, slack) -> None: + """Tokens equal up to a near-tie. The kernels round an FP32 dot over R to the BF16 bias, the + reference an FP64 product: the two biases differ by up to ``_bias_tol``. Where a request's + chain first differs, the kernel's token must score within the two candidates' tolerances + of the reference's; the chain is not compared past that position.""" + assert tokens.dtype == torch.int32 and tokens.shape == ref_tokens.shape + diverged = 0 + for b in range(tokens.shape[0]): + diff = (tokens[b] != ref_tokens[b]).nonzero() + if diff.numel() == 0: + continue + i = int(diff[0]) + k, r = int(tokens[b, i]), int(ref_tokens[b, i]) + assert 0 <= k < VOCAB, f"request {b} position {i}: token {k} outside the vocab" + tol = float(_bias_tol(bias[b, i, [k, r]], slack[b, i, [k, r]]).sum()) + gap = float(scores[b, i, r] - scores[b, i, k]) + assert gap <= tol, f"request {b} position {i}: {k} vs {r}, gap {gap:.3e} > {tol:.3e}" + diverged += 1 + assert diverged <= 1, f"{diverged} requests' chains differ" + + +def _block_logits(reqs: int, block: int, gen) -> torch.Tensor: + return (2.0 * torch.randn(reqs, block, VOCAB, device="cuda", generator=gen)).to(torch.bfloat16) + + +# --------------------------------------------------------------------------------------------- +# k3_markov_chain (Triton, one launch) +# --------------------------------------------------------------------------------------------- + +_CHAIN = k3_markov_chain.k3_markov_chain + + +@pytest.mark.parametrize("cell", _CHAIN.CELLS, ids=[c.why[:48] for c in _CHAIN.CELLS]) +def test_markov_chain_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.spec.vocab_parallel_kernels import ( # noqa: E501 + _CHAIN_MAX_REQS, + _CHAIN_MAX_STEPS, + _MARKOV_BLOCK_V, + ) + + _require_arch(_CHAIN) + spec, gen = cell.spec, _gen(cell) + tp, rank, reqs, block = spec["tp"], spec["rank"], spec["reqs"], spec["block"] + shard = VOCAB // tp + programs = -(-shard // _MARKOV_BLOCK_V) + if programs > torch.cuda.get_device_properties(0).multi_processor_count: + pytest.skip(f"{programs} programs must all be resident") + certifying(_CHAIN) + markov_w1, markov_w2 = _weights(gen) + full = _block_logits(reqs, block, gen) + anchors = _anchors(reqs, spec["anchor"], gen, out_of_vocab=reqs > 1) + blocks = [full[..., s * shard : (s + 1) * shard].contiguous() for s in range(tp)] + w2_shards = [markov_w2[s * shard : (s + 1) * shard] for s in range(tp)] + tokens = torch.full((reqs, block), -7, dtype=torch.int32, device="cuda") + set_floats = _CHAIN_MAX_STEPS * _CHAIN_MAX_REQS * tp * 2 + partial = torch.empty( + (_CHAIN_MAX_STEPS * programs * _CHAIN_MAX_REQS * 2,), dtype=torch.float32, device="cuda" + ) + tickets = torch.zeros((_CHAIN_MAX_STEPS,), dtype=torch.int32, device="cuda") + launch = torch.zeros((1,), dtype=torch.int32, device="cuda") + exchange = torch.full((3 * set_floats,), -0.0, dtype=torch.float32, device="cuda") + scratch = (partial, tickets, launch, exchange, exchange) + ref = _CHAIN.reference( + blocks, anchors, markov_w1, w2_shards, tokens, *scratch, set_floats, VOCAB, 0, tp, rank + ) + scores, bias, slack = _chain_scores(full, anchors, markov_w1, markov_w2, ref) + + # The other ranks' pair of every position, in set 0 (launch 0): [step][request][rank][2]. + cols = torch.arange(VOCAB, device="cuda").view(tp, shard).expand(reqs, block, tp, shard) + value, index = _reference.argmax_pairs(scores.view(reqs, block, tp, shard), cols) + value = torch.where(value == 0, 0.0, value) # -0.0 marks an empty slot + region = exchange[: block * _CHAIN_MAX_REQS * tp * 2].view(block, _CHAIN_MAX_REQS, tp, 2) + peers = [s for s in range(tp) if s != rank] + region[:, :reqs, peers, 0] = value.transpose(0, 1)[:, :, peers] + region[:, :reqs, peers, 1] = index.transpose(0, 1)[:, :, peers].float() + + out = _CHAIN( + blocks[rank], + anchors, + markov_w1, + w2_shards[rank], + tokens, + *scratch, + set_floats=set_floats, + vocab=VOCAB, + col_base=rank * shard, + tp=tp, + tp_rank=rank, + ) + torch.cuda.synchronize() + assert out is tokens and int(launch) == 1 + _assert_chain(tokens, ref, scores, bias, slack) + + +# --------------------------------------------------------------------------------------------- +# k3_dspark_markov_greedy (CuTe DSL, one launch) +# --------------------------------------------------------------------------------------------- + +_GREEDY = k3_dspark_markov_greedy.k3_dspark_markov_greedy + + +def _u32_to_i32(x: torch.Tensor) -> torch.Tensor: + return torch.where(x >= 2**31, x - 2**32, x).to(torch.int32) + + +def _keys(value: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """``(low, high)`` int32 words of the kernel's key of FP32 ``value`` at global ``index``: + high the value's bits in unsigned order (-0 as +0, every NaN on top), low ``~index``.""" + b = value.contiguous().view(torch.int32).to(torch.int64) & _U32 + b = torch.where(b == 0x80000000, 0, b) + high = torch.where(b >= 0x80000000, ~b & _U32, b | 0x80000000) + high = torch.where((b & 0x7FFFFFFF) > 0x7F800000, _U32, high) + low = ~index.to(torch.int64) & _U32 + return torch.stack([_u32_to_i32(low), _u32_to_i32(high)], dim=-1) + + +@pytest.mark.parametrize("cell", _GREEDY.CELLS, ids=[c.why[:48] for c in _GREEDY.CELLS]) +def test_dspark_markov_greedy_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.dspark_markov_greedy import ( # noqa: E501 + NUM_SETS, + ROWS, + STATE_WORDS, + ctas_for, + set_bytes, + width_for, + ) + + _require_arch(_GREEDY) + certifying(_GREEDY) + spec, gen = cell.spec, _gen(cell) + tp, rank, reqs, block = spec["tp"], spec["rank"], spec["reqs"], spec["block"] + shard = VOCAB // tp + ctas, width = ctas_for(shard), width_for(reqs) + markov_w1, markov_w2 = _weights(gen) + full = _block_logits(reqs, block, gen) + anchors = _anchors(reqs, spec["anchor"], gen, out_of_vocab=reqs > 1) + blocks = [full[..., s * shard : (s + 1) * shard].contiguous() for s in range(tp)] + w2_shards = [markov_w2[s * shard : (s + 1) * shard] for s in range(tp)] + tokens = torch.full((reqs, block), -7, dtype=torch.int32, device="cuda") + state = torch.zeros((STATE_WORDS,), dtype=torch.int32, device="cuda") + box = torch.zeros((NUM_SETS * set_bytes(tp * ctas) // 4,), dtype=torch.int32, device="cuda") + kwargs = dict( + vocab=VOCAB, + tp=tp, + rank=rank, + state=state, + box_uc=box.data_ptr(), + box_mc=box.data_ptr(), + multicast=False, + ) + ref = _GREEDY.reference(blocks, anchors, markov_w1, w2_shards, tokens, **kwargs) + scores, bias, slack = _chain_scores(full, anchors, markov_w1, markov_w2, ref) + + # Every other rank's CTA keys of every position in set 0 (launch 0): + # [position][rank][CTA][request slot][low, high]; padding slots hold a key of 0.0. + cols = torch.arange(VOCAB, device="cuda").view(tp, ctas, ROWS) + value, index = _reference.argmax_pairs( + scores.view(reqs, block, tp, ctas, ROWS), cols.expand(reqs, block, tp, ctas, ROWS) + ) + words = torch.empty((block, tp, ctas, width, 2), dtype=torch.int32, device="cuda") + words[:, :, :, :reqs] = _keys(value.float(), index).permute(1, 2, 3, 0, 4) + pad_index = cols[..., 0].expand(block, tp, ctas).unsqueeze(-1).expand(-1, -1, -1, width - reqs) + words[:, :, :, reqs:] = _keys(torch.zeros(pad_index.shape, device="cuda"), pad_index) + region = box[: block * tp * ctas * width * 2].view(block, tp, ctas, width, 2) + peers = [s for s in range(tp) if s != rank] + region[:, peers] = words[:, peers] + + _GREEDY(blocks[rank], anchors, markov_w1, w2_shards[rank], tokens, **kwargs) + torch.cuda.synchronize() + assert int(state[0]) == 1, "the launch ordinal advances once" + _assert_chain(tokens, ref, scores, bias, slack) + + +# --------------------------------------------------------------------------------------------- +# k3_fused_markov_chain_logits (full vocab) +# --------------------------------------------------------------------------------------------- + +_FUSED = k3_fused_markov_chain_logits.k3_fused_markov_chain_logits + + +@pytest.mark.parametrize("cell", _FUSED.CELLS, ids=[c.why[:48] for c in _FUSED.CELLS]) +def test_fused_markov_chain_logits_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + from tensorrt_llm._torch.models.modeling_speculative import ( + dspark_markov_chain, + dspark_markov_step_bias, + ) + + _require_arch(_FUSED) + certifying(_FUSED) + spec, gen = cell.spec, _gen(cell) + reqs, block = spec["reqs"], spec["block"] + markov_w1, markov_w2 = _weights(gen) + base = _block_logits(reqs, block, gen).float() + anchors = _anchors(reqs, spec["anchor"], gen, out_of_vocab=spec["anchor"] == "int32") + base_before = base.clone() + out = _FUSED(base, anchors, markov_w1, markov_w2) + assert torch.equal(base, base_before), "the op must not write the logits" + assert out.shape == base.shape and out.dtype == torch.float32 + + # Cross-check: the chain of the drafter's own head (the same F.linear bias), bit for bit. + _, chained = dspark_markov_chain( + base, + anchors, + lambda prev, _hidden: dspark_markov_step_bias(prev, markov_w1, markov_w2), + cast_bias_to_logits=True, + ) + assert torch.equal(out, chained) + + # The reference rounds an FP64 bias to BF16, the op an FP32 GEMV: the biases differ by up + # to ``_bias_tol``, and both FP32 sums with the logits round once (2**-24 of the result + # each, and the bias read back from the reference's sum once more). The positions are + # compared until the two chains first pick different tokens, which must be a near-tie + # within the candidates' tolerances (every later position has another previous token). + ref = _FUSED.reference(base, anchors, markov_w1, markov_w2) + ref_bias = ref - base + cols = torch.arange(VOCAB, device="cuda").expand(reqs, -1) + w2 = markov_w2.double() + prev = anchors.to(torch.int64) + for i in range(block): + valid = (prev >= 0) & (prev < VOCAB) + row = markov_w1[prev.clamp(0, VOCAB - 1)].double() * valid[:, None] + tol = _bias_tol(ref_bias[:, i], _dot_slack(row, w2)) + band = tol + 2.0**-22 * ref[:, i].abs() + assert bool(((out[:, i] - ref[:, i]).abs() <= band).all()), f"position {i}" + _, k_tok = _reference.argmax_pairs(out[:, i], cols) + _, r_tok = _reference.argmax_pairs(ref[:, i], cols) + if not torch.equal(k_tok, r_tok): + b = int((k_tok != r_tok).nonzero()[0]) + k, r = int(k_tok[b]), int(r_tok[b]) + gap = float(ref[b, i, r] - ref[b, i, k]) + limit = float(tol[b, k] + tol[b, r]) + assert gap <= limit, f"position {i}: {k} vs {r}, gap {gap:.3e} > {limit:.3e}" + break + prev = r_tok diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_spec_tail_entries.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_spec_tail_entries.py new file mode 100644 index 000000000000..451f70356f47 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_spec_tail_entries.py @@ -0,0 +1,574 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The K3 speculative-tail entries (``ops/spec``) against their references, over their cells. + +Every entry here is integer bookkeeping, a copy, or a masked store whose arithmetic is one +multiply by 0 or 1 and an optional FP8 conversion, so each written buffer must match its +reference bit for bit, including the rows and slots the op must leave alone. The selection +entry is also checked against ``KDAReplayState.record_acceptance`` for the replay record. +""" + +import importlib +import zlib + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("the speculative-tail kernels need a GPU", allow_module_level=True) + +__extra_import_path__ = ["../../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import KDAReplayState # noqa: E402 +from tensorrt_llm._torch.speculative.interface import ( # noqa: E402 + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SEED, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +_OPS = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec" + + +def _entry(name): + return getattr(importlib.import_module(f"{_OPS}.{name}"), name) + + +k3_draft_block_table = _entry("k3_draft_block_table") +k3_fused_strict_accept = _entry("k3_fused_strict_accept") +k3_fused_select_accept = _entry("k3_fused_select_accept") +k3_next_new_tokens = _entry("k3_next_new_tokens") +k3_dflash_ctx_positions = _entry("k3_dflash_ctx_positions") +k3_dflash_gen_prepare = _entry("k3_dflash_gen_prepare") +k3_dflash_paged_latent_prepare = _entry("k3_dflash_paged_latent_prepare") +k3_dflash_noise_embedding = _entry("k3_dflash_noise_embedding") +k3_store_context_latents = _entry("k3_store_context_latents") +k3_scatter_rows_to_pool = _entry("k3_scatter_rows_to_pool") +k3_vocab_shard_embedding = _entry("k3_vocab_shard_embedding") + +_ARCH = "sm_{}{}".format(*torch.cuda.get_device_capability()) +if not any( + _ARCH in {a.value for a in type(op).ARCHS} + for op in (k3_draft_block_table, k3_fused_select_accept, k3_dflash_gen_prepare) +): + pytest.skip(f"no speculative-tail entry is certified on {_ARCH}", allow_module_level=True) + +VOCAB = 163840 +HIDDEN = 7168 +LATENT = 576 # kv_lora 512 + rope 64 +LAYERS = 5 +PAGE = 128 +MAX_CTX = 997696 +EMPTY_INDEX = 2**24 + + +def _ids(op): + return [c.why[:44] for c in op.CELLS] + + +def _gen(cell) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + + +def _require(op) -> None: + if _ARCH not in {a.value for a in type(op).ARCHS}: + pytest.skip(f"{type(op).__name__} is not certified on {_ARCH}") + + +def _randint(low, high, shape, gen, dtype=torch.int64): + return torch.randint(low, high, shape, generator=gen, device="cuda", dtype=dtype) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + """The tensor's raw bits, so that -0.0 and +0.0 differ.""" + width = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}[t.element_size()] + return t.contiguous().view(width) + + +def _same(out: torch.Tensor, ref: torch.Tensor, what: str) -> None: + assert out.dtype == ref.dtype, (what, out.dtype, ref.dtype) + assert out.shape == ref.shape, (what, tuple(out.shape), tuple(ref.shape)) + assert torch.equal(_bits(out), _bits(ref)), f"{what} differs from the reference" + + +def _force(cell, draft_len: int, gen) -> tuple: + """The ``force`` operand ``K3DFlashWorker._force_accept_kernel_args`` builds, and the counter.""" + value = float(cell.spec["force"]) + if value == 0.0: + return None, None + int_part = int(value) + max_total = draft_len + 1 + base_total = min(int_part + 1, max_total) + frac = value - int_part + if frac == 0.0 or base_total >= max_total: + return (base_total,), None + cpu = torch.Generator(device="cpu").manual_seed(_FORCE_ACCEPT_RNG_SEED) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, dtype=torch.float32, generator=cpu).cuda() + counter = _randint(0, 1000, (1,), gen) + force = ( + base_total, + frac, + max_total, + pool, + counter, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + ) + return force, counter + + +def _drafts_agreeing(target: torch.Tensor, draft_len: int, gen) -> torch.Tensor: + """``[G, K]`` int32 drafts that match ``target[:, :K]`` on a random prefix per row, behind + a wider row stride.""" + num_gens = target.shape[0] + agree = _randint(0, draft_len + 1, (num_gens,), gen) + cols = torch.arange(draft_len, device="cuda") + wrong = (target[:, :draft_len] + 1) % VOCAB + picked = torch.where(cols[None, :] < agree[:, None], target[:, :draft_len], wrong) + buf = torch.zeros((num_gens, draft_len + 3), dtype=torch.int32, device="cuda") + buf[:, :draft_len] = picked.to(torch.int32) + return buf[:, :draft_len] + + +# -------------------------------------------------------------------------------------------- +# Page tables, acceptance and next tokens +# -------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cell", k3_draft_block_table.CELLS, ids=_ids(k3_draft_block_table)) +def test_k3_draft_block_table_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require(k3_draft_block_table) + certifying(k3_draft_block_table) + g = _gen(cell) + s = cell.spec + rows, blocks, divisor, block_offset = s["rows"], s["blocks"], s["divisor"], s["block_offset"] + # The metadata's [pools, rows, K/V, blocks] offsets; the op reads plane 0 of pool 1. + meta = _randint(-5, 2**20, (2, rows + 3, 2, blocks), g, torch.int32) + held = _randint(0, blocks + 1, (rows,), g) + held[0] = blocks + if rows > 1: + held[1] = 0 + encoded = _randint(0, 2**18, (rows, blocks), g) * divisor + block_offset + cols = torch.arange(blocks, device="cuda") + # Unallocated columns hold the -1 placeholder; one allocated column sits below the offset. + plane = torch.where(cols[None, :] < held[:, None], encoded, -1) + plane[0, 0] = block_offset - 1 + meta[1, :rows, 0] = plane.to(torch.int32) + src = meta[1, :rows, 0] + meta_before = meta.clone() + num_slots = rows + 1 + table = torch.full((num_slots, blocks), -7, dtype=torch.int32, device="cuda") + counts = torch.full((num_slots,), -7, dtype=torch.int64, device="cuda") + partial = None + if s["partial"]: + chunks = -(-blocks // 2048) + partial = torch.full((num_slots * chunks,), -7, dtype=torch.int32, device="cuda") + + ref_table, ref_counts = k3_draft_block_table.reference( + src, divisor, table, counts, partial, block_offset=block_offset + ) + k3_draft_block_table(src, divisor, table, counts, partial, block_offset=block_offset) + _same(table, ref_table, "table") + _same(counts, ref_counts, "counts") + assert int(counts[rows]) == -7 and bool((table[rows] == -7).all()), "a row past rows moved" + assert torch.equal(meta, meta_before), "the op must not write the metadata" + + +@pytest.mark.parametrize("cell", k3_fused_strict_accept.CELLS, ids=_ids(k3_fused_strict_accept)) +def test_k3_fused_strict_accept_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require(k3_fused_strict_accept) + certifying(k3_fused_strict_accept) + g = _gen(cell) + s = cell.spec + num_contexts, num_gens, draft_len = s["num_contexts"], s["num_gens"], s["draft_len"] + batch = num_contexts + num_gens + width = draft_len + 1 + target = _randint(0, VOCAB, (num_contexts + num_gens * width,), g) + draft = _drafts_agreeing(target[num_contexts:].view(num_gens, width), draft_len, g) + force, counter = _force(cell, draft_len, g) + before = (target.clone(), draft.clone()) + counter_before = None if counter is None else counter.clone() + + ref_accepted, ref_num = k3_fused_strict_accept.reference( + target, draft, num_contexts, batch, force + ) + accepted, num_accepted = k3_fused_strict_accept(target, draft, num_contexts, batch, force) + _same(num_accepted, ref_num, "num_accepted") + # Context rows define column 0 only. + _same(accepted[:num_contexts, 0], ref_accepted[:num_contexts, 0], "context tokens") + _same(accepted[num_contexts:], ref_accepted[num_contexts:], "generation tokens") + if counter is not None: + assert int(counter) == int(counter_before) + 1, "the force counter advances by one" + assert torch.equal(target, before[0]) and torch.equal(draft, before[1]) + + +def _pairs(cell, rows: int, gen) -> tuple[tuple, torch.Tensor]: + """Every rank's interleaved (value, global index) pairs as the MNNVL exchange lays them out, + and the token each row selects (largest value, lowest index among equal values).""" + s = cell.spec + tp, splits = s["tp"], s["splits"] + n_slots = tp * splits + stride = 1 << (2 * n_slots - 1).bit_length() + shard = VOCAB // tp + # Quarter steps make equal maxima across slots common; -0.0 against +0.0 in row 0. + vals = torch.randn((rows, n_slots), generator=gen, device="cuda").mul(4).round().div(4) + vals[0] = -1.0 + vals[0, n_slots // 2] = 0.0 + vals[0, n_slots // 3] = -0.0 + rank = torch.arange(n_slots, device="cuda") // splits + idx = rank[None, :] * shard + _randint(0, shard, (rows, n_slots), gen) + # A split without valid columns reports (-inf, EMPTY_INDEX). + empty = torch.rand((rows, n_slots), generator=gen, device="cuda") < 0.05 + vals = torch.where(empty, float("-inf"), vals) + idx = torch.where(empty, EMPTY_INDEX, idx) + # Lanes past the slots hold NaN: the op must not read them. + buf = torch.full((rows, stride), float("nan"), dtype=torch.float32, device="cuda") + buf[:, 0 : 2 * n_slots : 2] = vals + buf[:, 1 : 2 * n_slots : 2] = idx.to(torch.float32) + best = vals.amax(dim=1, keepdim=True) + selected = torch.where(vals == best, idx, 2**31 - 1).amin(dim=1) + return (buf, stride, 2, n_slots), selected + + +@pytest.mark.parametrize("cell", k3_fused_select_accept.CELLS, ids=_ids(k3_fused_select_accept)) +def test_k3_fused_select_accept_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require(k3_fused_select_accept) + certifying(k3_fused_select_accept) + g = _gen(cell) + s = cell.spec + num_contexts, num_gens, draft_len = s["num_contexts"], s["num_gens"], s["draft_len"] + batch = num_contexts + num_gens + width = draft_len + 1 + pairs, selected = _pairs(cell, num_contexts + num_gens * width, g) + draft = _drafts_agreeing(selected[num_contexts:].view(num_gens, width), draft_len, g) + force, counter = _force(cell, draft_len, g) + counter_before = None if counter is None else counter.clone() + + kda_replay = kv_lens = None + if s["replay"]: + # The replay count is the first int32 of each slot's 16-byte record. + records = _randint(0, 9, (3 * batch + 2, 4), g, torch.int32) + perm = torch.randperm(records.shape[0], generator=g, device="cuda") + state_indices = perm[:batch].to(torch.int32).contiguous() + dummy = torch.zeros(batch, dtype=torch.bool, device="cuda") + if num_gens > 1: + dummy[num_contexts + 1] = True + kda_replay = (records[:, 0], state_indices, dummy) + kv_lens = _randint(10, 9000, (batch,), g, torch.int32) + records_before = records.clone() + kv_before = None if kv_lens is None else kv_lens.clone() + pair_buf_before, draft_before = pairs[0].clone(), draft.clone() + + ref_accepted, ref_num, ref_rewind, ref_prev, ref_kv = k3_fused_select_accept.reference( + pairs, draft, num_contexts, batch, force, kda_replay, kv_lens + ) + accepted, num_accepted, rewind = k3_fused_select_accept( + pairs, draft, num_contexts, batch, force, kda_replay=kda_replay, kv_lens=kv_lens + ) + _same(num_accepted, ref_num, "num_accepted") + _same(accepted[:num_contexts, 0], ref_accepted[:num_contexts, 0], "context tokens") + _same(accepted[num_contexts:], ref_accepted[num_contexts:], "generation tokens") + assert torch.equal( + ref_accepted[num_contexts:, 0], selected[num_contexts::width].to(torch.int32) + ) + if force is None: + target = selected[num_contexts:].view(num_gens, width) + assert torch.equal(ref_accepted[num_contexts:], target.to(torch.int32)) + if counter is not None: + assert int(counter) == int(counter_before) + 1, "the force counter advances by one" + if s["replay"]: + _same(rewind, ref_rewind, "rewind") + _same(kv_lens, ref_kv, "kv_lens") + _same(records[:, 0], ref_prev, "replay count") + assert torch.equal(records[:, 1:], records_before[:, 1:]), "the rest of each record moved" + # Main's eager record of the same acceptance. + main = records_before.clone() + state = KDAReplayState(draft_len) + state.prev_num_accepted_tokens = main[:, 0] + state.record_acceptance( + kda_replay[1][num_contexts:], ref_num[num_contexts:] - 1, kda_replay[2][num_contexts:] + ) + _same(main[:, 0], ref_prev, "record_acceptance") + assert torch.equal(kv_before[:num_contexts], kv_lens[:num_contexts]) + else: + assert rewind is None and ref_rewind is None and ref_prev is None and ref_kv is None + assert torch.equal(_bits(pairs[0]), _bits(pair_buf_before)) and torch.equal(draft, draft_before) + + +@pytest.mark.parametrize("cell", k3_next_new_tokens.CELLS, ids=_ids(k3_next_new_tokens)) +def test_k3_next_new_tokens_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require(k3_next_new_tokens) + certifying(k3_next_new_tokens) + g = _gen(cell) + s = cell.spec + batch, draft_len = s["batch"], s["draft_len"] + # The worker's [max_batch, K + 1] accepted buffer, of which the batch rows are passed. + accepted = _randint(0, VOCAB, (batch + 2, draft_len + 1), g, torch.int32)[:batch] + num_accepted = _randint(1, draft_len + 2, (batch,), g, torch.int32) + drafts = _randint(0, VOCAB, (batch, draft_len), g, torch.int32) + kv_lens = rewind = None + if s["rewind"]: + kv_lens = _randint(100, 9000, (batch,), g, torch.int32) + rewind = (1 - num_accepted).contiguous() + before = (accepted.clone(), num_accepted.clone(), drafts.clone()) + + ref_out, ref_kv = k3_next_new_tokens.reference(accepted, drafts, num_accepted, kv_lens, rewind) + out = k3_next_new_tokens(accepted, drafts, num_accepted, kv_lens, rewind) + _same(out, ref_out, "next new tokens") + if kv_lens is not None: + _same(kv_lens, ref_kv, "kv_lens") + assert all(torch.equal(a, b) for a, b in zip((accepted, num_accepted, drafts), before)) + + +# -------------------------------------------------------------------------------------------- +# The DFlash generation step's context bookkeeping +# -------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cell", k3_dflash_ctx_positions.CELLS, ids=_ids(k3_dflash_ctx_positions)) +def test_k3_dflash_ctx_positions_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require(k3_dflash_ctx_positions) + certifying(k3_dflash_ctx_positions) + g = _gen(cell) + s = cell.spec + num_gens, width, max_batch = s["num_gens"], s["width"], s["max_batch"] + dummy = max_batch + ctx_len = _randint(0, MAX_CTX - width, (max_batch + 1,), g) + ctx_len[dummy] = 0 + slots = torch.randperm(max_batch, generator=g, device="cuda")[:num_gens].contiguous() + if num_gens > 2: + slots[-2:] = dummy # graph padding rows share the dummy slot + before = ctx_len.clone() + ref = k3_dflash_ctx_positions.reference(ctx_len, slots, width) + out = k3_dflash_ctx_positions(ctx_len, slots, width) + _same(out, ref, "positions") + assert torch.equal(ctx_len, before), "the op must not write ctx_len" + + +def _step_state(cell, num_rows: int, first_row: int, cp_size: int, gen) -> dict: + """A generation step's worker state: context lengths per slot, the decoded page table and + its counts per batch row, the accepted tokens. Generation request ``g`` is batch row + ``first_row + g``; the last ``dummies`` requests are padding on the dummy slot. One request + sits at the context ceiling and one has fewer pages than its block needs.""" + s = cell.spec + num_gens, width = s["num_gens"], s["draft_len"] + 1 + max_batch = num_gens + 2 + dummy = max_batch + live = num_gens - s["dummies"] + slots = torch.randperm(max_batch, generator=gen, device="cuda")[:num_gens].contiguous() + slots[live:] = dummy + ctx_len = _randint(0, 60000, (max_batch + 1,), gen) + ctx_len[dummy] = 0 + ctx_len[slots[0]] = MAX_CTX - 3 + old = ctx_len[slots] + span = PAGE * cp_size + counts = torch.zeros(num_rows + 2, dtype=torch.int64, device="cuda") + rows = first_row + torch.arange(num_gens, device="cuda") + need = (old + 2 * width + span - 1) // span + counts[rows] = need + _randint(0, 2, (num_gens,), gen) + if live > 1: + counts[first_row + 1] = (old[1] + span - 1) // span # the block does not fit + max_blocks = int(counts.max()) + 1 + tables = _randint(0, 2**20, (num_rows + 2, max_blocks + 3), gen, torch.int32)[:, :max_blocks] + return dict( + slots=slots, + ctx_len=ctx_len, + counts=counts, + tables=tables, + dummy=dummy, + num_accepted=_randint(1, width + 1, (num_rows + 2,), gen, torch.int32), + accepted=_randint(0, VOCAB, (num_rows + 2, width), gen, torch.int32), + seq_lens=_randint(1, 9000, (num_rows + 2,), gen, torch.int32), + ) + + +@pytest.mark.parametrize("cell", k3_dflash_gen_prepare.CELLS, ids=_ids(k3_dflash_gen_prepare)) +def test_k3_dflash_gen_prepare_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require(k3_dflash_gen_prepare) + certifying(k3_dflash_gen_prepare) + g = _gen(cell) + s = cell.spec + num_gens, width, cp_size = s["num_gens"], s["draft_len"] + 1, s["cp_size"] + st = _step_state(cell, num_gens, 0, cp_size, g) + args = ( + st["num_accepted"], + st["accepted"], + st["slots"], + st["ctx_len"], + st["counts"], + st["tables"], + st["seq_lens"], + 0, + num_gens, + st["dummy"], + MAX_CTX, + 163837, + PAGE, + width, + width, + ) + before = {k: st[k].clone() for k in ("num_accepted", "accepted", "slots", "counts", "tables")} + ref, ref_ctx, ref_seq = k3_dflash_gen_prepare.reference( + *args, cp_size=cp_size, cp_rank=s["cp_rank"] + ) + out = k3_dflash_gen_prepare(*args, cp_size=cp_size, cp_rank=s["cp_rank"]) + assert set(out) == set(ref), (sorted(out), sorted(ref)) + for name in out: + _same(out[name], ref[name], name) + _same(st["ctx_len"], ref_ctx, "ctx_len") + _same(st["seq_lens"], ref_seq, "seq_lens_cuda") + assert all(torch.equal(st[k], v) for k, v in before.items()), "an input moved" + + +@pytest.mark.parametrize( + "cell", k3_dflash_paged_latent_prepare.CELLS, ids=_ids(k3_dflash_paged_latent_prepare) +) +def test_k3_dflash_paged_latent_prepare_cells(cell, certifying) -> None: # noqa: F811 + _require(k3_dflash_paged_latent_prepare) + certifying(k3_dflash_paged_latent_prepare) + g = _gen(cell) + s = cell.spec + num_gens, width, first_row = s["num_gens"], s["draft_len"] + 1, s["first_row"] + st = _step_state(cell, first_row + num_gens, first_row, 1, g) + args = ( + st["num_accepted"], + st["accepted"], + st["slots"], + st["ctx_len"], + st["counts"], + st["tables"], + st["seq_lens"], + first_row, + st["dummy"], + MAX_CTX, + PAGE, + width, + width, + ) + ref, ref_ctx, ref_seq = k3_dflash_paged_latent_prepare.reference(*args) + out = k3_dflash_paged_latent_prepare(*args) + assert type(out) is type(ref) + for name in out._fields: + _same(getattr(out, name), getattr(ref, name), name) + _same(st["ctx_len"], ref_ctx, "ctx_len") + _same(st["seq_lens"], ref_seq, "seq_lens_cuda") + + +@pytest.mark.parametrize( + "cell", k3_dflash_noise_embedding.CELLS, ids=_ids(k3_dflash_noise_embedding) +) +def test_k3_dflash_noise_embedding_cells(cell, certifying) -> None: # noqa: F811 + _require(k3_dflash_noise_embedding) + certifying(k3_dflash_noise_embedding) + g = _gen(cell) + s = cell.spec + combined = torch.randn((s["num_gens"] + 1, HIDDEN), generator=g, device="cuda") + combined = combined.to(torch.bfloat16) + combined[0, :3] = -0.0 + before = combined.clone() + ref = k3_dflash_noise_embedding.reference(combined, s["block_size"]) + out = k3_dflash_noise_embedding(combined, s["block_size"]) + _same(out, ref, "noise embedding") + assert torch.equal(_bits(combined), _bits(before)) + + +# -------------------------------------------------------------------------------------------- +# Paged-pool stores and the vocab-shard lookup +# -------------------------------------------------------------------------------------------- + + +def _latent_rows(shape, gen) -> torch.Tensor: + """BF16 latents with a spread of magnitudes, a few past the FP8 E4M3 range (448).""" + x = torch.randn(shape, generator=gen, device="cuda") + x = x * torch.logspace(-3, 1, shape[-1], device="cuda") + x[..., :4] *= 300.0 + return x.to(torch.bfloat16) + + +def _arena(pages: int, dtype, gen) -> torch.Tensor: + """Every layer's pages interleaved, ``[pages, layers, 1, page, D]``.""" + x = torch.randn((pages, LAYERS, 1, PAGE, LATENT), generator=gen, device="cuda") + return x.to(dtype) + + +def _distinct_slots(n: int, pages: int, gen) -> tuple[torch.Tensor, torch.Tensor]: + flat = torch.randperm(pages * PAGE, generator=gen, device="cuda")[:n] + return flat // PAGE, flat % PAGE + + +@pytest.mark.parametrize("cell", k3_store_context_latents.CELLS, ids=_ids(k3_store_context_latents)) +def test_k3_store_context_latents_cells(cell, certifying) -> None: # noqa: F811 + _require(k3_store_context_latents) + certifying(k3_store_context_latents) + g = _gen(cell) + s = cell.spec + num_tokens = s["num_gens"] * s["width"] + # The drafter's layer-major latents, read through the [N, L, D] transpose. + latent = _latent_rows((LAYERS, num_tokens, LATENT), g).transpose(0, 1) + dtype = torch.float8_e4m3fn if s["pool"] == "fp8" else torch.bfloat16 + arena = _arena(24, dtype, g) + pools = [arena[:, layer : layer + 1] for layer in range(LAYERS)] + pool_ptrs = torch.tensor([p.data_ptr() for p in pools], dtype=torch.int64, device="cuda") + pages, offsets = _distinct_slots(num_tokens, arena.shape[0], g) + # 1 for an accepted token, 0 past the accepted prefix, -1 on another rank's page. + mask = _randint(0, 2, (num_tokens,), g, torch.int8) + if s["cp_size"] > 1: + other = _randint(0, s["cp_size"], (num_tokens,), g) != 0 + mask = torch.where(other, torch.full_like(mask, -1), mask) + latent_before = latent.clone() + + ref = k3_store_context_latents.reference( + latent, pools, pool_ptrs, pages, offsets, mask, mla=s["mla"] + ) + k3_store_context_latents(latent, pools, pool_ptrs, pages, offsets, mask, mla=s["mla"]) + for layer, (pool, expected) in enumerate(zip(pools, ref)): + _same(pool, expected, f"pool of layer {layer}") + assert torch.equal(_bits(latent), _bits(latent_before)) + + +@pytest.mark.parametrize("cell", k3_scatter_rows_to_pool.CELLS, ids=_ids(k3_scatter_rows_to_pool)) +def test_k3_scatter_rows_to_pool_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require(k3_scatter_rows_to_pool) + certifying(k3_scatter_rows_to_pool) + g = _gen(cell) + s = cell.spec + rows = s["rows"] + src = _latent_rows((rows, LATENT), g) + arena = _arena(16, torch.float8_e4m3fn, g) + pool = arena[:, 2, 0] # one layer's [pages, page, D] view of the interleaved pool + pages, offsets = _distinct_slots(rows, arena.shape[0], g) + rows_mask = None + if s["rows_mask"]: + # Block positions on another rank's pages map onto slots this rank's rows also write: + # only the masked-in rows may land. + rows_mask = (_randint(0, 3, (rows,), g) == 0).to(torch.int8) + rows_mask[0] = 1 + pages = torch.where(rows_mask.bool(), pages, pages[0]) + offsets = torch.where(rows_mask.bool(), offsets, offsets[0]) + src_before = src.clone() + ref = k3_scatter_rows_to_pool.reference(src, pool, pages, offsets, rows_mask) + k3_scatter_rows_to_pool(src, pool, pages, offsets, rows_mask=rows_mask) + _same(pool, ref, "pool") + assert torch.equal(_bits(src), _bits(src_before)) + + +@pytest.mark.parametrize("cell", k3_vocab_shard_embedding.CELLS, ids=_ids(k3_vocab_shard_embedding)) +def test_k3_vocab_shard_embedding_cells(cell, certifying) -> None: # noqa: F811 + _require(k3_vocab_shard_embedding) + certifying(k3_vocab_shard_embedding) + g = _gen(cell) + s = cell.spec + shard = VOCAB // s["tp"] + start = s["rank"] * shard + weight = torch.randn((shard, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + weight[0, :5] = -0.0 + # Bonus ids from anywhere in the vocab, this shard's first and last rows, then the mask id. + ids = _randint(0, VOCAB, (s["num_ids"],), g) + ids[0] = start + if s["num_ids"] > 2: + ids[1] = start + shard - 1 + ids[-1] = 163837 + ref = k3_vocab_shard_embedding.reference(ids, weight, start, start + shard) + out = k3_vocab_shard_embedding(ids, weight, start, start + shard) + _same(out, ref, "shard lookup") diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_splitk_dense.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_splitk_dense.py new file mode 100644 index 000000000000..0ba6965d1729 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_splitk_dense.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the K3 SM107 target's k3_splitk_dense entry: the certified cells against the +entry's FP64 reference.""" + +import zlib + +import pytest +import torch + +pytest.importorskip("cutlass") + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.dspark_draft_gemm_push import ( # noqa: E402, E501 + ensure_cute_runtime, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.low_m_bf16_splitk import ( # noqa: E402, E501 + SplitKTactic, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec import ( # noqa: E402, E501 + k3_splitk_dense as entry, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import ( # noqa: E402, F401 — path above; fixture by name + assert_within_ulp, + certifying, +) + +op = entry.k3_splitk_dense + + +def _on_archs() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in op.ARCHS} + + +pytestmark = pytest.mark.skipif(not _on_archs(), reason="not an ARCHS architecture") + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """2 ulp per element against the row scale, 0.5 ulp relative RMS. + + One BF16 rounding on each side: the kernel sums FP32 partial products per split-K slice and + then the slices, the reference is an FP64 product. With unit-scale operands the FP32 + accumulation over K <= 7168 is off by far less than a BF16 step of the output, so the two + differ only where the FP32 value sits at a BF16 rounding midpoint, by one step of that + element: at most 2 units of ``ULP[bf16]`` against the row scale. Those flips are rare enough + that the relative RMS stays far below half an ulp. A dropped or doubled K slice or a + transposed tile misses by O(1). + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.5) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + spec = cell.spec + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + a = torch.randn(spec["rows"], spec["k"], device="cuda", generator=gen).to(torch.bfloat16) + weight = (torch.randn(spec["n"], spec["k"], device="cuda", generator=gen) / 16).to( + torch.bfloat16 + ) + before = a.clone() + out = torch.full((spec["rows"], spec["n"]), 5.0, dtype=torch.bfloat16, device="cuda") + ensure_cute_runtime() + tactic = SplitKTactic(*spec["tactic"]) + result = op(a, weight.t(), None, out, True, tactic) + assert result.data_ptr() == out.data_ptr() + assert torch.equal(a, before), "the op must not write its input" + _compare(out, op.reference(a, weight.t(), None, out, True, tactic)) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_vocab_parallel_entries.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_vocab_parallel_entries.py new file mode 100644 index 000000000000..25ddb3f930d4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/spec/test_k3_vocab_parallel_entries.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The vocab-parallel selection entries (``ops/spec``) over their certified cells: one rank of a +vocab-sharded head on one GPU, the exchange emulated as the sum of every rank's zero-filled pair +buffer (what the MNNVL all-reduce computes).""" + +import types +import zlib + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("needs a CUDA GPU", allow_module_level=True) + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.dflash import ( # noqa: E402, E501 + _PairExchange, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec import ( # noqa: E402, E501 + _reference, + k3_markov_add_local_argmax, + k3_markov_reduce_embed, + k3_vocab_local_select, + k3_vocab_reduce_pairs, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +VOCAB = 163840 +MARKOV_RANK = 256 +_ARCH = "sm_{}{}".format(*torch.cuda.get_device_capability()) + + +def _require_arch(op) -> None: + if _ARCH not in {a.value for a in op.ARCHS}: + pytest.skip(f"{type(op).__name__} is not certified on {_ARCH}") + + +def _gen(cell) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + + +def _layout(tp: int, rank: int, rows: int, shard: int, sentinel: float = 3.0): + """The MNNVL exchange's layout of ``rank``, its buffer filled with ``sentinel`` so that the + zero fill shows.""" + exchange = types.SimpleNamespace(backend="mnnvl", tp_size=tp, tp_rank=rank) + layout = _PairExchange.layout(exchange, rows, shard, torch.device("cuda")) + layout.buffer.fill_(sentinel) + return layout + + +def _logits(rows: int, cols: int, gen, ties: bool) -> torch.Tensor: + """BF16 logits of the head's scale; with ``ties`` some rows repeat their maximum in other + columns (other splits and shards) and some hold NaNs.""" + x = (2.0 * torch.randn(rows, cols, device="cuda", generator=gen)).to(torch.bfloat16) + if ties: + cols_at = torch.randint(0, cols, (rows, 3), device="cuda", generator=gen) + for r in range(0, rows, 3): + top = x[r].float().max() + 1.0 + x[r, cols_at[r]] = top.to(torch.bfloat16) + for r in range(1, rows, 7): + x[r, cols_at[r, :2]] = float("nan") + return x + + +def _assert_pairs_equal(out: torch.Tensor, ref: torch.Tensor) -> None: + # Exact, NaN == NaN (the kernel keeps the logit's NaN payload) and -0 == +0. + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0, equal_nan=True) + + +def _own_slots(buf: torch.Tensor, layout) -> tuple[torch.Tensor, torch.Tensor]: + cols = buf[:, 2 * layout.slot0 : 2 * (layout.slot0 + layout.splits)] + return cols[:, 0::2], cols[:, 1::2].to(torch.int64) + + +# --------------------------------------------------------------------------------------------- +# k3_vocab_local_select +# --------------------------------------------------------------------------------------------- + +_SELECT = k3_vocab_local_select.k3_vocab_local_select + + +@pytest.mark.parametrize("cell", _SELECT.CELLS, ids=[c.why[:48] for c in _SELECT.CELLS]) +def test_vocab_local_select_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require_arch(_SELECT) + certifying(_SELECT) + spec, gen = cell.spec, _gen(cell) + tp, rank, rows = spec["tp"], spec["rank"], spec["rows"] + shard = VOCAB // tp + col_base = rank * shard + local = _logits(rows, shard, gen, spec["ties"]) + layout = _layout(tp, rank, rows, shard) + before = layout.buffer.clone() + local_before = local.clone() + sampling = None + if spec["sampling"]: + from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import GREEDY_TEMPERATURE_THRESHOLD + + temperatures = 0.3 + 1.2 * torch.rand(rows, device="cuda", generator=gen) + temperatures[::5] = 0.0 # greedy rows inside a sampled batch + seeds = torch.randint(0, 2**62, (rows,), device="cuda", generator=gen) + offsets = torch.randint(0, 2**40, (rows,), device="cuda", generator=gen) + sampling = (temperatures, seeds, offsets, float(GREEDY_TEMPERATURE_THRESHOLD)) + + _SELECT(local, col_base, layout, sampling) + # Bitwise: the tied cells hold NaNs. + assert torch.equal(local.view(torch.int16), local_before.view(torch.int16)), "logits written" + ref = _SELECT.reference(local, col_base, layout._replace(buffer=before), sampling) + out = layout.buffer + # Columns outside this rank's slots: zero (the all-reduce sums them away). + own = torch.zeros_like(out, dtype=torch.bool) + own[:, layout.zero_lo : layout.zero_hi] = True + assert torch.equal(out[~own], torch.zeros_like(out[~own])), "other ranks' columns not zeroed" + if sampling is None: + _assert_pairs_equal(out, ref) + return + _assert_sampled_pairs(out, ref, layout, local, col_base, sampling) + + +def _assert_sampled_pairs(out, ref, layout, local, col_base, sampling) -> None: + """The kernel scores in FP32, the reference in FP64. A score is at most about 32 in + magnitude (logits / T with T >= 0.3, plus Gumbel noise below 17) and the kernel's carries + at most four roundings of 1 ulp at that scale (the temperature product, the two logs of the + noise and the sum): 2**-16. The pairs agree within 2**-14 (4x that), and where the winning + columns differ the kernel's must score within 2**-14 of the reference's winner in the same + split (a near-tie the FP32 evaluation resolves the other way).""" + tol = 2.0**-14 + k_val, k_idx = _own_slots(out, layout) + r_val, r_idx = _own_slots(ref, layout) + scores = k3_vocab_local_select.sampled_scores(local, col_base, sampling) + same = k_idx == r_idx + torch.testing.assert_close(k_val[same], r_val[same], rtol=0.0, atol=tol) + if not bool(same.all()): + rows, splits = (~same).nonzero(as_tuple=True) + k_cols, r_cols = k_idx[rows, splits] - col_base, r_idx[rows, splits] - col_base + assert bool((k_cols // layout.chunk == splits).all()), "a winner outside its split" + gap = scores[rows, r_cols] - scores[rows, k_cols] + assert float(gap.max()) <= tol, f"winners differ beyond a near-tie: {gap.max():.3e}" + # A near-tie is rare: a handful of the row x split winners at most. + assert int((~same).sum()) <= max(2, same.numel() // 100) + + +# --------------------------------------------------------------------------------------------- +# k3_vocab_reduce_pairs +# --------------------------------------------------------------------------------------------- + +_REDUCE = k3_vocab_reduce_pairs.k3_vocab_reduce_pairs + + +def _exchanged_pairs(values: torch.Tensor, tp: int) -> tuple[torch.Tensor, object]: + """Every rank's published pairs of ``values`` ``[rows, vocab]`` summed: the exchange's + result (interleaved slots: row stride ``stride_out``, slot stride 2).""" + rows, vocab = values.shape + shard = vocab // tp + total, layout = None, None + for rank in range(tp): + layout = _layout(tp, rank, rows, shard, sentinel=0.0) + part = _reference.publish_pairs( + layout, values[:, rank * shard : (rank + 1) * shard].float(), rank * shard + ) + total = part if total is None else total + part + return total, layout + + +@pytest.mark.parametrize("cell", _REDUCE.CELLS, ids=[c.why[:48] for c in _REDUCE.CELLS]) +def test_vocab_reduce_pairs_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require_arch(_REDUCE) + certifying(_REDUCE) + spec, gen = cell.spec, _gen(cell) + tp, rows = spec["tp"], spec["rows"] + full = _logits(rows, VOCAB, gen, spec["ties"]) + pairs, layout = _exchanged_pairs(full, tp) + n_slots = tp * layout.splits + pairs_before = pairs.clone() + out_dtype = getattr(torch, spec["out_dtype"]) + out = torch.full((rows,), -5, dtype=out_dtype, device="cuda") + ret = _REDUCE(pairs, layout.stride_out, 2, n_slots, out) + assert ret is out + assert torch.equal(pairs.view(torch.int32), pairs_before.view(torch.int32)), "pairs written" + ref = _REDUCE.reference(pairs, layout.stride_out, 2, n_slots, out) + assert ref.dtype == out_dtype and torch.equal(out, ref) + # Cross-check: the argmax of the gathered logits. + cols = torch.arange(VOCAB, device="cuda").expand(rows, -1) + _, gathered = _reference.argmax_pairs(full.float(), cols) + assert torch.equal(out.to(torch.int64), gathered) + if not spec["ties"]: + assert torch.equal(out.to(torch.int64), torch.argmax(full.float(), dim=-1)) + + +# --------------------------------------------------------------------------------------------- +# The multi-launch Markov chain: k3_markov_reduce_embed and k3_markov_add_local_argmax +# --------------------------------------------------------------------------------------------- + +_EMBED = k3_markov_reduce_embed.k3_markov_reduce_embed +_ADD = k3_markov_add_local_argmax.k3_markov_add_local_argmax + + +def _markov_w1(gen) -> torch.Tensor: + # [vocab, R] with entries of 0.25: a bias GEMV over R = 256 of unit scale. + return (0.25 * torch.randn(VOCAB, MARKOV_RANK, device="cuda", generator=gen)).to(torch.bfloat16) + + +@pytest.mark.parametrize("cell", _EMBED.CELLS, ids=[c.why[:48] for c in _EMBED.CELLS]) +def test_markov_reduce_embed_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require_arch(_EMBED) + certifying(_EMBED) + spec, gen = cell.spec, _gen(cell) + tp, reqs, block, step = spec["tp"], spec["reqs"], spec["block"], spec["step"] + markov_w1 = _markov_w1(gen) + anchors = torch.randint(0, VOCAB, (reqs,), device="cuda", generator=gen, dtype=torch.int64) + anchors[3] = -1 + anchors[7] = VOCAB + emb = torch.full((reqs, MARKOV_RANK), 9.0, dtype=torch.bfloat16, device="cuda") + tokens = torch.full((reqs, block), -7, dtype=torch.int32, device="cuda") + if step == 0: + pairs, stride_row, stride_slot, n_slots = anchors, 0, 0, 1 + values = None + else: + values = _logits(reqs, VOCAB, gen, ties=False) + pairs, layout = _exchanged_pairs(values, tp) + stride_row, stride_slot, n_slots = layout.stride_out, 2, tp * layout.splits + args = (pairs, stride_row, stride_slot, n_slots, anchors, markov_w1) + ref_emb, ref_tokens = _EMBED.reference(*args, emb, tokens, VOCAB, step) + _EMBED(*args, emb, tokens, VOCAB, step) + assert torch.equal(emb, ref_emb) and torch.equal(tokens, ref_tokens) + if step == 0: + assert torch.equal(emb[3], torch.zeros_like(emb[3])), "an out-of-vocab anchor's row" + assert torch.equal(tokens, torch.full_like(tokens, -7)), "step 0 stores no token" + else: + prev = torch.argmax(values.float(), dim=-1) + assert torch.equal(tokens[:, step - 1].to(torch.int64), prev) + assert torch.equal(emb, markov_w1[prev]) + others = [c for c in range(block) if c != step - 1] + assert bool((tokens[:, others] == -7).all()), "only token step - 1 is stored" + + +@pytest.mark.parametrize("cell", _ADD.CELLS, ids=[c.why[:48] for c in _ADD.CELLS]) +def test_markov_add_local_argmax_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + _require_arch(_ADD) + certifying(_ADD) + spec, gen = cell.spec, _gen(cell) + tp, rank, reqs, block, step = ( + spec["tp"], + spec["rank"], + spec["reqs"], + spec["block"], + spec["step"], + ) + shard = VOCAB // tp + col_base = rank * shard + local_block = _logits(reqs * block, shard, gen, ties=False).view(reqs, block, shard) + w2_shard = (0.25 * torch.randn(shard, MARKOV_RANK, device="cuda", generator=gen)).to( + torch.bfloat16 + ) + emb = (0.25 * torch.randn(reqs, MARKOV_RANK, device="cuda", generator=gen)).to(torch.bfloat16) + bias = torch.nn.functional.linear(emb, w2_shard) + step_logits = local_block[:, step] + assert not step_logits.is_contiguous() + layout = _layout(tp, rank, reqs, shard) + before = layout.buffer.clone() + block_before = local_block.clone() + _ADD(step_logits, bias, col_base, layout) + assert torch.equal(local_block, block_before), "the op must not write the logits" + ref = _ADD.reference(step_logits, bias, col_base, layout._replace(buffer=before)) + _assert_pairs_equal(layout.buffer, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_fused_kda_conv_prefill.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_fused_kda_conv_prefill.py new file mode 100644 index 000000000000..71b52e455834 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_fused_kda_conv_prefill.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the fused_kda_conv_prefill entry: every cell of ``fused_kda_conv_prefill.CELLS``, +built the way the tp1 KDA layer calls it (strided QKVG column view, transposed weight copy, padded +conv pool slots), against the fp64 reference: Q/K/V within the stated bf16 band, beta and the +written conv states bitwise, untouched slots and slot padding unchanged.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.fused_kda_conv_prefill import ( # noqa: E501 + TP1_HEADS, + fused_kda_conv_prefill, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import assert_within_ulp, validating + +HEAD_DIM = 128 +WIDTH = 4 +SLOTS = 8 +SLOT_PAD = 64 # elements past each slot's [3 * H * D, W - 1] state, as the cache manager pads + + +def _driven_here() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in fused_kda_conv_prefill.ARCHS} + + +pytestmark = pytest.mark.skipif(not _driven_here(), reason="fused_kda_conv_prefill.ARCHS") + + +def _inputs(spec: dict, seed: int) -> dict: + gen = torch.Generator(device="cuda").manual_seed(seed) + heads, seq_lens = TP1_HEADS, spec["seq_lens"] + channels = 3 * heads * HEAD_DIM + tokens = sum(seq_lens) + + def randn(*shape, scale=1.0): + return torch.randn(*shape, generator=gen, device="cuda") * scale + + qkvg = randn(tokens, 4 * heads * HEAD_DIM).bfloat16() + weight_t = randn(WIDTH, channels, scale=0.5).bfloat16() + pool = randn(SLOTS, channels * (WIDTH - 1) + SLOT_PAD).bfloat16() + conv_states = pool[:, : channels * (WIDTH - 1)].view(SLOTS, channels, WIDTH - 1) + starts = [0] + for n in seq_lens: + starts.append(starts[-1] + n) + slots = torch.randperm(SLOTS, generator=torch.Generator().manual_seed(seed))[: len(seq_lens)] + args = dict( + qkv=qkvg[:, :channels], + conv_weight=weight_t.t(), + conv_states=conv_states, + query_start_loc=torch.tensor(starts, dtype=torch.int32, device="cuda"), + cache_indices=slots.to(torch.int32).cuda(), + has_initial_state=torch.tensor(spec["has_init"], dtype=torch.bool, device="cuda"), + num_heads=heads, + head_dim=HEAD_DIM, + ) + if spec["prepad"]: + bfa = randn(1, tokens, HEAD_DIM + heads).bfloat16() + args.update( + padded_tokens=-(-tokens // 64) * 64, + beta=bfa[..., HEAD_DIM : HEAD_DIM + heads], + ) + return args, pool + + +@pytest.mark.parametrize( + "cell", fused_kda_conv_prefill.CELLS, ids=[c.why[:60] for c in fused_kda_conv_prefill.CELLS] +) +@torch.no_grad() +def test_cells(cell) -> None: + args, pool = _inputs(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + pool_before = pool.clone() + ref_out, ref_states = fused_kda_conv_prefill.reference(**args) + with validating(fused_kda_conv_prefill): + out = fused_kda_conv_prefill(**args) + + assert len(out) == len(ref_out), cell.why + for name, got, want in zip("qkvb", out, ref_out): + assert got.shape == want.shape and got.dtype == want.dtype, (cell.why, name) + if name == "b": + assert torch.equal(got, want), (cell.why, name) + continue + # The convolution accumulates in fp32 (here fp64) before one bf16 rounding, and Q/K are + # normalized from that rounded value: an element can land one bf16 ulp away, and a + # normalized row inherits that ulp, so two ulps of the row scale bound both. + assert_within_ulp(got, want, element_ulp=2.0, rms_ulp=0.25) + + state_size = args["conv_states"].shape[1] * (WIDTH - 1) + assert torch.equal(args["conv_states"], ref_states), f"{cell.why}: conv states" + assert torch.equal(pool[:, state_size:], pool_before[:, state_size:]), "slot padding written" + written = set(args["cache_indices"].tolist()) + for slot in range(SLOTS): + if slot not in written: + assert torch.equal(pool[slot], pool_before[slot]), f"slot {slot} written" diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_fused_kda_post_conv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_fused_kda_post_conv.py new file mode 100644 index 000000000000..badc9cbe345c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_fused_kda_post_conv.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the fused_kda_post_conv entry: every cell of ``fused_kda_post_conv.CELLS`` against +its fp64 reference (Q/K within one bf16 ulp of the per-head scale, V bitwise), plus the +chunk-padded form with beta.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.fused_kda_post_conv import ( # noqa: E501 + fused_kda_post_conv, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import assert_within_ulp, validating + +HEAD_DIM = 128 + + +def _driven_here() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in fused_kda_post_conv.ARCHS} + + +pytestmark = pytest.mark.skipif(not _driven_here(), reason="fused_kda_post_conv.ARCHS") + + +def _packed(heads: int, tokens: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(3 * heads * HEAD_DIM, tokens, generator=gen, device="cuda").bfloat16() + + +def _check(out, ref, context): + assert len(out) == len(ref), context + for name, got, want in zip("qkvb", out, ref): + assert got.shape == want.shape and got.dtype == want.dtype, (context, name) + assert got.is_contiguous(), (context, name) + if name in "qk": + # fp32 norm in the kernel vs fp64 here, each rounded once to bf16: an element can + # land one bf16 ulp of its own value away, at most 2**-7 of the head's largest + # magnitude, i.e. two ULP units of the row scale (1.27 measured on GB300). + assert_within_ulp(got, want, element_ulp=2.0, rms_ulp=0.25) + else: + assert torch.equal(got, want), (context, name) + + +@pytest.mark.parametrize( + "cell", fused_kda_post_conv.CELLS, ids=[c.why[:60] for c in fused_kda_post_conv.CELLS] +) +@torch.no_grad() +def test_cells(cell) -> None: + heads, tokens = cell.spec["heads"], cell.spec["tokens"] + packed = _packed(heads, tokens, zlib.crc32(cell.why.encode()) % 2**31) + with validating(fused_kda_post_conv): + out = fused_kda_post_conv(packed, heads, HEAD_DIM) + _check(out, fused_kda_post_conv.reference(packed, heads, HEAD_DIM), cell.why) + + +@pytest.mark.parametrize("tokens", [1, 63, 65, 193]) +@torch.no_grad() +def test_chunk_padded_with_beta(tokens: int) -> None: + """features.kda_prepad_prefill's form: Q/K/V and beta (a column view of the [f_a | b] + output) zero-padded to the 64-token chunk.""" + heads = 6 + packed = _packed(heads, tokens, tokens) + projection = torch.randn(1, tokens, HEAD_DIM + 8, device="cuda").bfloat16() + beta = projection[..., HEAD_DIM : HEAD_DIM + heads] + padded = -(-tokens // 64) * 64 + with validating(fused_kda_post_conv): + out = fused_kda_post_conv(packed, heads, HEAD_DIM, padded_tokens=padded, beta=beta) + ref = fused_kda_post_conv.reference(packed, heads, HEAD_DIM, padded_tokens=padded, beta=beta) + _check(out, ref, f"tokens={tokens}") diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_k3_kda_fab.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_k3_kda_fab.py new file mode 100644 index 000000000000..6295e194a5f0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_k3_kda_fab.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the k3_kda_fab_proj and k3_kda_fab_fa entries: every cell of +``k3_kda_fab_proj.CELLS`` and ``k3_kda_fab_fa.CELLS`` against the fp64 GEMM references, and f_a / +beta of the two entries bitwise equal.""" + +import zlib + +import pytest +import torch + +pytest.importorskip("cutlass") + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.k3_kda_fab_fa import ( # noqa: E501, E402 + k3_kda_fab_fa, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.k3_kda_fab_proj import ( # noqa: E501, E402 + FA_DIM, + HIDDEN, + bfa_rows, + k3_kda_fab_proj, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import assert_within_ulp, validating + + +def _driven_here() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in k3_kda_fab_proj.ARCHS} + + +pytestmark = pytest.mark.skipif(not _driven_here(), reason="k3_kda_fab_proj.ARCHS") + + +def _operands(heads: int, rows: int, seed: int): + """``x`` with a few outlier channels (as the KDA input carries), ``[f_a; b]`` with its zero + padding rows, and ``f_b``.""" + gen = torch.Generator().manual_seed(seed) + x = torch.randn(rows, HIDDEN, generator=gen) + x[:, torch.randperm(HIDDEN, generator=gen)[:16]] *= 30.0 + w_bfa = torch.randn(bfa_rows(heads), HIDDEN, generator=gen) * 0.02 + w_bfa[FA_DIM + heads :] = 0 + w_fb = torch.randn(heads * FA_DIM, FA_DIM, generator=gen) * 0.05 + return (t.bfloat16().cuda() for t in (x, w_bfa, w_fb)) + + +def _seed(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +# The kernel sums fp32 partials of a 7168-long dot product in a fixed split order; this side +# sums in fp64. Both round once to bf16, so an element lands at most one ulp of its row's scale +# away, and the forget gate (a GEMM of those bf16 values) inherits that ulp: two ulps bound it. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.25 + + +@pytest.mark.parametrize( + "cell", k3_kda_fab_proj.CELLS, ids=[c.why[:60] for c in k3_kda_fab_proj.CELLS] +) +@torch.no_grad() +def test_fab_proj_cells(cell) -> None: + heads, rows = cell.spec["heads"], cell.spec["rows"] + x, w_bfa, w_fb = _operands(heads, rows, _seed(cell.why)) + with validating(k3_kda_fab_proj): + beta, gate = k3_kda_fab_proj(x, w_bfa, w_fb, heads) + ref_beta, ref_gate = k3_kda_fab_proj.reference(x, w_bfa, w_fb, heads) + assert beta.shape == (rows, heads) and gate.shape == (rows, heads * FA_DIM), cell.why + assert_within_ulp(beta, ref_beta, ELEMENT_ULP, RMS_ULP) + assert_within_ulp(gate, ref_gate, ELEMENT_ULP, RMS_ULP) + + +@pytest.mark.parametrize("cell", k3_kda_fab_fa.CELLS, ids=[c.why[:60] for c in k3_kda_fab_fa.CELLS]) +@torch.no_grad() +def test_fab_fa_cells(cell) -> None: + heads, rows = cell.spec["heads"], cell.spec["rows"] + x, w_bfa, w_fb = _operands(heads, rows, _seed(cell.why)) + with validating(k3_kda_fab_fa): + bfa = k3_kda_fab_fa(x, w_bfa, heads, w_fb) + assert bfa.shape == (rows, bfa_rows(heads)), cell.why + assert_within_ulp(bfa, k3_kda_fab_fa.reference(x, w_bfa, heads, w_fb), ELEMENT_ULP, RMS_ULP) + # f_a mode is the full kernel's GEMM 1 and reduction: bitwise its beta (and f_a). + beta, _ = k3_kda_fab_proj(x, w_bfa, w_fb, heads) + assert torch.equal(bfa[:, FA_DIM : FA_DIM + heads], beta), cell.why diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_k3_l2_prefetch.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_k3_l2_prefetch.py new file mode 100644 index 000000000000..dd6e0d1be1aa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_k3_l2_prefetch.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the k3_l2_prefetch entry: every cell of ``k3_l2_prefetch.CELLS`` runs (eagerly and +in a CUDA graph) and leaves the prefetched bytes and their neighbours unchanged, which is all its +no-op reference states.""" + +import pytest +import torch + +pytest.importorskip("cutlass") + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.k3_l2_prefetch import ( # noqa: E501, E402 + tensor_range, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.k3_l2_prefetch import ( # noqa: E501, E402 + k3_l2_prefetch, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import validating + +GUARD = 4096 # bytes on each side of every range + + +def _driven_here() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in k3_l2_prefetch.ARCHS} + + +pytestmark = pytest.mark.skipif(not _driven_here(), reason="k3_l2_prefetch.ARCHS") + + +@pytest.mark.parametrize( + "cell", k3_l2_prefetch.CELLS, ids=[c.why[:60] for c in k3_l2_prefetch.CELLS] +) +@torch.no_grad() +def test_cells(cell) -> None: + gen = torch.Generator(device="cuda").manual_seed(7) + buffers = [ + torch.randint( + 0, 256, (nbytes + 2 * GUARD,), generator=gen, device="cuda", dtype=torch.uint8 + ) + for nbytes in cell.spec["nbytes"] + ] + before = [b.clone() for b in buffers] + ranges = [tensor_range(b[GUARD : GUARD + n]) for b, n in zip(buffers, cell.spec["nbytes"])] + assert [n for _, n in ranges] == list(cell.spec["nbytes"]), "ranges must stay whole" + + with validating(k3_l2_prefetch): + assert k3_l2_prefetch(ranges) is None + torch.cuda.synchronize() + assert k3_l2_prefetch.reference(ranges) is None + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + k3_l2_prefetch(ranges) + graph.replay() + torch.cuda.synchronize() + for b, want in zip(buffers, before): + assert torch.equal(b, want), cell.why diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_kda_mtp_decode.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_kda_mtp_decode.py new file mode 100644 index 000000000000..480e45b9d1b0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_kda_mtp_decode.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the kda_mtp_decode entry: every cell of ``kda_mtp_decode.CELLS``, built as the KDA +layer's verify calls it (row-strided QKVG and [f_a | b] sections, dim-contiguous conv caches, +scattered slots with slot-indexed accepted counts, the variant's own arguments), against the +entry's fp64 reference: the output, the committed state and conv window, and the cached drafts.""" + +import zlib + +import pytest +import torch + +pytest.importorskip("cutlass") + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.kda_mtp_decode import ( # noqa: E501, E402 + kda_mtp_decode, +) + +__extra_import_path__ = ["../../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 + +K = 128 +W = 4 +E4M3_TOP_STEP = 32.0 # one E4M3 step in [256, 448], the binade a scaled head maximum lands in + + +def _driven_here() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in kda_mtp_decode.ARCHS} + + +pytestmark = pytest.mark.skipif(not _driven_here(), reason="kda_mtp_decode.ARCHS") + + +def _call_args(spec: dict, seed: int) -> dict: + """The arguments K3KDAKernelDispatch passes for this variant, on a pool that already went + through several verify rounds.""" + heads, requests, num_spec, variant = ( + spec["heads"], + spec["requests"], + spec["num_spec"], + spec["variant"], + ) + gen = torch.Generator().manual_seed(seed) + dim = heads * K + rows = requests * (1 + num_spec) + pool = max(2 * requests, 8) + + def rn(*shape, scale=1.0, dtype=torch.float32): + return (torch.randn(*shape, generator=gen) * scale).to(dtype).cuda() + + def conv_cache(): # dim-contiguous [slots, dim, W - 1 + num_spec] + return rn(pool, W - 1 + num_spec, dim, dtype=torch.bfloat16).transpose(1, 2) + + qkvg = rn(rows, 4 * dim, scale=0.5, dtype=torch.bfloat16) + bfa = rn(rows, (K + heads + 7) // 8 * 8, dtype=torch.bfloat16) # [f_a | b | pad] + # V-first state rows whose scales spread over two orders of magnitude. + state = rn(pool, heads, K, K, scale=0.05) * torch.logspace(-1, 1, K, device="cuda").view( + 1, 1, K, 1 + ) + slots = torch.randperm(pool, generator=gen)[:requests] + accepted = torch.zeros(pool, dtype=torch.int32) + accepted[slots] = torch.randint(0, num_spec + 1, (requests,), generator=gen, dtype=torch.int32) + args = dict( + x_q=qkvg[:, :dim].view(1, rows, heads, K), + x_k=qkvg[:, dim : 2 * dim].view(1, rows, heads, K), + x_v=qkvg[:, 2 * dim : 3 * dim].view(1, rows, heads, K), + w_q=rn(dim, W, scale=0.5), + w_k=rn(dim, W, scale=0.5), + w_v=rn(dim, W, scale=0.5), + cs_q=conv_cache(), + cs_k=conv_cache(), + cs_v=conv_cache(), + g=None, + beta=bfa[:, K : K + heads].view(1, rows, heads), + A_log=rn(heads, scale=0.5), + dt_bias=rn(dim, scale=0.5), + recurrent_state=state, + k_cache=rn(pool, num_spec, dim, scale=0.08, dtype=torch.bfloat16), + g_cache=-torch.rand(pool, num_spec, dim, generator=gen).cuda() * 0.5, + v_cache=rn(pool, num_spec, dim, dtype=torch.bfloat16), + beta_cache=torch.rand(pool, num_spec, heads, generator=gen).to(torch.bfloat16).cuda(), + ssm_state_indices=slots.to(torch.int32).cuda(), + cu_seqlens=None, + num_spec=num_spec, + num_accepted_tokens=accepted.cuda(), + lower_bound=-5.0, + scale=K**-0.5, + packed_token_layout=True, + onorm_g=qkvg[:, 3 * dim :].view(1, rows, heads, K), + onorm_weight=(0.75 + 0.5 * torch.rand(K, generator=gen)).cuda(), + onorm_eps=1e-5, + fuse_output_norm=True, + ) + if variant == "fused_fb": + # f_a in {-1, 0, 1} / 4 and f_b in {-2, ..., 2} / 16: every gate dot product is an + # integer of magnitude at most 128 * 2 = 256 over 64, exact in any fp32 summation order + # and in bf16, so the kernel's gate GEMM and the reference's agree bit for bit and the + # state comparison measures the recurrence alone. + fa = torch.randint(-1, 2, (rows, K), generator=gen) / 4 + bfa[:, :K] = fa.to(torch.bfloat16).cuda() + w_fb = torch.randint(-2, 3, (dim, K), generator=gen) / 16 + args.update(fa=bfa[:, :K].view(1, rows, K), w_fb=w_fb.to(torch.bfloat16).cuda()) + args["use_pdl"] = True + else: + args["g"] = rn(rows, dim, dtype=torch.bfloat16).view(1, rows, heads, K) + if variant == "bf16": + args["use_pdl"] = True + else: + args["quantize_output"] = True + return args + + +def _scale_offsets(rows: int, heads: int) -> torch.Tensor: + """``[rows, heads, 4]`` byte offsets of the R128c4 UE8M0 scales of rows ``[0, rows)``.""" + r = torch.arange(rows, device="cuda")[:, None] + h = torch.arange(heads, device="cuda")[None, :] + base = (r // 128) * heads * 512 + h * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + return base[..., None] + torch.arange(4, device="cuda") + + +def _check_mxfp8(out, ref, rows: int, heads: int, why: str) -> None: + act, scales = out + ref_act, ref_scales = ref + assert act.shape == ref_act.shape and act.dtype == ref_act.dtype, why + offsets = _scale_offsets(rows, heads) + got_sf = scales.view(torch.uint8).reshape(-1)[offsets] + want_sf = ref_scales[offsets] + # The four replicas of each (row, head) scale, and the round-up power of two itself. + assert torch.equal(got_sf, want_sf), f"{why}: UE8M0 scales" + unit = torch.exp2(want_sf[..., 0].float() - 127)[..., None] + got = act.float().view(rows, heads, K) * unit + want = ref_act.float().view(rows, heads, K) * unit + # The kernel's bf16 normalized value can land one bf16 ulp from the reference's (fp32 vs fp64 + # recurrence); re-rounded to E4M3 that is at most one step of the head's top binade, and it + # happens to a small fraction of the elements. + assert ((got - want).abs() <= E4M3_TOP_STEP * unit).all(), f"{why}: E4M3 values" + differing = (act.view(torch.uint8) != ref_act.view(torch.uint8)).float().mean().item() + assert differing <= 1e-2, f"{why}: {differing:.2%} of E4M3 codes differ" + + +def _check_post(args: dict, post: dict, why: str) -> None: + # Raw bf16 conv inputs are moved, never computed: bitwise. + for name in ("cs_q", "cs_k", "cs_v"): + assert torch.equal(args[name], post[name]), f"{why}: {name}" + # The fp32 state after 3xTF32 / fp32 recurrence steps, the fp32 golden's band. + torch.testing.assert_close( + args["recurrent_state"], post["recurrent_state"], rtol=0.0, atol=1e-4, msg=why + ) + # The fp32 log-decay: a few fast-math exp / reciprocal roundings of values in [-5, 0]. + torch.testing.assert_close(args["g_cache"], post["g_cache"], rtol=0.0, atol=2e-5, msg=why) + # bf16 k / v / beta rounded once from values that agree to fp32 accuracy: an element lands at + # most one bf16 ulp away, which is two ULP units of its row's scale. + for name in ("k_cache", "v_cache", "beta_cache"): + assert_within_ulp(args[name], post[name], element_ulp=2.0, rms_ulp=0.25) + + +@pytest.mark.parametrize( + "cell", kda_mtp_decode.CELLS, ids=[c.why[:72] for c in kda_mtp_decode.CELLS] +) +@torch.no_grad() +def test_cells(cell) -> None: + spec = cell.spec + args = _call_args(spec, zlib.crc32(cell.why.encode()) % 2**31) + ref, post = kda_mtp_decode.reference(**args) + with validating(kda_mtp_decode): + out = kda_mtp_decode(**args) + torch.cuda.synchronize() + + rows = spec["requests"] * (1 + spec["num_spec"]) + if spec["variant"] == "mxfp8": + _check_mxfp8(out, ref, rows, spec["heads"], cell.why) + else: + assert out.shape == (1, rows, spec["heads"], K) and out.dtype == torch.bfloat16, cell.why + # Two bf16 roundings: bf16(o) may land one ulp of its own value from the reference's, + # which moves the normalized element by as much, and the output rounds once more. Each + # is at most 2**-7 of the head's largest magnitude, two ULP units of the row scale + # (2.5 measured on GB300). + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=0.25) + _check_post(args, post, cell.why) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_kda_prefill.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_kda_prefill.py new file mode 100644 index 000000000000..8a18be3d34e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/ops/ssm/test_kda_prefill.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the kda_prefill entry: every cell of ``kda_prefill.CELLS``, built as the tp1 KDA +layer calls it (metadata hints, a chunk-padded single sequence, padded pool slots, in-forward +snapshot rows), against the entry's fp64 recurrence; and that recurrence against the family's +``_kda_reference.recurrent_reference``.""" + +import zlib + +import pytest +import torch + +pytest.importorskip("cutlass") + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.ssm.kda_prefill import ( # noqa: E501, E402 + kda_prefill, +) + +__extra_import_path__ = ["../../../../..", "../../../ops/ssm"] +import _kda_reference as kr # noqa: E402 +from _validating import validating # noqa: E402 + +HEAD_DIM, CHUNK = kr.HEAD_DIM, kr.CHUNK +SLOTS = 12 +SLOT_PAD = 4096 # floats past each slot's state: the cache manager's slot stride is wider +# Bands relative to the reference's largest magnitude: bf16 tensor-core chunks vs an fp64 +# token-by-token recurrence (the family prefill entries' bands). +OUT_TOL = 2e-2 +STATE_TOL = 1e-2 + + +def _driven_here() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in kda_prefill.ARCHS} + + +pytestmark = pytest.mark.skipif(not _driven_here(), reason="kda_prefill.ARCHS") + + +def _chunk_indices(seq_lens) -> torch.Tensor: + pairs = [(s, c) for s, n in enumerate(seq_lens) for c in range(-(-n // CHUNK))] + return torch.tensor(pairs, dtype=torch.long, device="cuda").view(-1, 2) + + +def _call_args(spec: dict, seed: int) -> tuple[dict, torch.Tensor, int]: + """The arguments K3KDAKernelDispatch.prefill_chunk_kda passes for this batch, the padded pool + storage behind ``state_pool``, and the real token count.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + heads, seq_lens = spec["heads"], list(spec["seq_lens"]) + tokens = sum(seq_lens) + raw = kr.kimi_k3_inputs(heads, tokens, gen) + pool = torch.randn(SLOTS, heads * HEAD_DIM * HEAD_DIM + SLOT_PAD, generator=gen, device="cuda") + state_pool = pool[:, : heads * HEAD_DIM * HEAD_DIM].view(SLOTS, heads, HEAD_DIM, HEAD_DIM) + state_pool.mul_(0.05) + order = torch.randperm(SLOTS, generator=torch.Generator().manual_seed(seed)) + slots = order[: len(seq_lens)] + for slot, fresh in zip(slots.tolist(), spec["fresh"]): + if fresh: + state_pool[slot].zero_() # the layer resets fresh rows before the op + # The metadata's int32 state indices sit at a 16-byte aligned offset. + state_indices = torch.empty(4 + len(seq_lens), dtype=torch.int32, device="cuda")[4:] + state_indices.copy_(slots) + single = tokens if len(seq_lens) == 1 else None + q, k, v, beta = raw["q"], raw["k"], raw["v"], raw["beta"] + if single is not None and tokens % CHUNK: + # features.kda_prepad_prefill: Q/K/V/beta arrive zero-padded to the chunk, g does not. + pad = CHUNK - tokens % CHUNK + q, k, v = (torch.nn.functional.pad(t, (0, 0, 0, 0, 0, pad)) for t in (q, k, v)) + beta = torch.nn.functional.pad(beta, (0, 0, 0, pad)) + args = dict( + q=q, + k=k, + v=v, + g=raw["g"], + beta=beta, + state_pool=state_pool, + state_indices=state_indices, + scale=kr.SCALE, + cu_seqlens=kr.cu_seqlens_of(seq_lens), + chunk_indices=_chunk_indices(seq_lens), + chunk_size=CHUNK, + safe_gate=True, + lower_bound=kr.LOWER_BOUND, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + A_log=raw["A_log"], + dt_bias=raw["dt_bias"], + varlen_is_aligned=all(n % CHUNK == 0 for n in seq_lens), + single_sequence_length=single, + ) + if spec.get("snapshot_chunks") is not None: + args["snapshot_rows"] = order[len(seq_lens) : 2 * len(seq_lens)].to(torch.int32).cuda() + args["snapshot_chunks"] = torch.tensor( + spec["snapshot_chunks"], dtype=torch.int32, device="cuda" + ) + return args, pool, tokens + + +@pytest.mark.parametrize("cell", kda_prefill.CELLS, ids=[c.why[:60] for c in kda_prefill.CELLS]) +@torch.no_grad() +def test_cells(cell) -> None: + args, pool, tokens = _call_args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + pool_before = pool.clone() + ref_out, ref_pool = kda_prefill.reference(**args) + with validating(kda_prefill): + out = kda_prefill(**args).clone() # a view of the op's shape-keyed scratch buffer + + assert out.shape == args["v"].shape and out.dtype == args["v"].dtype, cell.why + kr.assert_scaled_close(out[:, :tokens], ref_out[:, :tokens], OUT_TOL) + written = set(args["state_indices"].tolist()) + if "snapshot_rows" in args: + written |= { + row + for row, c in zip(args["snapshot_rows"].tolist(), args["snapshot_chunks"].tolist()) + if c >= 0 + } + state_pool = args["state_pool"] + for slot in range(SLOTS): + if slot in written: + kr.assert_scaled_close(state_pool[slot], ref_pool[slot], STATE_TOL) + else: + assert torch.equal(pool[slot], pool_before[slot]), f"{cell.why}: slot {slot} written" + state_size = state_pool[0].numel() + assert torch.equal(pool[:, state_size:], pool_before[:, state_size:]), "slot padding written" + + +@torch.no_grad() +def test_reference_is_the_family_recurrence() -> None: + """The entry's reference on two sequences equals the family's fp64 recurrence.""" + spec = dict(heads=6, seq_lens=(70, 33), fresh=(True, False)) + args, _, _ = _call_args(spec, 5) + ref_out, ref_pool = kda_prefill.reference(**args) + starts = args["cu_seqlens"].tolist() + for b, slot in enumerate(args["state_indices"].tolist()): + lo, hi = starts[b], starts[b + 1] + o, s = kr.recurrent_reference( + args["q"][0, lo:hi], + args["k"][0, lo:hi], + args["v"][0, lo:hi], + args["g"][0, lo:hi], + args["beta"][0, lo:hi], + args["state_pool"][slot], + args["A_log"], + args["dt_bias"], + ) + kr.assert_scaled_close(ref_out[0, lo:hi], o.to(ref_out.dtype), 1e-2) + kr.assert_scaled_close(ref_pool[slot], s.float(), 1e-6) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_dspark_draft_ar.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_dspark_draft_ar.py new file mode 100644 index 000000000000..93328f3bcdb7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_dspark_draft_ar.py @@ -0,0 +1,449 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DSpark drafter all-reduce sites: the pushed all-reduce + residual + +RMSNorm site kernel and the split-K GEMM push epilogue are bitwise the C++ MNNVL fused +all-reduce of the same partials, with C++ calls interleaved on the workspace and under CUDA-graph +replay; the owner-multicast embedding site is bitwise the vocab-sharded lookup + all-reduce. The +Lamport slots come from the Kimi K3 package's pushed all-reduce (``k3_rubin/comm/lamport.py``).""" + +import os +import pickle +import sys +import traceback +from types import SimpleNamespace + +import cloudpickle +import pytest +import torch +from mpi4py import MPI + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__( + cloudpickle.dumps, + cloudpickle.loads, + pickle.HIGHEST_PROTOCOL, +) + +pytestmark = pytest.mark.threadleak(enabled=False) + +HIDDEN = 7168 +EPS = 1e-6 + + +def _input(m, h, seed, rank, special=False): + gen = torch.Generator(device="cuda").manual_seed(seed * 1000 + rank) + x = (torch.randn(m, h, device="cuda", generator=gen) * 0.25).to(torch.bfloat16) + if special: + bits = x.view(torch.int16) + r = torch.rand(m, h, device="cuda", generator=gen) + bits[r < 0.05] = -32768 # -0 + bits[(r >= 0.05) & (r < 0.08)] = 0 + den = (r >= 0.08) & (r < 0.1) + bits[den] = torch.randint(1, 127, (int(den.sum()),), device="cuda", generator=gen).to( + torch.int16 + ) + return x + + +def _differ(ref, out): + return int((ref.view(torch.int16) != out.view(torch.int16)).sum()) + + +def _run_rank(world_size): + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.dflash import ( # noqa: E501 + DraftArSites, + push_tactic, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.dspark_draft_ar import ( # noqa: E501 + ar_norm_config, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.dspark_draft_gemm_push import ( # noqa: E501 + run_gemm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.layers import ( # noqa: E501 # noqa: E501 + TpAllReduce, + lamport_push_target, + lamport_push_tensor, + mnnvl_lamport_push_target, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec.k3_draft_ar_norm import ( # noqa: E501 + k3_draft_ar_norm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec.k3_draft_gemm_push import ( # noqa: E501 + k3_draft_gemm_push, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec.k3_splitk_dense import ( # noqa: E501 + k3_splitk_dense, + ) + from tensorrt_llm.functional import AllReduceFusionOp, AllReduceParams, AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + torch.cuda.set_device(rank) + try: + mapping = Mapping(world_size=world_size, tp_size=world_size, rank=rank) + allreduce = TpAllReduce(mapping, torch.bfloat16, AllReduceStrategy.MNNVL) + assert allreduce.mnnvl_allreduce is not None + gen = torch.Generator(device="cuda").manual_seed(5) + gamma = (1 + 0.1 * torch.randn(HIDDEN, device="cuda", generator=gen)).to(torch.bfloat16) + norm = SimpleNamespace(weight=gamma, variance_epsilon=EPS) + bad = 0 + + def reference(partial, residual): + return allreduce( + partial, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=gamma, + eps=EPS, + ), + ) + + def site(target, residual): + out = torch.empty_like(residual) + res_out = torch.empty_like(residual) + k3_draft_ar_norm( + target.workspace, + target.one_shot, + target.cpp_one_shot, + residual, + gamma, + EPS, + out, + res_out, + ) + return out, res_out + + # Push kernel + site against the C++ fused all-reduce, both transports, C++ calls between. + for m in (1, 7, 8, 14, 16, 21, 28, 32, 56, 64, 112, 128): + if not ar_norm_config(m, HIDDEN, torch.device("cuda", rank))[0]: + continue + for special in (False, True): + x = _input(m, HIDDEN, 7 * m + special, rank, special) + res = _input(m, HIDDEN, 11 * m + special, 0, special) + ref_out, ref_res = reference(x, res) + for one_shot in (True, False): + target = mnnvl_lamport_push_target( + allreduce.mnnvl_allreduce, m, HIDDEN, one_shot=one_shot + ) + lamport_push_tensor(allreduce, target, x) + out, res_out = site(target, res) + bad += _differ(ref_out, out) + _differ(ref_res, res_out) + again_out, again_res = reference(x, res) + bad += _differ(ref_out, again_out) + _differ(ref_res, again_res) + + # Split-K GEMM push epilogue + site against the same GEMM's store-only variant (== the + # vendored split-K kernel up to 32 rows) + C++ fused all-reduce. + for role, k, m in ( + ("o_proj", 1024, 7), + ("o_proj", 1024, 8), + ("o_proj", 1024, 14), + ("down", 1792, 16), + ("down", 1792, 21), + ("o_proj", 1024, 32), + ("o_proj", 1024, 56), + ("down", 1792, 56), + ("down", 1792, 64), + ): + w = (torch.randn(HIDDEN, k, device="cuda", generator=gen) * 0.02).to(torch.bfloat16) + tactic = push_tactic(role, m, HIDDEN, k) + assert tactic is not None, (role, m) + x = _input(m, k, 13 * m + k, rank) + res = _input(m, HIDDEN, 17 * m, 0) + y = run_gemm(x, w, torch.empty(m, HIDDEN, device="cuda", dtype=torch.bfloat16), tactic) + if m <= 32: + y_vendored = torch.empty_like(y) + k3_splitk_dense(x, w.t(), None, y_vendored, True, tactic) + bad += _differ(y_vendored, y) + ref_out, ref_res = reference(y, res) + target = lamport_push_target(allreduce, m, HIDDEN) + k3_draft_gemm_push(x, w, target.output, target, tactic, store_out=True) + bad += _differ(y, target.output) + out, res_out = site(target, res) + bad += _differ(ref_out, out) + _differ(ref_res, res_out) + sites = DraftArSites(allreduce) + out, res_out = sites.linear_ar_norm(x, w, role, res, norm) + bad += _differ(ref_out, out) + _differ(ref_res, res_out) + # The split-K GEMM is a reordered accumulation of the cuBLAS one. + torch.testing.assert_close(y.float(), (x.float() @ w.float().t()), atol=2e-2, rtol=2e-2) + + # One CUDA graph: o_proj site -> C++ fused all-reduce -> down site, replayed on new inputs. + m = 14 + w_o = (torch.randn(HIDDEN, 1024, device="cuda", generator=gen) * 0.02).to(torch.bfloat16) + w_d = (torch.randn(HIDDEN, 1792, device="cuda", generator=gen) * 0.02).to(torch.bfloat16) + x_o = _input(m, 1024, 1, rank) + x_d = _input(m, 1792, 2, rank) + p_c = _input(m, HIDDEN, 3, rank) + res = _input(m, HIDDEN, 4, 0) + sites = DraftArSites(allreduce) + + def chain(): + a, r = sites.linear_ar_norm(x_o, w_o, "o_proj", res, norm) + b, r = reference(p_c, r) + c, r = sites.linear_ar_norm(x_d, w_d, "down", r, norm) + return a, b, c, r + + def expected(): + y_o = run_gemm( + x_o, + w_o, + torch.empty(m, HIDDEN, device="cuda", dtype=torch.bfloat16), + push_tactic("o_proj", m, HIDDEN, 1024), + ) + a, r = reference(y_o, res) + b, r = reference(p_c, r) + y_d = run_gemm( + x_d, + w_d, + torch.empty(m, HIDDEN, device="cuda", dtype=torch.bfloat16), + push_tactic("down", m, HIDDEN, 1792), + ) + c, r = reference(y_d, r) + return a, b, c, r + + chain() + torch.cuda.synchronize() + comm.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = chain() + for rep in range(3): + x_o.copy_(_input(m, 1024, 100 + rep, rank)) + x_d.copy_(_input(m, 1792, 200 + rep, rank)) + p_c.copy_(_input(m, HIDDEN, 300 + rep, rank)) + torch.cuda.synchronize() + comm.Barrier() + graph.replay() + torch.cuda.synchronize() + for got, want in zip(outs, expected()): + bad += _differ(want, got) + del graph + torch.cuda.synchronize() + assert comm.allreduce(bad) == 0, f"rank {rank}: {bad} values differ from the C++ path" + except Exception: + traceback.print_exc() + raise + return True + + +def _run_rank_embed(world_size): + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.dflash import ( # noqa: E501 + draft_embed_site, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.layers import ( # noqa: E501 + EmbeddingTable, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec.k3_vocab_shard_embedding import ( # noqa: E501 + k3_vocab_shard_embedding, + ) + + def vocab_shard_embedding(embedding, ids): + local = k3_vocab_shard_embedding( + ids, embedding.weight, embedding.vocab_start, embedding.vocab_end + ) + return embedding.all_reduce(local) + + from tensorrt_llm.mapping import Mapping + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + torch.cuda.set_device(rank) + try: + mapping = Mapping(world_size=world_size, tp_size=world_size, rank=rank) + vocab = 1000 * world_size - 3 # a shorter last shard + embedding = EmbeddingTable( + vocab, + HIDDEN, + dtype=torch.bfloat16, + mapping=mapping, + vocab_parallel=True, + device="cuda", + ) + gen = torch.Generator(device="cuda").manual_seed(9) + full = torch.randn(vocab, HIDDEN, device="cuda", generator=gen).to(torch.bfloat16) + full.view(torch.int16)[:, ::7] = -32768 # -0 lanes: +0 after the all-reduce + embedding.load(full) + bad = 0 + for n in (1, 2, 9, 17, 129): + for dtype in (torch.int64, torch.int32): + ids = torch.randint(0, vocab, (n,), device="cuda", generator=gen).to(dtype) + ids[0] = vocab - 1 + if n > 2: + ids[1] = -1 # outside the vocab: a zero row + ids[2] = vocab + expected = vocab_shard_embedding(embedding, ids) + out = draft_embed_site(embedding, ids) + assert out is not None + bad += _differ(expected, out) + # CUDA graph replay. + ids = torch.randint(0, vocab, (9,), device="cuda", generator=gen) + draft_embed_site(embedding, ids) + torch.cuda.synchronize() + comm.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = draft_embed_site(embedding, ids) + for rep in range(3): + ids.copy_(torch.randint(0, vocab, (9,), device="cuda", generator=gen)) + torch.cuda.synchronize() + comm.Barrier() + graph.replay() + torch.cuda.synchronize() + bad += _differ(vocab_shard_embedding(embedding, ids), out) + del graph + assert comm.allreduce(bad) == 0, f"rank {rank}: {bad} values differ" + except Exception: + traceback.print_exc() + raise + return True + + +def _run_rank_mlp(world_size): + """The drafter glue: a TP GatedMlp whose down projection is deferred to the pushed site + (``_draft_mlp(defer_down=True)`` + ``_reduce_partial``) against the split-K down GEMM + C++ + fused all-reduce, and the fallback above the split-K rows against the module path.""" + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.dflash import ( # noqa: E501 + DraftArSites, + push_tactic, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.drafter import ( # noqa: E501 + K3MLADSparkForCausalLM, + _draft_mlp, + _PendingDown, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.kernels.dspark_draft_gemm_push import ( # noqa: E501 + run_gemm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.layers import ( # noqa: E501 # noqa: E501 # noqa: E501 + GatedMlp, + RMSNormWeight, + TpAllReduce, + ) + from tensorrt_llm.functional import AllReduceFusionOp, AllReduceParams, AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + torch.cuda.set_device(rank) + try: + mapping = Mapping(world_size=world_size, tp_size=world_size, rank=rank) + fused_ar = TpAllReduce(mapping, torch.bfloat16, AllReduceStrategy.MNNVL) + mlp = GatedMlp( + HIDDEN, + 1792 * world_size, + activation="silu", + dtype=torch.bfloat16, + mapping=mapping, + allreduce_strategy=AllReduceStrategy.MNNVL, + device="cuda", + ) + gen = torch.Generator(device="cuda").manual_seed(40 + rank) + with torch.no_grad(): + mlp.gate_up_proj.weight.copy_( + torch.randn(mlp.gate_up_proj.weight.shape, device="cuda", generator=gen) * 0.02 + ) + mlp.down_proj.weight.copy_( + torch.randn(mlp.down_proj.weight.shape, device="cuda", generator=gen) * 0.02 + ) + norm = RMSNormWeight(HIDDEN, EPS, torch.bfloat16, device="cuda") + stub = SimpleNamespace( + _fused_ar_residual_norm=K3MLADSparkForCausalLM._fused_ar_residual_norm + ) + sites = DraftArSites(fused_ar) + no_reduce = AllReduceParams(enable_allreduce=False) + bad = 0 + with torch.inference_mode(): + for m in (7, 14, 28, 56, 112): + x = _input(m, HIDDEN, 50 + m, 0) + res = _input(m, HIDDEN, 60 + m, 0) + module_partial = _draft_mlp(mlp, x, no_reduce) + ref_module = K3MLADSparkForCausalLM._reduce_partial( + stub, fused_ar, None, module_partial, res, norm + ) + pending = _draft_mlp(mlp, x, no_reduce, defer_down=True) + assert isinstance(pending, _PendingDown) + out = K3MLADSparkForCausalLM._reduce_partial( + stub, fused_ar, sites, pending, res, norm + ) + k = pending.h2.shape[1] + tactic = push_tactic("down", m, HIDDEN, k) + if tactic is None: + want = ref_module # no pushed site above the split-K rows + else: + y = run_gemm( + pending.h2, + mlp.down_proj.weight.detach(), + torch.empty(m, HIDDEN, device="cuda", dtype=torch.bfloat16), + tactic, + ) + want = fused_ar( + y, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=res, + norm_weight=norm.weight, + eps=norm.variance_epsilon, + ), + ) + torch.testing.assert_close( + out[0].float(), ref_module[0].float(), atol=0.1, rtol=0.05 + ) + bad += _differ(want[0], out[0]) + _differ(want[1], out[1]) + assert comm.allreduce(bad) == 0, f"rank {rank}: {bad} values differ" + except Exception: + traceback.print_exc() + raise + return True + + +_needs = pytest.mark.skipif( + torch.cuda.device_count() < 2 + or torch.cuda.get_device_capability()[0] != 10 + or os.environ.get("TLLM_TEST_SKIP_MNNVL") == "1", + reason="needs two SM100-family GPUs with MNNVL", +) +_WORLD_SIZES = [2, 4] if torch.cuda.device_count() >= 4 else [2] + + +@_needs +@pytest.mark.parametrize("mpi_pool_executor", _WORLD_SIZES, indirect=True) +def test_dspark_draft_ar_sites(mpi_pool_executor): + world = mpi_pool_executor.num_workers + for r in mpi_pool_executor.map(_run_rank, *zip(*[(world,)] * world)): + assert r is True + + +@_needs +@pytest.mark.parametrize("mpi_pool_executor", _WORLD_SIZES, indirect=True) +def test_dspark_draft_embed_site(mpi_pool_executor): + world = mpi_pool_executor.num_workers + for r in mpi_pool_executor.map(_run_rank_embed, *zip(*[(world,)] * world)): + assert r is True + + +@_needs +@pytest.mark.parametrize("mpi_pool_executor", _WORLD_SIZES, indirect=True) +def test_dspark_draft_mlp_site(mpi_pool_executor): + world = mpi_pool_executor.num_workers + for r in mpi_pool_executor.map(_run_rank_mlp, *zip(*[(world,)] * world)): + assert r is True diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_dspark_mla_rope.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_dspark_mla_rope.py new file mode 100644 index 000000000000..358aa0b7c1f1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_dspark_mla_rope.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The K3 MLA drafter's RMSNorm + RoPE entries (``ops/spec/k3_dspark_*``) write bitwise what +the launches they replace write: the ``TLLM_DSPARK_TAIL_V2`` kernels against the +per-layer / separate launches, and the row-strided ``k3_dspark_rmsnorm_rope`` against main's +``cute_dsl_dspark_rmsnorm_rope`` on a contiguous copy.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._utils import is_sm_100f + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not is_sm_100f(), reason="needs an SM100-family GPU" +) + +_OPS = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec" + +KV_LORA = 512 +ROPE = 64 +NOPE = 128 +HEAD_DIM = KV_LORA + ROPE +EPS = 1e-5 + + +def _entry(name): + import importlib + + return getattr(importlib.import_module(f"{_OPS}.{name}"), name) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _freqs_table(num_positions: int, gen) -> torch.Tensor: + """A ``[P, ROPE // 2, 2]`` FP32 (cos, sin) table with YaRN-like magnitudes.""" + angle = torch.rand(num_positions, ROPE // 2, device="cuda", generator=gen) * 6.3 + scale = 1.0 + 0.3 * torch.rand(num_positions, 1, device="cuda", generator=gen) + return torch.stack([angle.cos() * scale, angle.sin() * scale], dim=-1).contiguous() + + +def _rows(n, width, gen, row_stride=None): + """BF16 rows with a spread of magnitudes; optionally a row-strided view.""" + row_stride = row_stride or width + base = torch.randn(n, row_stride, device="cuda", generator=gen) + base = base * torch.logspace(-2, 2, n, device="cuda").unsqueeze(1) + return base.to(torch.bfloat16)[:, :width] + + +@pytest.mark.parametrize("num_tokens", [1, 7, 14, 28, 33]) +def test_ctx_norm_rope_matches_per_layer_kernel(num_tokens): + layers = 5 + gen = torch.Generator(device="cuda").manual_seed(num_tokens) + latent = _rows(num_tokens, layers * HEAD_DIM, gen).view(num_tokens, layers, HEAD_DIM) + weight = (1.0 + 0.5 * torch.randn(layers, KV_LORA, device="cuda", generator=gen)).to( + torch.bfloat16 + ) + table = _freqs_table(4096, gen) + positions = torch.randint(0, 4096, (num_tokens,), device="cuda", generator=gen) + + out = _entry("k3_dspark_mla_ctx_norm_rope")( + latent, torch.nn.Parameter(weight), positions, table, EPS + ) + + import importlib + + into = importlib.import_module(f"{_OPS}.k3_dspark_rmsnorm_rope_into") + compiled = into.build_rmsnorm_rope_into(HEAD_DIM, ROPE, 1, EPS, True, True, False, KV_LORA) + gathered = table[positions].contiguous() + ref = torch.empty_like(out) + for layer in range(layers): + into.k3_dspark_rmsnorm_rope_into( + latent[:, layer], weight[layer], gathered, ref[layer], compiled + ) + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(ref)) + + +@pytest.mark.parametrize("num_tokens", [7, 14, 28, 32]) +@pytest.mark.parametrize("page_size", [32, 64]) +def test_block_rope_matches_separate_launches(num_tokens, page_size): + from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_rope_into, + ) + + scatter_rows_to_pool = _entry("k3_scatter_rows_to_pool") + + heads = 4 + gen = torch.Generator(device="cuda").manual_seed(100 * num_tokens + page_size) + q = _rows(num_tokens * heads, NOPE + ROPE, gen).view(num_tokens, heads, NOPE + ROPE) + # The block latents are the kv half of the fused [q_a | kv_a] GEMM output. + kv = _rows(num_tokens, HEAD_DIM, gen, row_stride=1536 + HEAD_DIM) + weight = (1.0 + 0.5 * torch.randn(KV_LORA, device="cuda", generator=gen)).to(torch.bfloat16) + table = _freqs_table(8192, gen) + positions = torch.randint(0, 8192, (num_tokens,), device="cuda", generator=gen) + # An interleaved multi-layer page space, like the draft KV manager's view. + num_pages = 40 + space = torch.randn(num_pages, 3, page_size, HEAD_DIM, device="cuda", generator=gen) + pool = space.to(torch.bfloat16)[:, 1] + pages = torch.randperm(num_pages, device="cuda", generator=gen)[:num_tokens].long() + offsets = torch.randint(0, page_size, (num_tokens,), device="cuda", generator=gen) + + query = torch.full((num_tokens, heads, HEAD_DIM), 7.0, dtype=torch.bfloat16, device="cuda") + pool_new = pool.clone() + _entry("k3_dspark_mla_block_rope")( + q, query, kv, torch.nn.Parameter(weight), positions, table, pool_new, pages, offsets, EPS + ) + + gathered = table[positions].contiguous() + query_ref = torch.full_like(query, 7.0) + cute_dsl_dspark_rope_into(q, gathered, query_ref, heads, ROPE, KV_LORA) + blk_kv = _entry("k3_dspark_rmsnorm_rope")( + kv, weight, gathered, 1, ROPE, EPS, True, True, False, KV_LORA + ) + pool_ref = pool.clone() + scatter_rows_to_pool(blk_kv, pool_ref, pages, offsets) + torch.cuda.synchronize() + assert torch.equal(_bits(query), _bits(query_ref)) + assert torch.equal(_bits(pool_new), _bits(pool_ref)) + + +@pytest.mark.parametrize("num_rows", [1, 7, 33]) +@pytest.mark.parametrize("row_stride", [None, 1536 + HEAD_DIM]) +@pytest.mark.parametrize( + "norm_dim,apply_weight", [(KV_LORA, True), (None, True), (None, False)], ids=str +) +def test_strided_rmsnorm_rope_matches_main_op(num_rows, row_stride, norm_dim, apply_weight): + """``k3_dspark_rmsnorm_rope`` reads a row-strided view in place; main's op, compiled for + compact rows, gets a contiguous copy of the same rows. Same bits.""" + from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_rmsnorm_rope, + ) + + gen = torch.Generator(device="cuda").manual_seed(num_rows * 11 + (row_stride or 0)) + x = _rows(num_rows, HEAD_DIM, gen, row_stride=row_stride) + weight_len = HEAD_DIM if norm_dim is None else norm_dim + weight = (1.0 + 0.5 * torch.randn(weight_len, device="cuda", generator=gen)).to(torch.bfloat16) + table = _freqs_table(num_rows, gen) + args = (ROPE, EPS, apply_weight, True, False, norm_dim) + out = _entry("k3_dspark_rmsnorm_rope")(x, weight, table, 1, *args) + ref = cute_dsl_dspark_rmsnorm_rope(x.contiguous(), weight, table, 1, *args) + torch.cuda.synchronize() + assert out.is_contiguous() and out.shape == x.shape + assert torch.equal(_bits(out), _bits(ref)) + + +# The certified cells of the four entries, each against the entry's own FP64 reference. + +pytest.importorskip("cutlass") + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec import ( # noqa: E402, E501 + k3_dspark_mla_block_rope as _block_rope_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec import ( # noqa: E402, E501 + k3_dspark_mla_ctx_norm_rope as _ctx_norm_rope_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec import ( # noqa: E402, E501 + k3_dspark_rmsnorm_rope as _rmsnorm_rope_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec import ( # noqa: E402, E501 + k3_dspark_rmsnorm_rope_into as _rmsnorm_rope_into_entry, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402, F401 — path above; fixture by name + assert_within_ulp, + certifying, +) + +rmsnorm_rope_op = _rmsnorm_rope_entry.k3_dspark_rmsnorm_rope +rmsnorm_rope_into_op = _rmsnorm_rope_into_entry.k3_dspark_rmsnorm_rope_into +ctx_norm_rope_op = _ctx_norm_rope_entry.k3_dspark_mla_ctx_norm_rope +block_rope_op = _block_rope_entry.k3_dspark_mla_block_rope + + +def _on_archs(op) -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" in {arch.value for arch in op.ARCHS} + + +def _cells(op): + return pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) + + +def _cell_gen(cell) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """2 ulp per element against the row scale, 0.5 ulp relative RMS. + + One BF16 rounding on each side of an FP32 (kernel) against an FP64 (reference) RMSNorm and + rotation. The FP32 sum of squares over 512 columns and the approximate reciprocal square root + move a normalized value by a few FP32 ulps, the rotation's products by one; so the two differ + only where that value sits that close to a BF16 rounding midpoint, by one BF16 step of that + element: at most 2 units of ``ULP[bf16]`` against the row scale. Those flips are rare enough + that the relative RMS stays far below half an ulp. A wrong pair order, sign or position misses + by O(1). + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.5) + + +def _norm_weight(n, gen, layers=None): + shape = (n,) if layers is None else (layers, n) + return (1.0 + 0.5 * torch.randn(*shape, device="cuda", generator=gen)).to(torch.bfloat16) + + +@pytest.mark.skipif(not _on_archs(rmsnorm_rope_op), reason="not an ARCHS architecture") +@_cells(rmsnorm_rope_op) +def test_rmsnorm_rope_certified_cells(cell, certifying): # noqa: F811 — fixture by name + certifying(rmsnorm_rope_op) + gen = _cell_gen(cell) + rows = cell.spec["rows"] + x = _rows(rows, HEAD_DIM, gen, row_stride=cell.spec["row_stride"]) + weight = _norm_weight(KV_LORA, gen) + table = _freqs_table(rows, gen) + before = x.clone() + args = (weight, table, 1, ROPE, EPS, True, True, False, KV_LORA) + out = rmsnorm_rope_op(x, *args) + assert out.is_contiguous() and out.shape == x.shape and out.dtype == torch.bfloat16 + assert torch.equal(_bits(x), _bits(before)), "the op must not write its input" + _compare(out, rmsnorm_rope_op.reference(x, *args)) + + +@pytest.mark.skipif(not _on_archs(rmsnorm_rope_into_op), reason="not an ARCHS architecture") +@_cells(rmsnorm_rope_into_op) +def test_rmsnorm_rope_into_certified_cells(cell, certifying): # noqa: F811 — fixture by name + certifying(rmsnorm_rope_into_op) + gen = _cell_gen(cell) + rows, layers = cell.spec["rows"], cell.spec["layers"] + latent = _rows(rows, layers * HEAD_DIM, gen).view(rows, layers, HEAD_DIM) + weight = _norm_weight(KV_LORA, gen, layers) + table = _freqs_table(rows, gen) + compiled = _rmsnorm_rope_into_entry.build_rmsnorm_rope_into( + HEAD_DIM, ROPE, 1, EPS, True, True, False, KV_LORA + ) + before = latent.clone() + out = torch.full((layers, rows, HEAD_DIM), 3.0, dtype=torch.bfloat16, device="cuda") + for layer in range(layers): + # Each layer's slice of the GEMM output is read in place, as the drafter does. + assert ( + rmsnorm_rope_into_op( + latent[:, layer], weight[layer], table, out[layer], compiled + ).data_ptr() + == out[layer].data_ptr() + ) + assert torch.equal(_bits(latent), _bits(before)), "the op must not write its input" + for layer in range(layers): + _compare( + out[layer], + rmsnorm_rope_into_op.reference( + latent[:, layer], weight[layer], table, out[layer], compiled + ), + ) + + +@pytest.mark.skipif(not _on_archs(ctx_norm_rope_op), reason="not an ARCHS architecture") +@_cells(ctx_norm_rope_op) +def test_ctx_norm_rope_certified_cells(cell, certifying): # noqa: F811 — fixture by name + certifying(ctx_norm_rope_op) + gen = _cell_gen(cell) + rows, layers = cell.spec["rows"], cell.spec["layers"] + latent = _rows(rows, layers * HEAD_DIM, gen).view(rows, layers, HEAD_DIM) + weight = torch.nn.Parameter(_norm_weight(KV_LORA, gen, layers), requires_grad=False) + table = _freqs_table(16384, gen) + positions = torch.randint(0, 16384, (rows,), device="cuda", generator=gen) + before = latent.clone() + out = ctx_norm_rope_op(latent, weight, positions, table, EPS) + assert out.shape == (layers, rows, HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.equal(_bits(latent), _bits(before)), "the op must not write its input" + _compare(out, ctx_norm_rope_op.reference(latent, weight, positions, table, EPS)) + + +@pytest.mark.skipif(not _on_archs(block_rope_op), reason="not an ARCHS architecture") +@_cells(block_rope_op) +def test_block_rope_certified_cells(cell, certifying): # noqa: F811 — fixture by name + certifying(block_rope_op) + gen = _cell_gen(cell) + tokens, heads, page_size, num_pages = cell.spec["tokens"], 64, 128, 24 + q = _rows(tokens * heads, NOPE + ROPE, gen).view(tokens, heads, NOPE + ROPE) + kv_row_stride = cell.spec["kv_row_stride"] + kv = _rows( + tokens, HEAD_DIM, gen, row_stride=kv_row_stride if kv_row_stride != HEAD_DIM else None + ) + weight = torch.nn.Parameter(_norm_weight(KV_LORA, gen), requires_grad=False) + table = _freqs_table(16384, gen) + positions = torch.randint(0, 16384, (tokens,), device="cuda", generator=gen) + # One layer's view of an interleaved multi-layer page space; distinct (page, offset) slots. + space = torch.randn(num_pages, 3, page_size, HEAD_DIM, device="cuda", generator=gen) + pool = space.to(torch.bfloat16)[:, 1] + slots = torch.randperm(num_pages * page_size, device="cuda", generator=gen)[:tokens] + pages, offsets = slots // page_size, slots % page_size + query = torch.full((tokens, heads, HEAD_DIM), 7.0, dtype=torch.bfloat16, device="cuda") + ref_query, ref_pool = block_rope_op.reference( + q, query, kv, weight, positions, table, pool, pages, offsets, EPS + ) + before_q, before_kv = q.clone(), kv.clone() + block_rope_op(q, query, kv, weight, positions, table, pool, pages, offsets, EPS) + assert torch.equal(_bits(q), _bits(before_q)) and torch.equal(_bits(kv), _bits(before_kv)) + # Columns and slots the op does not write keep their bits. + assert torch.equal(_bits(query[..., :KV_LORA]), _bits(ref_query[..., :KV_LORA])) + written = torch.zeros(num_pages, page_size, dtype=torch.bool, device="cuda") + written[pages, offsets] = True + assert torch.equal(_bits(pool[~written]), _bits(ref_pool[~written])) + _compare(query[..., KV_LORA:], ref_query[..., KV_LORA:]) + _compare(pool[pages, offsets], ref_pool[pages, offsets]) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_fused_select_accept_kda_replay.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_fused_select_accept_kda_replay.py new file mode 100644 index 000000000000..7519763d49b9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_fused_select_accept_kda_replay.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``fused_select_accept`` with the KDA replay record, and the DSpark worker's gate for it. + +The kernel's replay write must equal ``KDAReplayState.record_acceptance`` on main's replay +count, which is the first int32 of each slot's 16-byte record (a strided 1-D view). +``K3DFlashWorker._kda_replay_operands`` hands the kernel those operands only when +``update_mamba_states`` is ``KimiK3HybridCacheManager``'s replay record and nothing else. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.dflash import ( + K3DFlashWorker, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.ops.spec.k3_fused_select_accept import ( # noqa: E501 + k3_fused_select_accept, +) +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAIntermediateState, + KDAReplayState, + KimiK3HybridCacheManager, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( + MambaHybridCacheManagerV2, +) + +_RECORD_INT32S = 4 # 16-byte replay-count record per slot + + +def _pairs(num_rows, n_slots, vocab, device, generator): + """Interleaved (value, global index) candidates with distinct values per row.""" + vals = torch.rand((num_rows, n_slots), generator=generator, device=device) + idx = torch.randint(0, vocab, (num_rows, n_slots), generator=generator, device=device) + buf = torch.stack((vals, idx.to(torch.float32)), dim=-1).reshape(num_rows, 2 * n_slots) + selected = idx.gather(1, vals.argmax(dim=1, keepdim=True)).squeeze(1).to(torch.int32) + return (buf.contiguous(), buf.stride(0), 2, n_slots), selected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("num_ctx,num_gens,draft_len", [(0, 5, 3), (2, 7, 5), (1, 1, 1)]) +def test_replay_record_matches_record_acceptance(num_ctx, num_gens, draft_len): + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(1234 + num_ctx * 10 + num_gens) + width = draft_len + 1 + batch = num_ctx + num_gens + pairs, selected = _pairs(num_ctx + num_gens * width, 4, 1000, device, gen) + + # Drafts agree with the target's selection on a per-request prefix. + target = selected[num_ctx:].view(num_gens, width) + agree = torch.randint(0, draft_len + 1, (num_gens,), generator=gen, device=device) + cols = torch.arange(draft_len, device=device) + draft = torch.where(cols[None, :] < agree[:, None], target[:, :draft_len], -1).to(torch.int32) + expected_num_accepted = 1 + agree.to(torch.int32) + + capacity = 3 * batch + 2 + records = torch.randint(0, 9, (capacity, _RECORD_INT32S), dtype=torch.int32, device=device) + perm = torch.randperm(capacity, generator=gen, device=device) + state_indices = perm[:batch].to(torch.int32).contiguous() + dummy = torch.zeros(batch, dtype=torch.bool, device=device) + dummy[num_ctx] = True # first generation request is a dummy: its slot keeps its value + kv_lens = torch.randint(10, 50, (batch,), dtype=torch.int32, device=device) + + reference = KDAReplayState(draft_len) + ref_records = records.clone() + reference.prev_num_accepted_tokens = ref_records[:, 0] + reference.record_acceptance(state_indices[num_ctx:], expected_num_accepted - 1, dummy[num_ctx:]) + + fused_records = records.clone() + prev = fused_records[:, 0] + assert prev.stride(0) == _RECORD_INT32S + fused_kv_lens = kv_lens.clone() + accepted, num_accepted, rewind = k3_fused_select_accept( + pairs, + draft, + num_ctx, + batch, + kda_replay=(prev, state_indices, dummy), + kv_lens=fused_kv_lens, + ) + torch.cuda.synchronize() + + assert torch.equal(num_accepted[num_ctx:], expected_num_accepted) + assert torch.equal( + num_accepted[:num_ctx], torch.ones(num_ctx, dtype=torch.int32, device=device) + ) + assert torch.equal(accepted[num_ctx:], target) + assert torch.equal(accepted[:num_ctx, 0], selected[:num_ctx]) + # The replay count matches main's eager record and the rest of each record is untouched. + assert torch.equal(fused_records, ref_records) + assert torch.equal(rewind, 1 - expected_num_accepted) + assert torch.equal(fused_kv_lens[:num_ctx], kv_lens[:num_ctx]) + assert torch.equal(fused_kv_lens[num_ctx:], kv_lens[num_ctx:] + 1) + + +def _k3_manager(*, state, state_indices=None, context_only=False, mamba_layers=2): + mgr = object.__new__(KimiK3HybridCacheManager) + mgr.kda_context_only = context_only + mgr.local_num_mamba_layers = mamba_layers + mgr._speculative_state = state + mgr._dummy_request_mask_buffer = SimpleNamespace(tensor=torch.zeros(8, dtype=torch.bool)) + mgr.get_state_indices = lambda: state_indices + return mgr + + +def _replay_state(): + state = KDAReplayState(3) + state.prev_num_accepted_tokens = torch.zeros((16, _RECORD_INT32S), dtype=torch.int32)[:, 0] + return state + + +def _metadata(mgr, state_indices): + return SimpleNamespace( + kv_cache_manager=mgr, mamba_metadata=SimpleNamespace(state_indices=state_indices) + ) + + +def test_gate_hands_over_the_replay_record(): + state = _replay_state() + indices = torch.arange(8, dtype=torch.int32) + mgr = _k3_manager(state=state) + fusable, operands = K3DFlashWorker._kda_replay_operands(_metadata(mgr, indices)) + assert fusable + prev, state_indices, dummy = operands + assert prev is state.prev_num_accepted_tokens + assert state_indices is indices + assert dummy is mgr._dummy_request_mask + + +def test_gate_falls_back_to_the_manager_state_indices(): + indices = torch.arange(8, dtype=torch.int32) + mgr = _k3_manager(state=_replay_state(), state_indices=indices) + fusable, operands = K3DFlashWorker._kda_replay_operands(_metadata(mgr, None)) + assert fusable and operands[1] is indices + + +@pytest.mark.parametrize("context_only,mamba_layers", [(True, 2), (False, 0)]) +def test_gate_fuses_with_nothing_to_record(context_only, mamba_layers): + mgr = _k3_manager(state=_replay_state(), context_only=context_only, mamba_layers=mamba_layers) + assert K3DFlashWorker._kda_replay_operands(_metadata(mgr, None)) == (True, None) + + +def test_gate_refuses_without_a_replay_record(): + mgr = _k3_manager(state=KDAIntermediateState()) + indices = torch.arange(8, dtype=torch.int32) + assert K3DFlashWorker._kda_replay_operands(_metadata(mgr, indices)) == (False, None) + + +def test_gate_refuses_strided_state_indices(): + mgr = _k3_manager(state=_replay_state()) + indices = torch.arange(16, dtype=torch.int32)[::2] + assert K3DFlashWorker._kda_replay_operands(_metadata(mgr, indices)) == (False, None) + + +def test_gate_refuses_other_managers(): + other = object.__new__(MambaHybridCacheManagerV2) + assert K3DFlashWorker._kda_replay_operands(_metadata(other, None)) == (False, None) + + class _Overrides(KimiK3HybridCacheManager): + def update_mamba_states(self, *args, **kwargs): + pass + + mgr = object.__new__(_Overrides) + assert K3DFlashWorker._kda_replay_operands(_metadata(mgr, None)) == (False, None) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_drafter_slot_no_sync.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_drafter_slot_no_sync.py new file mode 100644 index 000000000000..805477bd351d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_drafter_slot_no_sync.py @@ -0,0 +1,278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Kimi K3 DSpark drafter's slot bookkeeping on the unified draft cache must not wait for the GPU. + +``DFlashSpecMetadata.prepare`` (main's) runs before every forward of a disaggregated generation +server, while the previous decode step is still in flight: it releases the slots of requests that +left the batch and assigns the slots of requests that joined, seeding a joining request's context +length at its first position (its prompt's drafter context arrived in the unified cache's draft +layers). Every write there has to stay stream-ordered (pinned uploads), or the host waits for the +step in flight. + +``K3DFlashWorker._store_prefill_context`` runs right after the target forward of a context server +has been enqueued. On the unified draft cache it bounds each chunk by the row's page count, which +``_refresh_ctx_block_tables`` read from the manager's caches on the host; reading the counts back +from the device would make the host wait for that whole forward before the KV of the step can be +handed to the transceiver. +""" + +import collections +import contextlib +from functools import partial +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._device_copy import DeviceCopyBuffer +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.dflash import ( + K3DFlashWorker, +) +from tensorrt_llm._torch.speculative import DFlashSpecMetadata + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +FLOOR = 1 << 40 # request ids from here on are CUDA-graph padding dummies +NUM_SLOTS = 8 +TOKENS_PER_GEN = 3 # K + 1 target tokens of a generation request + + +@contextlib.contextmanager +def _no_host_sync(): + """Fail on any host-device synchronization inside the block.""" + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + yield + finally: + torch.cuda.set_sync_debug_mode("default") + + +def _bind(worker, *names): + for name in names: + setattr(worker, name, partial(getattr(K3DFlashWorker, name), worker)) + + +def _worker(): + """A worker holding only the slot state main's ``DFlashSpecMetadata.prepare`` touches.""" + lengths = list(range(10, 10 + NUM_SLOTS + 1)) + worker = SimpleNamespace( + _ctx_buf_inited=True, + _ctx_unified_cache=True, + _max_ctx=8192, + _req_to_slot={}, + _req_ctx_pos={}, + _free_slots=collections.deque(range(NUM_SLOTS)), + _graph_dummy_id_floor=FLOOR, + _dummy_slot=NUM_SLOTS, + _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"), + _ctx_len_host=list(lengths), + _batch_to_slot_buffer=DeviceCopyBuffer( + torch.zeros(NUM_SLOTS, dtype=torch.long, device="cuda"), pin_memory=True + ), + ) + _bind( + worker, "_assign_slot", "_seed_generation_slot", "_write_ctx_len", "_rebuild_batch_to_slot" + ) + return worker + + +def _prepare(worker, request_ids, first_positions): + """Main's ``DFlashSpecMetadata.prepare`` for a generation-only batch.""" + metadata = SimpleNamespace( + request_ids=request_ids, + num_generations=len(request_ids), + seq_lens=[TOKENS_PER_GEN] * len(request_ids), + host_position_ids=[p + j for p in first_positions for j in range(TOKENS_PER_GEN)], + _batch_indices_buffer=DeviceCopyBuffer( + torch.zeros(NUM_SLOTS, dtype=torch.int, device="cuda"), pin_memory=True + ), + _dflash_worker=worker, + ) + DFlashSpecMetadata.prepare(metadata) + + +def test_join_and_leave_on_the_unified_draft_cache_are_stream_ordered(): + worker = _worker() + with _no_host_sync(): + _prepare(worker, [100, 101], [40, 7]) + # Request 100 leaves and request 102 joins; 101 keeps its slot and its length. + _prepare(worker, [101, 102], [8, 20]) + torch.cuda.synchronize() + + assert worker._req_to_slot == {101: 1, 102: 2} + assert list(worker._free_slots) == [3, 4, 5, 6, 7, 0] + expected_lengths = list(range(10, 10 + NUM_SLOTS + 1)) + expected_lengths[0], expected_lengths[1], expected_lengths[2] = 0, 7, 20 + assert worker._ctx_len.tolist() == expected_lengths + assert worker._ctx_len_host == expected_lengths + assert worker._batch_to_slot_buffer.tensor[:2].tolist() == [1, 2] + + +def test_slot_reset_on_the_unified_draft_cache_is_stream_ordered(): + worker = _worker() + worker._req_to_slot = {7: 4} + worker._free_slots.remove(4) + with _no_host_sync(): + slot = worker._assign_slot(7, reset=True) + torch.cuda.synchronize() + assert slot == 0 + assert worker._ctx_len[4].item() == 0 and worker._ctx_len[0].item() == 0 + + +PAGE = 32 +HIDDEN = 16 + + +def _store_worker(host_counts): + """A worker that also holds what ``_store_prefill_context`` touches, on the unified cache.""" + worker = _worker() + stored = [] + worker.__dict__.update( + _helix_mapping=None, + _ctx_page_size=PAGE, + _ctx_block_tables=torch.zeros(NUM_SLOTS, 4, dtype=torch.int32, device="cuda"), + _ctx_block_counts=torch.tensor([1, 2] + [0] * (NUM_SLOTS - 2), device="cuda"), + _ctx_block_counts_host=host_counts, + _ctx_reuse_addressable=True, + _ctx_paged=True, + _ctx_kv_buf=[torch.empty(1, dtype=torch.bfloat16, device="cuda")], + _ctx_k_buf=None, + _store_context_kv_paged=lambda k, v, rows, positions: stored.append((k, rows, positions)), + ) + _bind(worker, "_helix_shard", "_store_prefill_context") + return worker, stored + + +def _store(worker, seq_lens): + num_tokens = sum(seq_lens) + hidden = torch.randn(num_tokens, HIDDEN, dtype=torch.bfloat16, device="cuda") + draft_model = SimpleNamespace( + fc=None, + hidden_norm=None, + project_target_hidden=lambda captured: captured, + precompute_context_kv=lambda proj, positions: (proj, None), + ) + spec_metadata = SimpleNamespace( + get_hidden_states=lambda total: hidden, + # Each prompt starts at position 0: no reused prefix. + host_position_ids=[p for n in seq_lens for p in range(n)], + request_ids=[100, 101], + ) + attn_metadata = SimpleNamespace( + num_ctx_tokens=num_tokens, + num_contexts=len(seq_lens), + _seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + ) + position_ids = torch.arange(num_tokens, device="cuda") + worker._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, num_tokens + ) + + +def test_prefill_store_on_the_unified_draft_cache_reads_page_counts_on_the_host(): + worker, stored = _store_worker(host_counts=[1, 2]) + with _no_host_sync(): + # Row 1 holds 2 pages, so 40 tokens fit; row 0 holds 1 page, so 33 overflow. + _store(worker, [33, 40]) + torch.cuda.synchronize() + + # The overflowing request drafts nothing: its slot goes back, emptied. + assert 100 not in worker._req_to_slot + slot1 = worker._req_to_slot[101] + assert worker._ctx_len_host[0] == 0 and worker._ctx_len_host[slot1] == 40 + assert worker._ctx_len[slot1].item() == 40 + assert len(stored) == 1 and stored[0][1].tolist() == [1] * 40 + assert stored[0][2].tolist() == list(range(40)) + + +def test_prefill_store_without_host_page_counts_waits_for_the_device(): + """Control: the device read-back that the host counts replace does synchronize.""" + worker, _ = _store_worker(host_counts=None) + with pytest.raises(RuntimeError, match="synchroniz"), _no_host_sync(): + _store(worker, [33, 40]) + + +OFFSET = 400 # the draft layers' K-buffer offset in a per-layer page table entry +WIDTH = 4 + + +def _refresh_worker(num_blocks, checks): + """A worker with what ``_refresh_ctx_block_tables`` touches, on the unified cache.""" + worker = SimpleNamespace( + _ctx_block_tables=torch.zeros(NUM_SLOTS + 1, WIDTH, dtype=torch.int32, device="cuda"), + _ctx_block_counts=torch.zeros(NUM_SLOTS + 1, dtype=torch.long, device="cuda"), + _ctx_block_partial=torch.zeros(NUM_SLOTS + 1, dtype=torch.int32, device="cuda"), + _ctx_block_counts_host=None, + _ctx_unified_cache=True, + _ctx_unified_host_counts=True, + _draft_table_checks_left=checks, + _ctx_pool_idx=0, + _ctx_block_divisor=1, + _ctx_block_offset=OFFSET, + _ctx_kv_manager=SimpleNamespace( + kv_cache_map={ + request_id: SimpleNamespace(num_blocks=count) + for request_id, count in num_blocks.items() + } + ), + ) + _bind(worker, "_refresh_ctx_block_tables", "_check_host_block_counts") + return worker + + +def _offsets(rows): + """``kv_cache_block_offsets`` ``[pools, rows, K/V, blocks]``; 0 where a row holds no page, + as the native backend writes an unallocated entry.""" + src = torch.zeros(1, NUM_SLOTS + 1, 2, WIDTH, dtype=torch.int32) + for i, blocks in enumerate(rows): + for j, block in enumerate(blocks): + src[0, i, :, j] = OFFSET + block + return src.cuda() + + +def test_block_tables_take_the_page_counts_from_the_manager_on_the_host(): + worker = _refresh_worker({100: 1, 101: 2}, checks=1) + attn_metadata = SimpleNamespace(num_contexts=2, kv_cache_block_offsets=_offsets([[5], [7, 9]])) + # The first iteration checks the host counts against the decoded table (it reads it back). + assert worker._refresh_ctx_block_tables(attn_metadata, 2, [100, 101]) + assert worker._ctx_block_counts_host == [1, 2] + with _no_host_sync(): + assert worker._refresh_ctx_block_tables(attn_metadata, 2, [100, 101]) + torch.cuda.synchronize() + assert worker._ctx_block_counts_host == [1, 2] + assert worker._ctx_block_tables[:2, :2].tolist() == [[5, 0], [7, 9]] + assert worker._ctx_block_counts[:2].tolist() == [1, 2] + + +def test_a_saturated_decoded_count_bounds_nothing_and_is_accepted(): + worker = _refresh_worker({100: 1, 101: 2}, checks=1) + # A zero layer offset: unallocated entries decode as page 0 and count as allocated. + worker._ctx_block_offset = 0 + src = torch.zeros(1, NUM_SLOTS + 1, 2, WIDTH, dtype=torch.int32, device="cuda") + attn_metadata = SimpleNamespace(num_contexts=2, kv_cache_block_offsets=src) + assert worker._refresh_ctx_block_tables(attn_metadata, 2, [100, 101]) + assert worker._ctx_block_counts[:2].tolist() == [WIDTH, WIDTH] + assert worker._ctx_block_counts_host == [1, 2] + + +def test_host_counts_that_disagree_with_the_table_are_dropped(): + worker = _refresh_worker({100: 3, 101: 2}, checks=1) + attn_metadata = SimpleNamespace(num_contexts=2, kv_cache_block_offsets=_offsets([[5], [7, 9]])) + assert worker._refresh_ctx_block_tables(attn_metadata, 2, [100, 101]) + assert worker._ctx_block_counts_host is None + assert worker._ctx_unified_host_counts is False + assert worker._refresh_ctx_block_tables(attn_metadata, 2, [100, 101]) + assert worker._ctx_block_counts_host is None diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_drafter_split_plan.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_drafter_split_plan.py new file mode 100644 index 000000000000..421183211b20 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_drafter_split_plan.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The DSpark drafter's split-KV plan (``_launch_mla_split_plan``) is main's planner kernel on the +CuTe DSL MLA decode's geometry at the drafter's piece cost. At main's piece cost it must write +the plan main's own launcher (``cute_dsl_mla.launch_mla_split_kv_plan``) writes.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.fmha import cute_dsl_mla + + +def _runner(): + try: + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLNVMlaDecodeBlackwellRunner, + ) + except ImportError: + return None + return CuteDSLNVMlaDecodeBlackwellRunner + + +requires_blackwell_mla_decode = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] != 10 + or _runner() is None, + reason="needs an SM100-family GPU and the CuTe DSL MLA decode runner", +) + + +@requires_blackwell_mla_decode +@pytest.mark.parametrize("num_heads, seq_len_q", [(64, 8), (64, 6), (32, 8), (128, 1)]) +@pytest.mark.parametrize("num_requests", [1, 3, 17, 64, 200]) +def test_drafter_plan_matches_mains_launcher_at_mains_cost(num_heads, seq_len_q, num_requests): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + drafter, + ) + + runner = _runner() + # The CuTe DSL occupancy query runs on the current CUDA context, which the first + # allocation creates (a serving drafter's weights are on the device by then). + torch.empty(1, device="cuda") + rows, waves = drafter._mla_split_plan_geometry(num_heads, seq_len_q) + assert rows == runner.get_num_q_tiles(num_heads, seq_len_q) + assert waves == max(runner._get_max_active_blocks() // runner._CLUSTER_SHAPE_MNK[0], 1) + + gen = torch.Generator(device="cuda").manual_seed(num_requests * 13 + num_heads + seq_len_q) + seq_lens = torch.randint( + 1, 200_000, (num_requests,), device="cuda", generator=gen, dtype=torch.int32 + ) + expected = torch.full((num_requests + 5,), -1, dtype=torch.int32, device="cuda") + cute_dsl_mla.launch_mla_split_kv_plan( + seq_lens, expected, num_requests=num_requests, rows_per_request=rows, wave_clusters=waves + ) + got = torch.full_like(expected, -1) + drafter._launch_mla_split_plan( + seq_lens, + got, + num_heads=num_heads, + seq_len_q=seq_len_q, + piece_cost=cute_dsl_mla._SPLIT_PLAN_PIECE_COST_TILES, + ) + torch.cuda.synchronize() + assert torch.equal(got, expected) + written = got[:num_requests] + assert ((written >= 1) & (written <= cute_dsl_mla._SPLIT_PLAN_MAX_SPLIT)).all() diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_unified_draft_binding.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_unified_draft_binding.py new file mode 100644 index 000000000000..41ef927cf3ee --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/spec/test_k3_unified_draft_binding.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Kimi K3 DSpark worker binds the drafter layers KV cache manager V2 appends to the target's. + +With ``TRTLLM_DISAGG_SHARED_DRAFT_KV`` the engine appends the drafter's context KV layers to +``KimiK3HybridCacheManager`` after the target's (``_appended_draft_layer_specs``) and hands the +worker the target manager itself as the draft cache (``get_draft_kv_cache_view``). The worker +then has to find those layers by their global ids (``draft_layer_ids``), take the per-layer +pool views, and decode the target metadata's ``kv_cache_block_offsets`` with the draft layers' +own pool index and K-buffer offset. + +The manager below keeps its real layer bookkeeping, so ``draft_layer_ids``, +``get_layer_kv_factor`` and ``get_draft_kv_cache_view`` are the manager's own; only the native +pool (``get_buffers`` and the page-index converter) is replaced by CPU tensors laid out like V2's +per-layer views (``stride(0)`` is one layer slot). +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.dflash import ( + K3DFlashWorker, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.drafter import ( # noqa: E501 + K3MLADSparkForCausalLM, +) +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import KimiK3HybridCacheManager +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + DataType, + KVCacheLayerOwner, + KVCacheLayerSpec, + Role, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +# The engine appends the drafter's layers to the target manager only on a main with the unified +# draft KV; without it the drafter keeps its own manager and these bindings never happen. +requires_appended_draft_layers = pytest.mark.skipif( + not getattr(KimiK3HybridCacheManager, "accepts_appended_draft_layers", False), + reason="main does not append the drafter's layers to KimiK3HybridCacheManager", +) + +TARGET_LAYER_IDS = (3, 7, 11, 15) # the target's MLA layers; its KDA layers hold no KV pages +DRAFT_LAYER_IDS = tuple(range(61, 66)) # appended after the target's num_hidden_layers +HEAD_DIM = 576 # one MLA latent per token: kv_lora_rank + qk_rope_head_dim +PAGE = 32 +PAGES = 6 +LAYER_SLOT = 100 # encoded-offset distance between the K buffers of consecutive local layers + + +class _UnifiedK3Manager(KimiK3HybridCacheManager): + """``KimiK3HybridCacheManager`` with the drafter's layers appended, minus its native pools.""" + + def __init__(self, *, per_layer_page_tables=True, draft_dtype=torch.float8_e4m3fn, cp=None): + # No super().__init__(): it builds the native pools. The layer bookkeeping below is what + # MambaHybridCacheManagerV2.__init__ derives from draft_layer_specs. + self.layer_specs = None + self.pp_layers = list(TARGET_LAYER_IDS + DRAFT_LAYER_IDS) + self.layer_offsets = {layer_id: local for local, layer_id in enumerate(self.pp_layers)} + self._legacy_draft_layer_ids = () + self._appended_draft_layer_specs = { + self.layer_offsets[layer_id]: KVCacheLayerSpec( + global_layer_id=layer_id, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.FP8, + num_kv_heads=1, + head_dim=HEAD_DIM, + buffer_roles=(Role.KEY,), + ) + for layer_id in DRAFT_LAYER_IDS + } + self.kv_factor = 1 + num_target = len(TARGET_LAYER_IDS) + # Target layers in layer group 0 / pool 0, the drafter's in group 1 / pool 1. + groups = [0] * num_target + [1] * len(DRAFT_LAYER_IDS) + self.layer_to_pool_mapping_dict = dict(enumerate(groups)) + self.kv_cache_pool_mapping = torch.tensor([[group, 0] for group in groups]) + self.enable_swa_scratch_reuse = False + self._use_per_layer_page_tables = per_layer_page_tables + self.max_blocks_per_seq = 4 + self.kv_cache_map = {} + self.mapping = cp or Mapping() + self._pools = { + layer_id: torch.zeros(PAGES, 1, 1, PAGE, HEAD_DIM, dtype=draft_dtype) + for layer_id in DRAFT_LAYER_IDS + } + self.impl = SimpleNamespace( + get_page_index_converter=lambda local, role: SimpleNamespace( + scale=1, layer_offset=LAYER_SLOT * local + ), + get_page_index_scale=lambda local, role: 1, + ) + + def get_buffers(self, layer_idx, kv_layout="NHD"): + assert kv_layout == "HND" + return self._pools[layer_idx] + + +def _worker(manager=None): + worker = K3DFlashWorker( + SimpleNamespace(attention_backend="CUTEDSL", max_draft_len=7), Mapping() + ) + # What _lazy_init_ctx_buffers sets before binding: a manager that is the target metadata's + # own is the unified cache, and the MLA drafter keeps one latent per token. + worker._ctx_unified_cache = manager is not None + worker._ctx_kv_factor = 1 + return worker + + +def _resource_manager(manager): + managers = {ResourceManagerType.KV_CACHE_MANAGER: manager} + return SimpleNamespace(get_resource_manager=managers.get) + + +def _bind(worker, manager, dtype=torch.float8_e4m3fn): + pool = worker._managed_ctx_pool(manager, len(DRAFT_LAYER_IDS), 1, HEAD_DIM, dtype, 1) + return pool, pool is not None and worker._init_ctx_block_offsets( + manager, pool, len(DRAFT_LAYER_IDS), 1, HEAD_DIM, PAGE + ) + + +@requires_appended_draft_layers +def test_the_target_manager_is_the_draft_cache(): + manager = _UnifiedK3Manager() + assert manager.draft_layer_ids == DRAFT_LAYER_IDS + assert manager.get_layer_kv_factor(DRAFT_LAYER_IDS[0]) == 1 + assert _worker().get_draft_kv_cache_manager(_resource_manager(manager)) is manager + + +@requires_appended_draft_layers +@pytest.mark.parametrize("per_layer_page_tables", [True, False]) +def test_binds_the_appended_draft_layers(per_layer_page_tables): + manager = _UnifiedK3Manager(per_layer_page_tables=per_layer_page_tables) + worker = _worker(manager) + pool, bound = _bind(worker, manager) + + assert bound + assert worker._ctx_layer_ids == DRAFT_LAYER_IDS + assert [layer.data_ptr() for layer in pool] == [ + manager._pools[layer_id].data_ptr() for layer_id in DRAFT_LAYER_IDS + ] + assert worker._ctx_pool_idx == 1 + # A V2 view's stride(0) is one layer slot, so an entry decodes as offset // kv_factor. + assert worker._ctx_block_divisor == 1 + # Per-layer page tables address the slot origin; the view starts at the first draft layer's + # K buffer (local layer 4). + expected_offset = LAYER_SLOT * len(TARGET_LAYER_IDS) if per_layer_page_tables else 0 + assert worker._ctx_block_offset == expected_offset + + +@requires_appended_draft_layers +def test_a_unified_pool_of_the_wrong_dtype_raises_instead_of_falling_back(): + manager = _UnifiedK3Manager(draft_dtype=torch.bfloat16) + with pytest.raises(ValueError, match="unified draft KV pool does not match"): + _bind(_worker(manager), manager) + + +@requires_appended_draft_layers +def test_a_separate_pool_of_the_wrong_dtype_falls_back_to_the_arena(): + manager = _UnifiedK3Manager(draft_dtype=torch.bfloat16) + pool, bound = _bind(_worker(), manager) + assert pool is None and not bound + + +@requires_appended_draft_layers +def test_draft_layers_of_mixed_lifecycles_cannot_share_one_table(): + manager = _UnifiedK3Manager() + manager.layer_to_pool_mapping_dict[len(TARGET_LAYER_IDS) + 1] = 2 + with pytest.raises(NotImplementedError, match="same lifecycle"): + _bind(_worker(manager), manager) + + +def _helix(cp_size, cp_rank): + return SimpleNamespace(cp_size=cp_size, cp_rank=cp_rank, has_cp_helix=lambda: cp_size > 1) + + +@pytest.mark.parametrize( + "target, drafter, ok", + [ + ((1, 0), None, True), + ((8, 3), (8, 3), True), + ((8, 3), None, False), + ((8, 3), (8, 4), False), + ((1, 0), (8, 3), False), + ], +) +def test_the_drafter_shards_its_pages_as_the_target_manager_does(target, drafter, ok): + manager = _UnifiedK3Manager() + manager.mapping = _helix(*target) + worker = _worker(manager) + worker._helix_mapping = None if drafter is None else _helix(*drafter) + if ok: + worker._check_unified_helix(manager) + else: + with pytest.raises(ValueError, match="helix-sharded"): + worker._check_unified_helix(manager) + + +@pytest.mark.parametrize( + "tokens_per_block, reads", + [(32, True), (64, True), (128, True), (48, False), (256, False), (1, False)], +) +def test_the_drafter_declares_which_manager_pages_it_reads(tokens_per_block, reads): + drafter = SimpleNamespace(dflash_attention_backend="CUTEDSL", _helix_mapping=None) + assert K3MLADSparkForCausalLM.reads_managed_ctx_pages(drafter, tokens_per_block) is reads + assert K3MLADSparkForCausalLM.reads_helix_sharded_ctx_pages.fget(drafter) is False + drafter._helix_mapping = _helix(8, 0) + assert K3MLADSparkForCausalLM.reads_helix_sharded_ctx_pages.fget(drafter) is True diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_engine_probes.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_engine_probes.py new file mode 100644 index 000000000000..00ffc952a572 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_engine_probes.py @@ -0,0 +1,688 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Engine names the private Kimi-K3 package probes behind the switches its Rubin deployment runs. + +The package reaches main's engine through duck-typed probes (``hasattr``, ``getattr(..., None)``, +``inspect.signature``). A probe for a name the runtime object does not have returns its fallback, +so the switch behind it stops working without an error: ``TRTLLM_K3_LD_FUSED`` on main's FC12 leaf +was one such case. Each test pins the probes of one switch against the class the package gets at +runtime (the leaf ``create_moe`` builds, the KV cache manager the executor creates, ...), so a +rename on either side fails here rather than as a slower step in a perf run. + +CPU only: classes and signatures, no kernels. +""" + +import dataclasses +import importlib +import inspect +from types import SimpleNamespace + +import pytest +from torch import nn + +K3_RUBIN = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default" +) +TARGETS = ("tp1", "helix8", "helix16_moetp4ep4", "dep16") + + +def _k3_module(name): + """A module of the Kimi K3 target package (``modeling``, ``drafter``).""" + return importlib.import_module(f"{K3_RUBIN}.{name}") + + +def _bare(cls): + """An instance of the ``nn.Module`` ``cls`` without running its constructor.""" + instance = cls.__new__(cls) + nn.Module.__init__(instance) + return instance + + +def _is_property(cls, name): + return isinstance(inspect.getattr_static(cls, name, None), property) + + +def _select_profile(monkeypatch, target): + """Make ``target``'s feature profile the active one for fields read when called.""" + features = _k3_module("features") + profile = features.PROFILES[target] + monkeypatch.setattr(features, "_ACTIVE", profile) + return profile + + +@pytest.mark.parametrize("target", TARGETS) +def test_runtime_hooks_reach_the_engine(monkeypatch, target): + """``TLLM_SPEC_HEAD_FUSION`` (the decode step's head) and ``TRTLLM_K3_ULL_STEP_PREP`` (the + packed step inputs): the engine arms them only from the ``RuntimeHooks`` it reads off the + outermost model (``runtime_hooks_of``), which is the package's VL wrapper. The wrapper must + hand on the text model's hooks, and the text model must declare exactly the hooks its target's + profile turns on; everything else stays at the default.""" + from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, runtime_hooks_of + + spec_decode_head = _k3_module("kernels.spec_decode_head").spec_decode_head + profile = _select_profile(monkeypatch, target) + text = _bare(_k3_module("modeling").KimiK3DefaultForCausalLM) + wrapper = _bare(_k3_module("modeling").ModelingV2KimiK3Nvfp4Sm107Default) + wrapper.llm = text + + hooks = runtime_hooks_of(wrapper) + assert hooks == runtime_hooks_of(text) + assert (hooks.decode_prologue is spec_decode_head) == profile.embedding_gather + assert hooks.pack_step_inputs == profile.pack_step_inputs + assert hooks.pack_spec_rng_uploads == profile.pack_step_inputs + assert ( + dataclasses.replace( + hooks, decode_prologue=None, pack_step_inputs=False, pack_spec_rng_uploads=False + ) + == DEFAULT_RUNTIME_HOOKS + ) + + +def test_runtime_hooks_are_on_where_the_recipes_set_them(): + """The recipes set ``TLLM_SPEC_HEAD_FUSION=1`` on helix8 and helix16 and + ``TRTLLM_K3_ULL_STEP_PREP=1`` on helix16 only.""" + features = _k3_module("features") + head = {t for t in TARGETS if features.PROFILES[t].embedding_gather} + pack = {t for t in TARGETS if features.PROFILES[t].pack_step_inputs} + assert head == {"helix8", "helix16_moetp4ep4"} + assert pack == {"helix16_moetp4ep4"} + + +def test_spec_decode_head_reads_main_names(): + """``kernels.spec_decode_head`` (``decode_prologue`` on helix8 / helix16) reads main's preparer + and metadata directly: the overlap switch, the previous step's offsets, the helix buffers and + the per-token validity flag. ``preprocess_inputs`` must run the prologue before its own head + and read the same names, so the two write the same step.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer + + source = inspect.getsource(DecoderInputPreparer.preprocess_inputs) + for name in ( + "self.decode_prologue", + "self._config.disable_overlap_scheduler", + "self._buffers.previous_pos_id_offsets_cuda", + "self._buffers.previous_kv_lens_offsets_cuda", + "num_chunked_ctx_requests", + "recompute_helix_spec_buffers", + "_helix_spec_tokens_valid", + ): + assert name in source, name + assert source.index("self.decode_prologue") < source.index( + "self._config.disable_overlap_scheduler" + ) + assert "get_runtime_tokens_per_gen_step" in inspect.signature(DecoderInputPreparer).parameters + assert "self.get_runtime_tokens_per_gen_step = get_runtime_tokens_per_gen_step" in ( + inspect.getsource(DecoderInputPreparer.__init__) + ) + for name in ( + "_helix_spec_tokens_valid", + "helix_position_offsets", + "helix_local_slots", + "helix_kv_bounds", + "on_update_kv_lens", + "recompute_helix_spec_buffers", + ): + assert hasattr(TrtllmAttentionMetadata, name), name + + +def test_routed_experts_tune_the_routing_tile_with_mains_runners(): + """The routed experts (every target but dep16) run their own kernel chains; the AutoTuner picks + the routing tile through main's two CuteDSL outer runners, called with the experts' chain as + ``forward_impl``. The experts read the tile of a tactic (``tile_size_of``) and the tuning config + off them, and FC12 offers only the tiles its fused kernel serves.""" + from tensorrt_llm._torch.moe.fused_moe.cutedsl.fused_fc12 import CuteDslFc12FusedMoENvfp4Runner + from tensorrt_llm._torch.moe.fused_moe.cutedsl.runners import CuteDslFusedMoENvfp4Runner + + routed = _k3_module("moe") + runners = routed._tile_runners() + assert runners["CUTEDSL"] == (CuteDslFusedMoENvfp4Runner, "CuteDslFusedMoE::run_moe_nvfp4") + assert runners["CUTEDSL_FC12"] == ( + CuteDslFc12FusedMoENvfp4Runner, + "CuteDslFc12FusedMoE::run_moe_nvfp4", + ) + params = inspect.signature(CuteDslFusedMoENvfp4Runner.__init__).parameters + for name in ( + "forward_impl", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "enable_finalize_fusion", + "enable_alltoall", + "workload_identity", + ): + assert name in params, name + assert CuteDslFusedMoENvfp4Runner.tile_size_of(-1) == 128 + assert CuteDslFc12FusedMoENvfp4Runner._tile_sizes() == [128, 256] + + +def test_latent_down_fusion_feeds_the_fc12_experts(): + """``TRTLLM_K3_LD_FUSED`` (helix8, helix16): the latent down projection hands the routed + experts its own NVFP4 output only where they run the FC12 kernels, whose input is quantized + with ``fp4_quantize`` against ``fc31_input_scale``, as k3_latent_down_fp4 quantizes it.""" + features = _k3_module("features") + assert {t for t in TARGETS if features.PROFILES[t].ld_fused} == {"helix8", "helix16_moetp4ep4"} + source = inspect.getsource(_k3_module("moe_runtime").K3MoERuntime._fused_latent_input) + assert '"CUTEDSL_FC12"' in source and "fc31_input_scale" in source + + +@pytest.mark.parametrize( + "target, expected", + [ + ("tp1", 0), + ("helix8", 16), + ("helix16_moetp4ep4", 16), + ("dep16", 0), + ], +) +def test_routed_experts_sm_reserve_follows_the_profile(monkeypatch, target, expected): + """Where B's recipes set the FC12 switches: ``TRTLLM_K3_FC12_SM_RESERVE=16`` on helix8 + (beside the shared expert confined to ``TRTLLM_K3_SHARED_SM=16`` SMs); die-major weights + (``TRTLLM_K3_FC12_DIE_LAYOUT``) are not available and only logged. On helix16 + B's small-batch experts kernel ran on 196 / 204 / 196 of the 212 SMs at 8 / 16 / 32 rows, + beside the 16-CTA persistent shared expert of the column-parallel tail. Only the FC12 kernels + take a reserve.""" + moe_runtime = _k3_module("moe_runtime") + profile = _k3_module("features").PROFILES[target] + monkeypatch.setattr(moe_runtime, "_K3_SHARED_SM", profile.shared_sm) + monkeypatch.setattr(moe_runtime, "_K3_FC12_DIE_MAJOR", profile.fc12_die_major) + fc12 = SimpleNamespace(kernels="CUTEDSL_FC12", sm_reserve=0) + moe_runtime._k3_routed_backend_opt_ins(fc12, column_parallel_tail=profile.ull2_coll_cp_tail) + assert fc12.sm_reserve == expected + other = SimpleNamespace(kernels="CUTEDSL", sm_reserve=0) + moe_runtime._k3_routed_backend_opt_ins(other, column_parallel_tail=profile.ull2_coll_cp_tail) + assert other.sm_reserve == 0 + + +class _RoutedCalled(Exception): + """Raised by a stand-in for the routed experts to stop the MoE tail right after them.""" + + +def _routed_experts_push(runtime, num_tokens): + """The ``output_push`` the MoE tail of a ``num_tokens``-row forward calls the routed experts + with (``None``: none).""" + import torch + + pushes = [] + + def routed_output(output_push=None): + pushes.append(output_push) + raise _RoutedCalled + + rows = torch.zeros(num_tokens, runtime.hidden_size) + with pytest.raises(_RoutedCalled): + runtime._rp_tail_forward(routed_output, rows, None, rows, None, True) + assert len(pushes) == 1 + return pushes[0] + + +@pytest.mark.parametrize("target", TARGETS) +def test_one_kernel_tail_requests_the_output_push(target, monkeypatch): + """``TRTLLM_K3_ULL_TAIL1`` with ``TRTLLM_K3_ULL2_COLL_CP_TAIL`` (helix16 generation): on the + column-parallel tail of at most ``k3_moe_tail1.MAX_M`` rows the routed experts are called with + ``output_push``, the latent workspace whose slots k3_moe_tail1 polls, where they push those + rows (``K3SmallBatchExperts.applies``); every other tail calls them without one.""" + import torch + + routed = _k3_module("moe") + assert "output_push" in inspect.signature(routed.K3RoutedExperts.forward).parameters + + LamportWorkspaceView = _k3_module("layers").LamportWorkspaceView + moe_runtime = _k3_module("moe_runtime") + tail1 = _k3_module("kernels.k3_moe_tail1") + profile = _k3_module("features").PROFILES[target] + monkeypatch.setattr(moe_runtime, "_K3_ULL2_COLL_CP_TAIL", profile.ull2_coll_cp_tail) + world = 16 + latent = LamportWorkspaceView( + world=world, + rank=0, + buffer=torch.empty(0), + mc_base=0, + peer_ptrs=torch.empty(0, dtype=torch.int64), + flags=torch.zeros(9, dtype=torch.uint32), + bytes_per_buffer=0, + ) + runtime = _bare(moe_runtime.K3MoERuntime) + runtime.hidden_size = tail1.HIDDEN + runtime.moe_hidden_size = tail1.LATENT + runtime.shared_expert_stream = None + # The small-batch kernel pushes up to 32 rows. + runtime.small_experts = SimpleNamespace(applies=lambda num_tokens: 1 <= num_tokens <= 32) + # The tail's workspaces and k3_moe_tail1's operands as construction and loading build them; + # the column-parallel tail serves every row count here. + runtime._cp_ws = SimpleNamespace(world=world) + runtime._cp_latent_ws = SimpleNamespace(world=world, rank=0) + runtime._cp_latent_push = latent + runtime._tail1_wexp = torch.ones(tail1.HIDDEN // 128, tail1.LATENT // 128, dtype=torch.uint8) + runtime._cp_tail_applies = lambda num_tokens, target: True + + pushed = [m for m in range(1, 33) if _routed_experts_push(runtime, m) is latent] + expected = list(range(1, tail1.MAX_M + 1)) if profile.ull2_coll_cp_tail else [] + assert pushed == expected + + # Experts that do not push (no small-batch kernel, a communication strategy) keep the tail. + runtime.small_experts = SimpleNamespace(applies=lambda num_tokens: False) + assert _routed_experts_push(runtime, 1) is None + + +@pytest.mark.parametrize( + "change, max_tokens", + [ + ({}, 32), + ({"kernels": "CUTEDSL"}, 0), + ({"intermediate_size_per_partition": 100}, 0), + ], +) +def test_small_batch_experts_serve_the_fc12_calls_of_at_most_32_tokens(change, max_tokens): + """``K3SmallBatchExperts`` (helix8, helix16 generation) takes the forwards of 1..32 tokens of + routed experts running the FC12 kernels; the published routing word exists only on MoE-TP + slices once the counters do.""" + import torch + + small_batch = _k3_module("moe") + op = _k3_module("kernels.k3_small_batch_moe_op") + kernel = _k3_module("kernels.k3_small_batch_moe") + backend = SimpleNamespace( + intermediate_size_per_partition=change.get("intermediate_size_per_partition", 768), + hidden_size=3584, + expert_size_per_partition=224, + tp_size=4, + ) + routed = SimpleNamespace(backend=backend, kernels=change.get("kernels", "CUTEDSL_FC12")) + experts = small_batch.K3SmallBatchExperts(routed) + assert experts.max_tokens == max_tokens + assert [n for n in range(40) if experts.applies(n)] == list(range(1, max_tokens + 1)) + assert experts.published_routing_word(8) is None + experts.counters = op.new_small_batch_moe_counters(torch.device("cpu")) + word = experts.published_routing_word(8) + if max_tokens: + assert word.data_ptr() == experts.counters[kernel.CTR_ROUTE].data_ptr() + assert experts.published_routing_word(33) is None + backend.tp_size = 1 + assert experts.published_routing_word(8) is None + else: + assert word is None + + +def test_private_persistent_grids_take_the_shared_sm_reserve(monkeypatch): + """``TRTLLM_DWDP_PREFETCH_GATHER`` with ``TRTLLM_K3_ATTN_RES_SM_RESERVE`` (tp1 context): the + DWDP gather's resident SMs reach persistent grids through ``cute_dsl_utils.sm_reserve()`` + (``DwdpManager`` sets it before the model is built). The KDA prefill grid sizes itself by it + and the persistent attention-residual op takes it as its ``sm_reserve`` argument.""" + import torch + + from tensorrt_llm._torch import cute_dsl_utils + + assert _k3_module("attn_res").sm_reserve is cute_dsl_utils.sm_reserve + schema = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd.default._schema + assert [arg.name for arg in schema.arguments][-1] == "sm_reserve" + prefill_ops = pytest.importorskip(f"{K3_RUBIN}.ops.ssm._kda_prefill") + monkeypatch.setattr(cute_dsl_utils, "_sm_reserve", 16) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda device: SimpleNamespace(multi_processor_count=212), + ) + assert prefill_ops._kda_persistent_grid_blocks(torch.device("cuda", 0)) == 196 + + +def test_fp8_kv_layer_selector_on_the_cache_policy(): + """``TLLM_KIMI_K3_FP8_KV_LAYERS`` (both roles): the MLA layers take their FP8 KV-cache + fallback only where ``kimi_k3_cache_policy`` provides the selector the cache manager reads; + otherwise every layer keeps the model-wide NVFP4 dtype and a warning names the gap.""" + from tensorrt_llm._torch import kimi_k3_cache_policy + + selector = getattr(kimi_k3_cache_policy, "get_kimi_k3_fp8_kv_layer_ids", None) + assert callable(selector) + assert "nvfp4_enabled" in inspect.signature(selector).parameters + + +def test_fp4_mla_probe_on_trtllm_attention(): + """The FP4 MLA paths (packed decode, ``TRTLLM_K3_ULL4_MLA_PLAN_EARLY``) key off + ``mqa.uses_fp4_mla_attention`` on main's ``TrtllmAttention``.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + assert _is_property(TrtllmAttention, "uses_fp4_mla_attention") + + +def test_dspark_unified_draft_cache_probes_on_the_k3_hybrid_manager(): + """DSpark's drafter context KV in the target manager (``TRTLLM_DISAGG_SHARED_DRAFT_KV``). + The engine appends the drafter's layers to + ``KimiK3HybridCacheManager`` only when the class opts in, and only when the loaded drafter + says it reads the manager's pages (``reads_managed_ctx_pages`` / + ``reads_helix_sharded_ctx_pages``, which the private drafter must override). The private + worker finds the layers through ``draft_layer_ids`` and ``get_buffers(..., kv_layout=...)``. + Main's ``DFlashSpecMetadata.prepare`` hands a disaggregated generation request's new slot to + the worker's ``_seed_generation_slot``, and the private worker seeds the slot's context + length only when it reports ``_ctx_unified_cache``. The KDA bookkeeping after accept runs + only when the manager is a ``MambaHybridCacheManager``.""" + from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM + from tensorrt_llm._torch.modules.kimi_kda.cache_manager import KimiK3HybridCacheManager + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager + from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata + + assert issubclass(KimiK3HybridCacheManager, MambaHybridCacheManager) + assert issubclass(KimiK3HybridCacheManager, KVCacheManagerV2) + if not hasattr(KimiK3HybridCacheManager, "accepts_appended_draft_layers"): + pytest.skip("main has no drafter layers in the target manager (unified draft KV)") + assert KimiK3HybridCacheManager.accepts_appended_draft_layers is True + assert _is_property(KVCacheManagerV2, "draft_layer_ids") + assert "kv_layout" in inspect.signature(KVCacheManagerV2.get_buffers).parameters + for name in ("get_layer_kv_factor", "get_layer_page_index_scale", "get_draft_kv_cache_view"): + assert callable(getattr(KVCacheManagerV2, name, None)), name + assert "_seed_generation_slot" in inspect.getsource(DFlashSpecMetadata.prepare) + worker = _k3_module("drafter").K3DSparkWorker + assert "_ctx_unified_cache" in inspect.getsource(worker._seed_generation_slot) + drafter = _k3_module("drafter").K3MLADSparkForCausalLM + assert drafter.reads_managed_ctx_pages is not DFlashForCausalLM.reads_managed_ctx_pages + assert ( + drafter.reads_helix_sharded_ctx_pages is not DFlashForCausalLM.reads_helix_sharded_ctx_pages + ) + + +def test_kda_inforward_snapshot_plan_on_the_private_metadata(): + """In-forward KDA snapshots (context, block reuse): ``KimiK3HybridCacheManager`` plans them + whenever main's KDA kernels are present, whatever the model, and + ``finish_inforward_snapshots`` then commits the prefix from the live slot and copies the + staging rows back. ``Mamba2Metadata.prepare`` hands the plan to + ``_prepare_inforward_split``, which the private metadata turns into a ``kda_snapshot`` plan + on the live slots with the plan's staging rows; the KDA layers read those rows through + ``inforward_snapshot_stage_views``.""" + import pytest + import torch + + from tensorrt_llm._torch.modules.kimi_kda.cache_manager import KimiK3HybridCacheManager + from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + InforwardSnapshotEntry, + MambaHybridCacheManagerV2, + ) + + assert issubclass(KimiK3HybridCacheManager, MambaHybridCacheManagerV2) + for name in ("inforward_snapshot_stage_views", "state_index", "finish_inforward_snapshots"): + assert callable(getattr(MambaHybridCacheManagerV2, name, None)), name + assert "_prepare_inforward_split(" in inspect.getsource(Mamba2Metadata.prepare) + metadata_cls = _k3_module("kda").K3MambaMetadata + assert list(inspect.signature(metadata_cls._prepare_inforward_split).parameters) == list( + inspect.signature(Mamba2Metadata._prepare_inforward_split).parameters + ) + + metadata = metadata_cls.__new__(metadata_cls) + metadata._state_indices_long = torch.zeros(4, dtype=torch.long) + metadata.kda_snapshot = None + manager = SimpleNamespace(conv_state_shape=(12, 3), state_index={11: 7, 13: 2}.__getitem__) + attn_metadata = SimpleNamespace(seq_lens=torch.tensor([200, 130, 96]), request_ids=[11, 12, 13]) + plan = { + 11: InforwardSnapshotEntry(row=1, point=1152, num_before=128, num_after=72), + 13: InforwardSnapshotEntry(row=0, point=64, num_before=64, num_after=32), + } + metadata._prepare_inforward_split(attn_metadata, manager, plan, 3, 3) + snapshot = metadata.kda_snapshot + assert snapshot.num_snapshots == 2 + assert snapshot.slots.tolist() == [7, 2] + assert snapshot.stage_rows.tolist() == [1, 0] + assert snapshot.token_indices.tolist() == [128, 330 + 64] + assert snapshot.chunks.tolist() == [2, -1, 1] + with pytest.raises(RuntimeError, match="generation requests"): + metadata._prepare_inforward_split(attn_metadata, manager, plan, 3, 4) + misaligned = {13: InforwardSnapshotEntry(row=0, point=48, num_before=48, num_after=48)} + with pytest.raises(RuntimeError, match="does not match the plan"): + metadata._prepare_inforward_split(attn_metadata, manager, misaligned, 3, 3) + + +@pytest.mark.parametrize("takes_flag", [False, True]) +def test_dspark_helix_drafter_probes(takes_flag, monkeypatch): + """The DSpark drafter under Helix (generation): its Helix exchange keeps the NCCL reformats + eager only where main's ``_helix_post_process`` takes ``compile_nccl_reformat``, and calls it + without the keyword otherwise.""" + drafter = _k3_module("drafter") + fn, _ = drafter._helix_post_process_fn() + assert callable(fn) + calls = [] + + def fake(*args, **kwargs): + calls.append((args, kwargs)) + return "out" + + monkeypatch.setattr(drafter, "_helix_post_process_fn", lambda: (fake, takes_flag)) + mask = object() + assert drafter._draft_helix_post_process(1, 2, 3, 4, 5, mask) == "out" + expected = {"zero_kv_mask": mask} + if takes_flag: + expected["compile_nccl_reformat"] = False + assert calls == [((1, 2, 3, 4, 5), expected)] + + +def test_dspark_latent_cache_helix_mapping(): + """The DSpark drafter under Helix with ``TRTLLM_DISAGG_SHARED_DRAFT_KV=1`` (generation): main + unifies the draft KV with a CP > 1 target only when the drafter reports + ``reads_helix_sharded_ctx_pages`` (``_dflash_unify_blocker``), which the private drafter does + when its latent cache is page-sharded. The sharding comes from the target's helix CP mapping: + tp = 1 and cp over the repurposed TP group, on the target's CP ranks, with the target's + all-to-all transport for ``TRTLLM_K3_ULL2_DRAFT_A2A``.""" + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + from tensorrt_llm.mapping import CpType, Mapping + + shard = _k3_module("drafter")._latent_cache_helix_mapping + if hasattr(KvCacheCreator, "_unified_drafter_reads_helix_pages"): + assert "_unified_drafter_reads_helix_pages()" in inspect.getsource( + KvCacheCreator._dflash_unify_blocker + ) + assert "reads_helix_sharded_ctx_pages" in inspect.getsource( + KvCacheCreator._unified_drafter_reads_helix_pages + ) + else: + # Main without the helix unified draft KV refuses to unify a CP > 1 target; the drafter + # then keeps its own manager and the sharding below serves its latent cache alone. + assert "cp_size" in inspect.getsource(KvCacheCreator._dflash_unify_blocker) + for rank in (0, 5): + helix = Mapping( + world_size=8, + rank=rank, + gpus_per_node=4, + tp_size=1, + cp_size=8, + cp_config={"cp_type": CpType.HELIX, "use_nccl_for_alltoall": False}, + ) + repurposed = helix.repurpose_helix_cp_to_tp() + inherited = shard(repurposed, helix, inherit_cp_config=True) + assert (inherited.tp_size, inherited.cp_size, inherited.pp_size) == (1, 8, 1) + assert inherited.has_cp_helix() + assert (inherited.cp_rank, inherited.rank) == (helix.cp_rank, rank) + assert inherited.cp_config["use_nccl_for_alltoall"] is False + assert shard(repurposed, helix).cp_config == {"cp_type": CpType.HELIX} + assert shard(repurposed, None) is None + plain = Mapping(world_size=8, rank=0, gpus_per_node=4, tp_size=8) + assert shard(plain, None) is None + + +def test_router_picks_the_routing_kernel(): + """The routing each forward hands the routed experts (``K3MoERuntime._router``): the CuTe + routing tail where ``_route_tail_target`` set a target, with its ``moe_sort`` tables where they + are of the experts asked for; else ``noaux_tc_route_sort`` where it serves; else + ``noaux_tc_op`` (the experts then run ``moe_sort``).""" + import torch + + moe_runtime = _k3_module("moe_runtime") + routing_types = _k3_module("moe") + assert _k3_module("kernels.k3_route_tail").TILE == 128 + ids, weights = torch.zeros(2, 16, dtype=torch.int32), torch.ones(2, 16) + tables = routing_types.SortTables( + 128, tuple(torch.zeros(1, dtype=torch.int32) for _ in range(6)) + ) + calls = [] + + def route_tail(logits, offset, num): + calls.append(("tail", offset, num)) + return routing_types.SortedRouting(ids, weights, tables) + + def route_with_sort(logits, offset, num): + calls.append(("sort", offset, num)) + return routing_types.SortedRouting(ids, weights, tables) if logits.shape[0] <= 256 else None + + def route(logits): + calls.append(("topk",)) + return ids, weights + + runtime = _bare(moe_runtime.K3MoERuntime) + runtime.gate = SimpleNamespace( + route=route, route_tail=route_tail, route_with_sort=route_with_sort + ) + runtime.routed_experts = SimpleNamespace(consumes_sort_tables=True) + logits = torch.empty(2, 896) + + assert runtime._router((0, 896))(logits, 0, 896).sort_tables is tables + assert runtime._router((0, 896))(logits, 448, 448).sort_tables is None + assert runtime._router(None)(logits, 0, 896).sort_tables is tables + assert runtime._router(None)(torch.empty(300, 896), 0, 896).sort_tables is None + assert calls == [ + ("tail", 0, 896), + ("tail", 0, 896), + ("sort", 0, 896), + ("sort", 0, 896), + ("topk",), + ] + calls.clear() + runtime.routed_experts = SimpleNamespace(consumes_sort_tables=False) + assert runtime._router(None)(logits, 0, 896).sort_tables is None + assert calls == [("topk",)] + + +def test_dspark_split_plan_planner(): + """The DSpark drafter's split-KV plan (generation) runs the planner kernel of + ``kernels/spec/mla_split_plan.py`` at the drafter's piece cost (``_launch_mla_split_plan``) on the + CuTe DSL MLA decode runner's geometry (``_mla_split_plan_geometry``): the kernel takes the + piece cost as an argument, and the runner exposes what the geometry reads.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + drafter = _k3_module("drafter") + plan = _k3_module("kernels.spec.mla_split_plan") + params = inspect.signature(plan._mla_split_kv_plan_kernel.fn).parameters + for name in ("PIECE_COST", "MAX_PIECES", "BLOCK", "KV_TILE", "MIN_TILES"): + assert name in params, name + for name in ( + "_SPLIT_PLAN_MAX_REQUESTS", + "_SPLIT_PLAN_MAX_SPLIT", + "_SPLIT_PLAN_KV_TILE", + "_SPLIT_PLAN_MIN_TILES", + ): + assert isinstance(getattr(plan, name, None), int), name + runner = CuteDSLNVMlaDecodeBlackwellRunner + assert list(inspect.signature(runner.get_num_q_tiles).parameters) == ["num_heads", "seq_len_q"] + assert callable(runner._get_max_active_blocks) + assert runner._CLUSTER_SHAPE_MNK[0] >= 1 + assert isinstance(drafter._MLA_SPLIT_PLAN_PIECE_COST_TILES, int) + assert drafter._MLA_SPLIT_PLAN_PIECE_COST_TILES > 0 + + +def test_dspark_drafter_declares_its_slot_convention(): + """DSpark's ``shift_label`` slot convention (generation): the engine sizes the unified draft + KV from ``draft_model.shift_label`` (``get_draft_block_width``, K slots instead of K + 1), which + main's ``DFlashForCausalLM`` initializes off and the drafter's DSpark heads set.""" + from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM + + drafter = _k3_module("drafter") + assert issubclass(drafter.K3MLADSparkForCausalLM, DFlashForCausalLM) + heads = drafter._DSparkHeadMixin.__new__(drafter._DSparkHeadMixin) + heads._init_dspark_heads(SimpleNamespace()) + assert heads.shift_label is True + heads._init_dspark_heads(SimpleNamespace(dflash_config={"shift_label": False})) + assert heads.shift_label is False + + +def test_dspark_worker_takes_the_shared_slot_calls(): + """Main's ``DFlashSpecMetadata.prepare()`` drives the drafter's context slots by calling + methods on the worker it is bound to, and on K3 that is the private ``K3DSparkWorker``: every + ``worker.(...)`` call there binds to the private method of that name. A missing + method or a changed signature raises ``TypeError`` on the first step after the context + buffers exist, so no DSpark run starts.""" + import ast + import textwrap + + from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata + + tree = ast.parse(textwrap.dedent(inspect.getsource(DFlashSpecMetadata.prepare))) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "worker" + ] + names = {call.func.attr for call in calls} + assert {"_assign_slot", "_write_ctx_len"} <= names, names + + worker = _k3_module("drafter").K3DSparkWorker + unbound = [] + for call in calls: + method = getattr(worker, call.func.attr, None) + if not callable(method): + unbound.append((call.func.attr, "missing")) + continue + args = [object()] * (1 + len(call.args)) + kwargs = {kw.arg: object() for kw in call.keywords if kw.arg is not None} + try: + inspect.signature(method).bind(*args, **kwargs) + except TypeError as error: + unbound.append((call.func.attr, str(error))) + assert not unbound, unbound + + +# The FP8 block-scale projections of K3's MLA layer; ``KimiK3ForCausalLM._finalize_weight_load`` +# swaps each one for ``_Fp8BlockScaleWeightReadLinear``. +_K3_MLA_PROJECTIONS = ("kv_a_proj_with_mqa", "q_b_proj", "kv_b_proj", "o_proj", "g_proj") +_K3_MLA_MODULES = ("mla",) +# Read through ``getattr(proj, name, None)``; the bias-free FP8 projection has none. +_OPTIONAL_PROJECTION_READS = frozenset({"bias"}) + + +def test_fp8_weight_read_linear_serves_the_k3_mla_projection_reads(): + """K3's MLA chains read attributes and methods of their projections, and on K3 those are + ``_Fp8BlockScaleWeightReadLinear``: every one read exists on the private module. + ``K3MLAAttention.cache_derived_state`` calls ``prepare_output_slice`` at post-load on every + Helix rank, so one missing there stops every Helix GEN worker before it serves.""" + import re + + import torch + + source = "\n".join(inspect.getsource(_k3_module(name)) for name in _K3_MLA_MODULES) + names = "|".join(_K3_MLA_PROJECTIONS) + reads = set(re.findall(rf"(?:self|layer)\.(?:{names})\.(\w+)\b", source)) + reads |= set(re.findall(r"\bproj\.(\w+)\b", source)) + reads |= set( + re.findall(rf"getattr\((?:(?:self|layer)\.(?:{names})|kv_b_proj), \"(\w+)\"", source) + ) + assert { + "prepare_output_slice", + "forward_slice", + "forward_mxfp8", + "forward_gated_mxfp8", + "gated_mxfp8_applies", + "takes_mid_mxfp8_input", + "supports_prequantized_input", + } <= reads, reads + + linear = _k3_module("fp8_weight_read")._Fp8BlockScaleWeightReadLinear + projection = linear(torch.zeros(256, 128, dtype=torch.float8_e4m3fn), torch.ones(2, 1), 256) + # Set by ``from_linear``, the constructor K3's swap uses. + projection.quant_config, projection.tp_size, projection.tp_rank = None, 1, 0 + missing = sorted( + attr + for attr in reads + if attr not in _OPTIONAL_PROJECTION_READS and not hasattr(projection, attr) + ) + assert not missing, missing diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_features.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_features.py new file mode 100644 index 000000000000..fc7f932710f9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_features.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-deployment feature profiles of the Kimi K3 SM107 target (``features.py``).""" + +import dataclasses +import subprocess +import sys +import textwrap +import types + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + features, +) + +pytestmark = pytest.mark.cpu_only + +_TARGETS = {"tp1", "helix8", "helix16_moetp4ep4", "dep16"} + + +@pytest.fixture +def fresh(monkeypatch): + """A process in which no profile is selected and no other module of the target is imported.""" + monkeypatch.setattr(features, "_ACTIVE", features.K3Features()) + monkeypatch.setattr(features, "_SELECTED", None) + monkeypatch.setattr(features, "_PACKAGE", "k3_features_test_no_such_package") + warnings = [] + from tensorrt_llm.logger import logger + + monkeypatch.setattr(logger, "warning", lambda msg, *args, **kwargs: warnings.append(msg)) + return warnings + + +def test_one_profile_per_routing_target(): + assert set(features.PROFILES) == _TARGETS + assert all(isinstance(p, features.K3Features) for p in features.PROFILES.values()) + + +def test_every_field_replaces_one_retired_variable(): + fields = {f.name for f in dataclasses.fields(features.K3Features)} + mapped = [f for f in features.RETIRED_ENV.values() if f is not None] + assert sorted(mapped) == sorted(fields) + + +def test_context_profile_only_sets_context_features(): + assert set(features.PROFILES["tp1"].changed()) == { + "kda_fused_conv_prefill", + "kda_prepad_prefill", + } + + +def test_every_field_is_set_by_some_target_and_left_off_by_another(): + for f in dataclasses.fields(features.K3Features): + values = {repr(getattr(p, f.name)) for p in features.PROFILES.values()} + assert len(values) > 1, f"{f.name}: every target uses {values}; the code should not ask" + + +def test_active_is_the_plain_path_before_selection(fresh): + assert features.active() == features.K3Features() + assert features.active().changed() == {} + assert features.selected() is None + + +def test_select_is_once_per_process(fresh): + assert features.select("helix8") is features.PROFILES["helix8"] + assert features.active() is features.PROFILES["helix8"] + assert features.selected() == "helix8" + assert features.select("helix8") is features.PROFILES["helix8"] + with pytest.raises(RuntimeError, match="helix8"): + features.select("dep16") + + +def test_select_rejects_unknown_targets(fresh): + with pytest.raises(KeyError): + features.select("tp8") + + +def test_select_after_a_k3_rubin_import_raises(fresh, monkeypatch): + monkeypatch.setitem(sys.modules, f"{features._PACKAGE}.modeling", types.ModuleType("modeling")) + with pytest.raises(RuntimeError, match="modeling"): + features.select("tp1") + assert features.selected() is None + + +def test_select_warns_about_retired_variables(fresh, monkeypatch): + for name in features.RETIRED_ENV: + monkeypatch.delenv(name, raising=False) + features.select("dep16") + assert fresh == [] + monkeypatch.setattr(features, "_SELECTED", None) + monkeypatch.setenv("TRTLLM_K3_ROUTE_TAIL", "1") + features.select("dep16") + assert len(fresh) == 1 and "TRTLLM_K3_ROUTE_TAIL" in fresh[0] + + +def test_select_names_the_variable_main_reads_for_a_renamed_one(fresh, monkeypatch): + """The helix recipes set the MNNVL one-shot limit under a name nothing reads any more; the + warning names the variable main's all-reduce reads instead.""" + for name in features.RETIRED_ENV: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("TRTLLM_MNNVL_AR_ONESHOT_MAX_BYTES", "2097152") + features.select("helix8") + assert len(fresh) == 1 + assert "TRTLLM_MNNVL_AR_ONESHOT_MAX_BYTES (set TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES" in fresh[0] + + +# What the routing does, in a process of its own: select the deployment's profile, then import the target. +_PACKAGE = features.__name__.rpartition(".")[0] +# The target's modules that read the profile at import (each holds the profile as ``_F``). +_PROFILE_READERS = ("modeling", "moe_runtime", "attn_res", "fp8_weight_read") + + +def _python(code): + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(code)], capture_output=True, text=True, timeout=900 + ) + assert result.returncode == 0, result.stdout[-2000:] + result.stderr[-6000:] + return result.stdout + + +@pytest.mark.parametrize("target", sorted(_TARGETS)) +def test_select_then_import_the_target(target): + """``select`` and then ``modeling``, as ``route()`` does: the modules import without a cycle, every + module that reads the profile at import holds the selected one, and each weight-loading mixin's + ``load_weights`` is the one its class resolves.""" + out = _python( + f""" + import importlib + from {_PACKAGE} import features + + profile = features.select({target!r}) + modeling = importlib.import_module("{_PACKAGE}.modeling") + weights = importlib.import_module("{_PACKAGE}.weights") + for name in {_PROFILE_READERS!r}: + module = importlib.import_module("{_PACKAGE}." + name) + assert module._F is profile, name + assert importlib.import_module("{_PACKAGE}.moe_runtime")._K3_SHARED_SM == profile.shared_sm + target_cls, text_cls = modeling.ModelingV2KimiK3Nvfp4Sm107Default, modeling.KimiK3DefaultForCausalLM + assert target_cls.__mro__[1] is weights.KimiK3DefaultWeights + assert text_cls.__mro__[1] is weights.KimiK3DefaultTextWeights + assert target_cls.load_weights is weights.KimiK3DefaultWeights.load_weights + assert text_cls.load_weights is weights.KimiK3DefaultTextWeights.load_weights + print("imported", features.selected()) + """ + ) + assert f"imported {target}" in out + + +@pytest.mark.parametrize("target", sorted(_TARGETS)) +def test_import_before_select_is_rejected(target): + """Importing ``modeling`` first would leave its import-time settings at the plain path's: the + profile's selection must refuse such a process.""" + out = _python( + f""" + import importlib + from {_PACKAGE} import features + + importlib.import_module("{_PACKAGE}.modeling") + try: + features.select({target!r}) + except RuntimeError as error: + assert "{_PACKAGE}.modeling" in str(error), error + assert features.selected() is None + print("rejected") + """ + ) + assert "rejected" in out diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_glue.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_glue.py new file mode 100644 index 000000000000..8a73edaa560e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_glue.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decoder glue of the Kimi K3 target against main's helpers it replaced: bitwise. + +* ``_Fp8BlockScaleWeightReadLinear.from_linear`` on a ``core.Fp8BlockScaleLinear`` (loaded, or + loaded and prepared; one rank or a COLUMN / ROW shard) gives the projection built from the codes + and scales that main's ``Linear(FP8_BLOCK_SCALES)`` loads from the same checkpoint tensors. +* ``_shard_dim`` slices as main's ``load_weight_shard`` on every split K3 loads through it. +""" + +import pytest +import torch +from utils.util import getSMVersion + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + layers as linear, +) +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode, load_weight_shard +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + +FP8 = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) +COL, ROW = TensorParallelMode.COLUMN, TensorParallelMode.ROW +ROWS = (1, 3, 8, 16, 31, 64, 100, 128, 129, 256, 1000) + +requires_block_fp8 = pytest.mark.skipif( + getSMVersion() not in (100, 103, 107) + or (getSMVersion() == 107 and not IS_CUTLASS_DSL_RUBIN_AVAILABLE), + reason="needs a supported block-scaled FP8 GPU and toolchain", +) + + +def _fp8_weight_read(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + fp8_weight_read, + ) + + return fp8_weight_read + + +def _weights(): + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default import ( + weights, + ) + + return weights + + +def _checkpoint_fp8(n, k, seed): + """E4M3 codes and FP32 (non power-of-two) 128x128 block scales of a random ``[n, k]`` weight.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + weight = torch.randn(n, k, device="cuda", generator=gen) * 0.02 + blocks = torch.nn.functional.pad(weight, (0, -k % 128, 0, -n % 128)).view( + -(-n // 128), 128, -(-k // 128), 128 + ) + scale = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-8) / 448.0 + expanded = scale.repeat_interleave(128, 0)[:n].repeat_interleave(128, 1)[:, :k] + return (weight / expanded).to(torch.float8_e4m3fn), scale + + +def _swapped(source): + proj = _fp8_weight_read()._Fp8BlockScaleWeightReadLinear.from_linear(source) + proj.post_load_weights() + return proj + + +def _from_main_linear(main): + """The FP8 weight read over the codes, scales and shard that main's ``Linear`` loaded.""" + proj = _fp8_weight_read()._Fp8BlockScaleWeightReadLinear( + main.weight.detach(), main.weight_scale.detach(), main.out_features + ) + proj.quant_config = main.quant_config + proj.tp_size, proj.tp_rank = main.tp_size, main.tp_rank + # main's ``Linear`` holds the checkpoint codes and FP32 grid as loaded: post-load prepares them. + proj._weights_transformed = False + proj.post_load_weights() + return proj + + +def _assert_same_projection(got, ref, k): + assert (got.out_features, got.in_features) == (ref.out_features, ref.in_features) + assert (got.tp_size, got.tp_rank) == (ref.tp_size, ref.tp_rank) + assert got.quant_config.quant_algo == ref.quant_config.quant_algo + assert torch.equal(got.weight.view(torch.uint8), ref.weight.view(torch.uint8)) + assert torch.equal(got.weight_scale, ref.weight_scale) + assert got._low_m_max == ref._low_m_max + for a, b in ( + (got._low_m_weight_scale, ref._low_m_weight_scale), + (got.lowm_weight_scale, ref.lowm_weight_scale), + ): + assert (a is None) == (b is None) + if a is not None: + assert torch.equal(a, b) + for m in ROWS: + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + assert torch.equal(got(x), ref(x)), m + + +# K3's checkpoint FP8 projection shapes (q_a + kv_a fused, q_b / kv_b / o / g heads, KDA q/k/v/g/o). +FP8_SHAPES = [(2112, 7168), (1152, 1536), (2304, 512), (7168, 1024), (512, 7168), (4096, 7168)] + + +@requires_block_fp8 +@pytest.mark.parametrize("prepared", [False, True]) +@pytest.mark.parametrize("n,k", FP8_SHAPES) +@torch.inference_mode() +def test_weight_read_from_core_linear_matches_main_linear(n, k, prepared): + codes, scale = _checkpoint_fp8(n, k, seed=n + k) + with torch.device("cuda"): + main = Linear(k, n, bias=False, dtype=torch.bfloat16, quant_config=FP8, reduce_output=False) + main.load_weights([{"weight": codes.clone(), "weight_scale": scale.clone()}]) + core = linear.Fp8BlockScaleLinear(k, n, device="cuda") + core.load(codes.clone(), scale.clone()) + if prepared: + core.prepare() + _assert_same_projection(_swapped(core), _from_main_linear(main), k) + + +# MLA projections K3 splits over TP ranks: kv_b / q_b COLUMN over 64 heads x 256 / 192, o ROW over +# 64 heads x 128. +SHARDED = [(16384, 512, COL), (12288, 1536, COL), (7168, 8192, ROW)] + + +@requires_block_fp8 +@pytest.mark.parametrize("tp", [2, 4]) +@pytest.mark.parametrize("n,k,mode", SHARDED) +@torch.inference_mode() +def test_weight_read_from_sharded_core_linear_matches_main_linear(n, k, mode, tp): + codes, scale = _checkpoint_fp8(n, k, seed=n * k + tp) + for rank in range(tp): + mapping = Mapping(world_size=tp, tp_size=tp, rank=rank) + with torch.device("cuda"): + main = Linear( + k, + n, + bias=False, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=mode, + quant_config=FP8, + reduce_output=False, + ) + main.load_weights([{"weight": codes.clone(), "weight_scale": scale.clone()}]) + core = linear.Fp8BlockScaleLinear( + k, n, tp_size=tp, tp_rank=rank, tp_mode=mode, device="cuda" + ) + core.load(codes.clone(), scale.clone()) + _assert_same_projection(_swapped(core), _from_main_linear(main), core.in_features) + + +# (shape, dim) of the splits K3 loads through ``_shard_dim``: KDA head-major rows (dim 0) and o_proj +# columns (dim 1), 1-D head vectors, width-1 dims, the dense / shared-expert gate, up (dim 0) and +# down (dim 1) halves. +SPLITS = [ + ((4096, 7168), 0), + ((7168, 4096), 1), + ((64,), 0), + ((64,), 1), + ((1, 1, 64, 1), 0), + ((18432, 7168), 0), + ((7168, 18432), 1), +] + + +@pytest.mark.parametrize("tp", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("shape,dim", SPLITS) +def test_shard_dim_matches_main_load_weight_shard(shape, dim, tp): + tensor = torch.randn(shape) + mode = COL if dim == 0 else ROW + for rank in range(tp): + got = _weights()._shard_dim(tensor, dim, tp, rank, torch.device("cpu")) + assert torch.equal(got, load_weight_shard(tensor, tp, rank, mode)), (rank, tp) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_ull_shared_fork.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_ull_shared_fork.py new file mode 100644 index 000000000000..fd4df2e52288 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/test_k3_ull_shared_fork.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Where the Kimi-K3 ULL MoE front forks the shared expert's GEMMs, and in what order it issues the fork. + +Under CUDA graph capture the order matters: the kernel captured first among an early-releasing kernel's +dependents keeps that kernel's stream and its release. A_mid's must be the front's head, the head's the +route kernel; the shared GEMMs depend on A_mid above ``_K3_ULL_FRONT_FORK_AT_HEAD_MAX_M`` rows and on the +head up to it, and are issued after the front either way. These tests drive +``K3MoERuntime._ull_fork_shared`` with a stand-in object that logs the event records and the shared +launch, and run the returned callbacks where the front would. + +CPU only: no kernels, no CUDA graphs. +""" + +import importlib +from types import SimpleNamespace + +import pytest + +MOE_RUNTIME = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__default.moe_runtime" # noqa: E501 + + +@pytest.fixture +def moe_runtime(): + return importlib.import_module(MOE_RUNTIME) + + +@pytest.mark.parametrize( + "lead, rows, expected", + [ + (False, 8, None), + (False, 32, None), + (True, 8, "head"), + (True, 16, "head"), + (True, 24, "after_head"), + (True, 32, "after_head"), + ], +) +def test_fork_placement(moe_runtime, monkeypatch, lead, rows, expected): + monkeypatch.setattr(moe_runtime, "_K3_ULL_FRONT_FORK_AT_HEAD_MAX_M", 16) + assert moe_runtime._ull_shared_fork(lead, rows) == expected + + +def _fork_log(moe_runtime, monkeypatch, lead, rows): + """Issue order of a front on ``rows`` rows: event records, the shared launch and the head.""" + monkeypatch.setattr(moe_runtime, "_K3_ULL_FRONT_FORK_AT_HEAD_MAX_M", 16) + log = [] + partial = object() + + def shared_first(identity, persist): + log.append(("shared", identity, persist)) + return partial + + moe = SimpleNamespace( + moe_main_event=SimpleNamespace(record=lambda: log.append("record")), + _ull_shared_first=shared_first, + _shared_gemms_can_lead=lambda: lead, + ) + after_head, after_front, forked = moe_runtime.K3MoERuntime._ull_fork_shared( + moe, "x", False, rows + ) + log.append("head") + if after_head is not None: + after_head() + log.append("route") + if after_front is not None: + after_front() + return log, forked, partial + + +@pytest.mark.parametrize( + "rows, expected", + [ + # Few rows: the event after the head (the shared GEMMs depend on it), the shared GEMMs after + # the route kernel, so the head's first dependent is the route kernel. + (16, ["head", "record", "route", "shared"]), + # More rows: the event before the head (the shared GEMMs depend on A_mid only, so the head is + # A_mid's first dependent), the shared GEMMs after the front, and no second record. + (32, ["record", "head", "route", "shared"]), + ], +) +def test_fork_order(moe_runtime, monkeypatch, rows, expected): + log, forked, partial = _fork_log(moe_runtime, monkeypatch, True, rows) + assert [e if isinstance(e, str) else e[0] for e in log] == expected + assert ("shared", "x", False) in log + assert forked == [partial] + + +def test_no_fork_without_lead(moe_runtime, monkeypatch): + log, forked, _ = _fork_log(moe_runtime, monkeypatch, False, 32) + assert log == ["head", "route"] + assert forked == [] diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_dflash_unified_kv_gate.py index e3fc35e2edb6..e812fcb765ce 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_dflash_unified_kv_gate.py @@ -120,7 +120,9 @@ def test_dspark_checkpoint_width_reaches_shared_cache_cost_and_construction( ) creator._speculative_config.max_draft_len = 7 creator._speculative_config.draft_is_embedded_in_target = False - creator._model_engine.model.draft_model = SimpleNamespace(_dspark_shift_label=shift_label) + creator._model_engine.model.draft_model = SimpleNamespace( + shift_label=shift_label, reads_managed_ctx_pages=lambda tokens_per_block: True + ) creator._kv_cache_config = KvCacheConfig() creator._tokens_per_block = 32 creator._max_batch_size = 4 @@ -276,6 +278,24 @@ def test_launcher_rejects_dflash_on_generation_worker_only( ) +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_embedded_dspark_draft_keeps_its_own_window(is_disagg: bool) -> None: + """An embedded DSpark draft (DeepSeek-V4) is not a DFlashForCausalLM: the + gate reports why it keeps its own window instead of calling hooks it lacks.""" + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout) + creator._speculative_config.spec_dec_mode = SimpleNamespace( + is_dflash=lambda: False, is_dspark=lambda: True, is_external_drafter=lambda: False + ) + creator._speculative_config.draft_is_embedded_in_target = True + creator._model_engine.model.draft_model = object() + + assert creator._loaded_standalone_dspark_drafter() is None + assert creator._dflash_unify_blocker() is not None + assert not creator._should_unify_dflash_kv_cache() + creator._check_shared_draft_kv_request() + + def test_disagg_without_shared_draft_kv_env_keeps_separate_cache( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -289,7 +309,7 @@ def test_disagg_without_shared_draft_kv_env_keeps_separate_cache( def test_dflash_with_different_tp_keeps_separate_cache(monkeypatch, tmp_path): spec = {"decoding_type": "DFlash"} - with pytest.raises(click.UsageError, match="same tensor_parallel_size and enable_attention_dp"): + with pytest.raises(click.UsageError, match="same attention tensor parallelism"): _launch_from_disagg_config( monkeypatch, tmp_path, @@ -336,9 +356,14 @@ def test_shared_cache_requested_with_unsupported_layout_fails_at_startup(): "context_servers,generation_servers,match", [ ( - {"enable_attention_dp": True}, - {"enable_attention_dp": False}, - "same tensor_parallel_size and enable_attention_dp", + {"tensor_parallel_size": 4, "enable_attention_dp": True}, + {"tensor_parallel_size": 4, "enable_attention_dp": False}, + "same attention tensor parallelism", + ), + ( + {}, + {"context_parallel_size": 8}, + "context_parallel_size 1", ), ( {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, @@ -351,7 +376,12 @@ def test_shared_cache_requested_with_unsupported_layout_fails_at_startup(): "attention_backend TRTLLM", ), ], - ids=["mismatched_attention_dp", "vanilla_dflash_backend", "default_dflash_backend"], + ids=[ + "mismatched_attention_dp", + "helix_generation", + "vanilla_dflash_backend", + "default_dflash_backend", + ], ) def test_launcher_rejects_shared_cache_for_ineligible_config( monkeypatch, tmp_path, context_servers, generation_servers, match @@ -376,6 +406,39 @@ def test_launcher_accepts_eligible_shared_cache_config(monkeypatch, tmp_path): ) +def test_launcher_shares_attention_dp_mla_drafter_across_tp(monkeypatch, tmp_path): + # Kimi K3 dep16: one-GPU context workers and a DEP16 generation worker hold + # every KV head per rank, as does a TP1 context worker without attention DP. + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV, raising=False) + spec = {"decoding_type": "DSpark", "attention_backend": "CUTEDSL"} + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + {"speculative_config": spec, "tensor_parallel_size": 1, "enable_attention_dp": True}, + {"speculative_config": spec, "tensor_parallel_size": 16, "enable_attention_dp": True}, + ) + + assert os.environ[_SHARED_DRAFT_KV_ENV] == "auto" + + +def test_launcher_shares_the_cutedsl_drafter_with_helix_generation(monkeypatch, tmp_path): + # Kimi K3 helix8: one-GPU context workers, TP1 x CP8 generation. + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV, raising=False) + spec = {"decoding_type": "DSpark", "attention_backend": "CUTEDSL"} + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + {"speculative_config": spec, "enable_attention_dp": True}, + {"speculative_config": spec, "context_parallel_size": 8}, + ) + + assert os.environ[_SHARED_DRAFT_KV_ENV] == "auto" + + @pytest.mark.parametrize( "request_mode,draft_heads,expected_mode", [("0", 8, "0"), (None, 4, "auto")], @@ -621,3 +684,17 @@ def test_launcher_checks_all_peer_configs( draft=_model_config(kv_heads=8, head_dim=128), ) assert creator._should_unify_dflash_kv_cache() is shared + + +@pytest.mark.parametrize( + "backend, tokens_per_block, reads", + [("TRTLLM", 32, True), ("TRTLLM", 64, True), ("TRTLLM", 128, False), ("FA4", 32, False)], +) +def test_dflash_drafter_reads_managed_pages_only_at_trtllm_gen_page_sizes( + backend: str, tokens_per_block: int, reads: bool +) -> None: + from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM + + # GQA DSpark drafters inherit this; at 128-token pages they keep their own pool. + drafter = SimpleNamespace(dflash_attention_backend=backend) + assert DFlashForCausalLM.reads_managed_ctx_pages(drafter, tokens_per_block) is reads diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_hybrid_unified_draft_kv.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_hybrid_unified_draft_kv.py new file mode 100644 index 000000000000..227b5751ec7c --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_hybrid_unified_draft_kv.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A standalone drafter's KV layers in a hybrid (recurrent + MLA) target manager.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaV2CacheLayoutPolicy +from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, + KVCacheManagerV2, + Role, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import mamba_cache_manager_v2 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( + MambaHybridCacheManagerV2, + _insert_draft_pool_ratio, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.quantization.mode import QuantMode +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + GpuCacheTierConfig, + KVCacheManagerConfig, + LayerId, + SsmLayerConfig, +) + +pytestmark = pytest.mark.cpu_only + +_NUM_TARGET_LAYERS = 6 +_NUM_DRAFT_LAYERS = 2 + + +class _K3LikeManager(MambaHybridCacheManagerV2): + accepts_appended_draft_layers = True + + +def _mla_config(*, num_layers, kv_quant_algo=None): + pretrained = SimpleNamespace( + num_hidden_layers=num_layers, + num_attention_heads=64, + kv_lora_rank=512, + qk_rope_head_dim=64, + hidden_size=7168, + torch_dtype=torch.bfloat16, + ) + return SimpleNamespace( + pretrained_config=pretrained, + quant_config=SimpleNamespace( + quant_mode=QuantMode.from_quant_algo(kv_cache_quant_algo=kv_quant_algo) + ), + ) + + +def _creator(*, manager_cls=_K3LikeManager, reads_pages=True, tp_size=16, is_disagg=False): + """A Kimi K3-like creator: NVFP4 MLA target, FP8 MLA drafter.""" + target = _mla_config(num_layers=_NUM_TARGET_LAYERS, kv_quant_algo=QuantAlgo.NVFP4) + draft = _mla_config(num_layers=_NUM_DRAFT_LAYERS, kv_quant_algo=QuantAlgo.FP8) + creator = object.__new__(KvCacheCreator) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace( + is_dflash=lambda: False, is_dspark=lambda: True, is_external_drafter=lambda: True + ), + attention_backend="CUTEDSL", + draft_is_embedded_in_target=False, + _allow_separate_draft_kv_cache=True, + draft_len_schedule=None, + ) + creator._draft_config = draft + creator._mapping = SimpleNamespace( + pp_size=1, + cp_size=1, + tp_size=tp_size, + enable_attention_dp=True, + pp_layers=lambda num_layers: list(range(num_layers)), + ) + creator._tokens_per_block = 128 + creator._max_seq_len = 4096 + creator._unification_max_seq_len = 4096 + creator._kv_cache_config = SimpleNamespace(max_attention_window=None) + creator._get_draft_max_attention_window = lambda max_seq_len, kv_cache_config: None + creator._llm_args = SimpleNamespace(attn_backend="TRTLLM", kv_cache_compression_config=None) + creator._model_engine = SimpleNamespace( + model=SimpleNamespace( + model_config=target, + config=SimpleNamespace(torch_dtype=torch.bfloat16), + draft_model=SimpleNamespace( + reads_managed_ctx_pages=lambda tokens_per_block: reads_pages + ), + ) + ) + creator._kv_cache_manager_cls = manager_cls + creator._is_disagg = is_disagg + creator._get_draft_kv_model_config = lambda: draft + return creator + + +@pytest.fixture(autouse=True) +def _shared_draft_kv_auto(monkeypatch): + monkeypatch.setenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", "auto") + + +def test_k3_like_drafter_shares_the_hybrid_manager(): + creator = _creator() + + assert creator._dflash_unify_blocker() is None + assert creator._should_unify_dflash_kv_cache() + assert not creator._should_create_separate_draft_kv_cache() + + +def test_hybrid_manager_without_opt_in_keeps_separate_cache(): + creator = _creator(manager_cls=MambaHybridCacheManagerV2) + + assert "cannot hold draft layers" in creator._dflash_unify_blocker() + + +def test_drafter_not_reading_manager_pages_keeps_separate_cache(): + creator = _creator(reads_pages=False) + + assert "does not read paged draft KV" in creator._dflash_unify_blocker() + + +def test_hybrid_target_gets_only_the_draft_layer_specs(): + specs = _creator()._build_unified_dflash_layer_specs(SimpleNamespace(max_attention_window=None)) + + assert [spec.global_layer_id for spec in specs] == [ + _NUM_TARGET_LAYERS + i for i in range(_NUM_DRAFT_LAYERS) + ] + for spec in specs: + assert spec.owner is KVCacheLayerOwner.DRAFT + assert (spec.dtype, spec.num_kv_heads, spec.head_dim) == (DataType.FP8, 1, 576) + assert spec.buffer_roles == (Role.KEY,) + + +def test_hybrid_cache_size_per_token_adds_the_draft_layer_bytes(monkeypatch): + estimate_kwargs = [] + + def fake_estimate(*args, **kwargs): + estimate_kwargs.append(kwargs) + return 100, 7 + + monkeypatch.setattr(mamba_cache_manager_v2, "_estimate_mamba_hybrid_cache_cost", fake_estimate) + specs = (_draft_spec(_NUM_TARGET_LAYERS), _draft_spec(_NUM_TARGET_LAYERS + 1)) + common = dict(max_batch_size=4, kv_cache_config=KvCacheConfig(), tokens_per_block=32) + + without_draft = MambaHybridCacheManagerV2.get_cache_size_per_token( + SimpleNamespace(), Mapping(), **common + ) + with_draft = MambaHybridCacheManagerV2.get_cache_size_per_token( + SimpleNamespace(), Mapping(), layer_specs=specs, shared_draft_block_width=16, **common + ) + + draft_bytes = sum(spec.bytes_per_token() for spec in specs) + assert draft_bytes > 0 + assert with_draft == (without_draft[0] + draft_bytes, without_draft[1]) + for kwargs in estimate_kwargs: + assert "layer_specs" not in kwargs + assert "shared_draft_block_width" not in kwargs + + +def test_plain_target_still_gets_target_and_draft_specs(): + creator = _creator(manager_cls=KVCacheManagerV2) + specs = creator._build_unified_dflash_layer_specs(SimpleNamespace(max_attention_window=[4096])) + + owners = [spec.owner for spec in specs] + assert ( + owners + == [KVCacheLayerOwner.TARGET] * _NUM_TARGET_LAYERS + + [KVCacheLayerOwner.DRAFT] * _NUM_DRAFT_LAYERS + ) + + +def test_k3_like_drafter_shares_the_hybrid_manager_in_disagg(): + # NVFP4 target pages and FP8 draft pages: both K-only, so both transfer. + creator = _creator(is_disagg=True) + + assert creator._dflash_unify_blocker() is None + assert creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize( + ("target", "draft", "compatible"), + [ + # K-only latents move whole on every rank, whatever each side's dtype. + ((1, 576, DataType.NVFP4, True), (1, 576, DataType.FP8, True), True), + ((1, 576, DataType.NVFP4, True), (1, 576, DataType.BF16, True), True), + # Head-sliced NVFP4 views need an identical layout. + ((8, 128, DataType.NVFP4, False), (8, 128, DataType.BF16, False), False), + ((8, 128, DataType.FP8, False), (8, 128, DataType.BF16, False), True), + ], +) +def test_disagg_layout_compatibility(target, draft, compatible): + assert KvCacheCreator._dflash_disagg_layouts_compatible(target, draft) is compatible + + +def _draft_spec(layer_id, dtype=DataType.FP8): + return KVCacheLayerSpec( + global_layer_id=layer_id, + owner=KVCacheLayerOwner.DRAFT, + dtype=dtype, + num_kv_heads=1, + head_dim=576, + buffer_roles=(Role.KEY,), + ) + + +def test_append_draft_layers_follows_the_target_layers(): + specs = (_draft_spec(4), _draft_spec(5)) + + mamba_mask, attention_mask, heads = MambaHybridCacheManagerV2._append_draft_layers( + specs, [True, False, True, False], [False, True, False, True], 1, k_only=True + ) + + assert mamba_mask == [True, False, True, False, False, False] + assert attention_mask == [False, True, False, True, True, True] + assert heads == [1, 1, 1, 1, 1, 1] + + +@pytest.mark.parametrize( + "specs, match", + [ + ((_draft_spec(5),), "after the target layers"), + ( + ( + KVCacheLayerSpec( + global_layer_id=4, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=1, + head_dim=576, + buffer_roles=(Role.KEY,), + ), + ), + "only draft-owned layers", + ), + ], +) +def test_append_draft_layers_rejects_other_layouts(specs, match): + with pytest.raises(ValueError, match=match): + MambaHybridCacheManagerV2._append_draft_layers( + specs, [True, False, True, False], [False, True, False, True], 1, k_only=True + ) + + +def test_draft_pool_ratio_holds_the_scarcest_target_token_capacity(): + # Lifecycles in registration order: recurrent state, FP4 MLA, FP8 MLA, + # draft, HP ring. + block_bytes = [1000, 400, 50, 300, 50] + full_history = [False, True, True, True, False] + target_ratio = [0.45, 0.4, 0.1, 0.05] + + ratio = _insert_draft_pool_ratio(target_ratio, block_bytes, full_history, 3) + + assert sum(ratio) == pytest.approx(1.0) + # Blocks per lifecycle are proportional to share / block bytes; the FP4 + # lifecycle is the scarcer full-history one (0.4/400 < 0.1/50). + assert ratio[3] / block_bytes[3] == pytest.approx(ratio[1] / block_bytes[1]) + assert [ratio[i] for i in (0, 1, 2, 4)] == pytest.approx( + [share * (1 - ratio[3]) for share in target_ratio] + ) + + +def test_draft_pool_ratio_needs_one_share_per_target_lifecycle(): + with pytest.raises(ValueError, match="2 entries for 3 target layer groups"): + _insert_draft_pool_ratio([0.5, 0.5], [1, 1, 1, 1], [True] * 4, 1) + + +# Kimi K3 with a DSpark drafter, one representative layer per layer group in +# registration order: recurrent state, FP4 MLA, FP8 MLA, draft, HP ring. The +# FP8 MLA and draft layers are windowed just past max_seq_len, so each is a +# full-history layer group of its own. +_K3_MAX_SEQ_LEN = 996672 +_K3_FP4_BLOCK_BYTES = 22 * 128 * 360 +_K3_FP8_BLOCK_BYTES = 2 * 128 * 576 +_K3_DRAFT_BLOCK_BYTES = 128 * 576 +_K3_TARGET_POOL_RATIO = [0.14, 0.635, 0.205, 0.02] + + +def _k3_draft_pool_ratio_config(pool_ratio): + manager = object.__new__(MambaHybridCacheManagerV2) + manager._appended_draft_layer_specs = {3: None} + manager.max_seq_len = _K3_MAX_SEQ_LEN + config = KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(quota=96 << 30)], + layers=[ + SsmLayerConfig(layer_id=LayerId(0), buffers=[BufferConfig(role="ssm", size=64 << 20)]), + AttentionLayerConfig( + layer_id=LayerId(1), buffers=[BufferConfig(role="key", size=_K3_FP4_BLOCK_BYTES)] + ), + AttentionLayerConfig( + layer_id=LayerId(2), + buffers=[BufferConfig(role="key", size=_K3_FP8_BLOCK_BYTES)], + sliding_window_size=_K3_MAX_SEQ_LEN + 2, + ), + AttentionLayerConfig( + layer_id=LayerId(3), + buffers=[BufferConfig(role="key", size=_K3_DRAFT_BLOCK_BYTES)], + sliding_window_size=_K3_MAX_SEQ_LEN + 4, + ), + AttentionLayerConfig( + layer_id=LayerId(4), + buffers=[BufferConfig(role=Role.MLA_HP_TAIL, size=22 * 19 * 576 * 2)], + sliding_window_size=19, + ), + ], + initial_pool_ratio=pool_ratio, + # SSM layers require it; MambaHybridCacheManagerV2._build_cache_config sets it. + commit_min_snapshot=True, + ) + return manager, config + + +def test_k3_target_pool_ratio_gets_the_draft_share_inserted(): + manager, config = _k3_draft_pool_ratio_config(_K3_TARGET_POOL_RATIO) + + resolved = manager._with_draft_pool_ratio(config).initial_pool_ratio + + assert len(resolved) == 5 and sum(resolved) == pytest.approx(1.0) + draft = resolved[3] + assert 0 < draft < 1 + assert [resolved[i] for i in (0, 1, 2, 4)] == pytest.approx( + [share * (1 - draft) for share in _K3_TARGET_POOL_RATIO] + ) + # Token parity with the scarcer full-history target group (FP4 MLA here). + assert draft / _K3_DRAFT_BLOCK_BYTES == pytest.approx( + min(resolved[1] / _K3_FP4_BLOCK_BYTES, resolved[2] / _K3_FP8_BLOCK_BYTES) + ) + + +def test_k3_pool_ratio_with_one_share_per_layer_group_is_used_as_given(): + manager, config = _k3_draft_pool_ratio_config([0.1, 0.5, 0.2, 0.18, 0.02]) + assert manager._with_draft_pool_ratio(config) is config + + +def test_k3_pool_ratio_of_another_length_is_rejected(): + manager, config = _k3_draft_pool_ratio_config([0.5, 0.5]) + with pytest.raises(ValueError, match="one per target layer group \\(4\\)"): + manager._with_draft_pool_ratio(config) + + +def test_base_manager_hooks_read_appended_draft_specs(): + manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = None + manager.dtype = DataType.NVFP4 + manager.kv_factor = 1 + manager.pp_layers = [0, 1, 6] + manager.layer_offsets = {0: 0, 1: 1, 6: 2} + manager._legacy_draft_layer_ids = set() + manager._appended_draft_layer_specs = {2: _draft_spec(6)} + + assert manager.get_layer_owner(6) is KVCacheLayerOwner.DRAFT + assert manager.get_layer_owner(1) is KVCacheLayerOwner.TARGET + assert manager.draft_layer_ids == (6,) + assert manager.get_draft_kv_cache_view() is manager + assert manager._get_layer_dtype_by_local_id(2) == DataType.FP8 + assert manager._get_layer_dtype_by_local_id(1) == DataType.NVFP4 + assert manager._get_buffer_roles_for_layer(2) == [Role.KEY] + assert manager.get_layer_bytes_per_token(2, Role.ALL) == 576 + assert manager.get_layer_cache_spec(6) == _draft_spec(6) + + +def test_fp4_mla_transfer_ids_of_appended_draft_layers(): + # Target layer 0 is recurrent, 1 is FP4 MLA (HP ring at virtual layer 4), + # 2 is an FP8 fallback layer, local 3 is draft layer 6. + manager = SimpleNamespace( + pp_layers=[0, 1, 2, 6], + _appended_draft_layer_specs={3: _draft_spec(6)}, + _cache_manager_layer_ids=[1], + _fp4_mla_compact_to_local=[1], + _fp8_mla_manager_layer_ids=[2], + _hp_manager_layer_ids=[4], + ) + ids = [ + Fp4MlaV2CacheLayoutPolicy._fp4_mla_disagg_global_layer_id(manager, layer) + for layer in range(5) + ] + + assert ids == [None, 2, 4, 12, 3] + + +@pytest.mark.parametrize("reads_helix_pages", [True, False]) +def test_helix_target_shares_only_with_a_helix_sharding_drafter(reads_helix_pages): + # Kimi K3 helix8 generation: TP1 x CP8, no attention DP. + creator = _creator(tp_size=1, is_disagg=True) + creator._mapping.cp_size = 8 + creator._mapping.enable_attention_dp = False + creator._mapping.has_cp_helix = lambda: True + creator._model_engine.model.draft_model.reads_helix_sharded_ctx_pages = reads_helix_pages + + blocker = creator._dflash_unify_blocker() + + if reads_helix_pages: + assert blocker is None + else: + assert "does not read a helix-sharded context" in blocker diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_manager_v2_device_page_table.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_manager_v2_device_page_table.py index cb6dc86d03f8..a78169b0b8e1 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_manager_v2_device_page_table.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_manager_v2_device_page_table.py @@ -538,6 +538,50 @@ def test_pageable_inputs_to_the_native_kernel_fail_fast(name): _check_page_table_is_gpu_addressable(**inputs) +@requires_native_page_table_kernel +@pytest.mark.parametrize("num_blocks", [4, 1024, 1028, 15584]) +def test_native_kernel_splits_long_rows_over_ctas(num_blocks): + """Rows longer than one CTA's 1024 entries are copied by several CTAs.""" + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + copy_batch_block_offsets_to_device, + ) + + generator = torch.Generator().manual_seed(num_blocks) + host_table = torch.randint( + 0, 1 << 20, (_NUM_POOLS, _CAPACITY, 2, num_blocks), dtype=torch.int32, generator=generator + ) + host_table[:, :, 0, num_blocks // 2 :: 7] = BAD_PAGE_INDEX + host_table = host_table.pin_memory() + copy_idx = torch.tensor([3, 0, 4], dtype=torch.int32).pin_memory() + index_scales = torch.tensor([2, 3], dtype=torch.int32).pin_memory() + kv_offset = torch.tensor([100, 200], dtype=torch.int32).pin_memory() + output = torch.full( + (_NUM_POOLS, copy_idx.shape[0] + 1, 2, num_blocks), + _UNTOUCHED, + dtype=torch.int32, + device="cuda", + ) + + copy_batch_block_offsets_to_device( + host_table, + output, + copy_idx, + index_scales, + kv_offset, + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + pages = host_table[:, copy_idx.long(), 0, :] + bad = pages == BAD_PAGE_INDEX + keys = pages * index_scales.view(-1, 1, 1) + values = keys + kv_offset.view(-1, 1, 1) + expected = torch.stack((keys.masked_fill(bad, 0), values.masked_fill(bad, 0)), dim=2) + result = output.cpu() + assert torch.equal(result[:, : copy_idx.shape[0]], expected) + assert (result[:, copy_idx.shape[0] :] == _UNTOUCHED).all() + + @requires_native_page_table_kernel def test_materializer_matches_the_native_kernel(monkeypatch, scales): """The device path must be a drop-in for ``copy_batch_block_offsets_to_device``.""" diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_mamba_cache_manager.py index c700f12768ee..780ca8a9dfae 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_mamba_cache_manager.py @@ -66,6 +66,8 @@ from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, + KVCacheLayerOwner, + KVCacheLayerSpec, KVCacheManagerV2, Role, ) @@ -2793,8 +2795,11 @@ def _build_v2_hybrid_fp4_mla( enable_block_reuse: bool = False, periodic_snapshot_interval: int = 0, additional_snapshot_offsets_from_end: tuple[int, ...] = (), + layer_specs=None, + pool_ratio=None, ) -> MambaHybridCacheManagerV2: - """A real hybrid V2 manager: KDA layer 0 and FP4 MLA layer 1.""" + """A real hybrid V2 manager: KDA layer 0 and FP4 MLA layer 1, then any + appended draft layers.""" monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") monkeypatch.delenv(KIMI_K3_FP8_KV_LAYERS_ENV, raising=False) return MambaHybridCacheManagerV2( @@ -2812,6 +2817,7 @@ def _build_v2_hybrid_fp4_mla( dtype="nvfp4", enable_block_reuse=enable_block_reuse, enable_partial_reuse=False, + pool_ratio=pool_ratio, mamba_state_config=MambaStateConfig( periodic_snapshot_interval=periodic_snapshot_interval, additional_snapshot_offsets_from_end=list(additional_snapshot_offsets_from_end), @@ -2832,6 +2838,7 @@ def _build_v2_hybrid_fp4_mla( num_hidden_layers=2, kv_lora_rank=32, qk_rope_head_dim=64 ), is_disagg=is_disagg, + layer_specs=layer_specs, ) @@ -2874,6 +2881,72 @@ def test_v2_hybrid_fp4_disagg_page_table_covers_every_role( mgr.shutdown() +def _fp8_draft_layer_spec(global_layer_id: int) -> KVCacheLayerSpec: + return KVCacheLayerSpec( + global_layer_id=global_layer_id, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.FP8, + num_kv_heads=1, + head_dim=96, + buffer_roles=(Role.KEY,), + ) + + +@skip_no_cuda +def test_v2_hybrid_fp4_hosts_fp8_draft_layers_in_their_own_lifecycle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + mgr = _build_v2_hybrid_fp4_mla( + monkeypatch, is_disagg=True, layer_specs=(_fp8_draft_layer_spec(2),) + ) + try: + assert mgr.get_draft_kv_cache_view() is mgr + assert mgr.draft_layer_ids == (2,) + assert mgr.get_kv_cache_dtype(2) == DataType.FP8 + draft_pool = int(mgr.kv_cache_pool_mapping[2][0]) + assert draft_pool not in {int(mgr.kv_cache_pool_mapping[1][0])} + assert mgr._draft_manager_layer_ids == [LayerId(2)] + view = mgr.get_buffers(2, kv_layout="HND") + assert view.dtype == torch.float8_e4m3fn + assert tuple(view.shape[1:]) == (1, 1, FP4_MLA_TOKENS_PER_BLOCK, 96) + + page_table = build_page_table_from_manager(mgr) + draft_groups = [ + lg + for lg in page_table.layer_groups + if isinstance(lg, AttentionLayerGroup) + and [ll.global_layer_id for ll in lg.local_layers] == [4] + ] + # FP4 MLA transfer ids: draft layer 2 travels as 2 * 2. + assert len(draft_groups) == 1 + assert [view.pool_role for view in draft_groups[0].pool_views] == [ + frozenset({str(Role.KEY)}) + ] + + # The warmup scrub reaches the draft pool as well as the target pools. + view.view(torch.uint8).fill_(0x3C) + mgr.check_invalid_values_in_kv_cache(fill_with_zero=True) + assert int(view.view(torch.uint8).count_nonzero()) == 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_fp4_target_pool_ratio_gains_the_draft_share( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Lifecycles: KDA state, FP4 MLA, draft, HP ring; the ratio lists the three targets. + mgr = _build_v2_hybrid_fp4_mla( + monkeypatch, layer_specs=(_fp8_draft_layer_spec(2),), pool_ratio=[0.3, 0.6, 0.1] + ) + try: + # The native manager takes one ratio per lifecycle; the draft share was inserted. + assert len(mgr.impl.layer_grouping) == 4 + assert mgr.draft_layer_ids == (2,) + finally: + mgr.shutdown() + + def test_v2_hybrid_fp4_mla_snapshot_realignment_logs_once( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unittest/_torch/pyexecutor/test_torch_mem_fraction.py b/tests/unittest/_torch/pyexecutor/test_torch_mem_fraction.py new file mode 100644 index 000000000000..18c2a186d7f9 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_torch_mem_fraction.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The PyTorch memory cap leaves up to 1 GiB of device memory to the driver.""" + +import pytest + +from tensorrt_llm._torch.pyexecutor import _util + +pytestmark = pytest.mark.cpu_only + +GiB = 1024**3 +MiB = 1024**2 + + +def _torch_cap(monkeypatch, *, mem_free, mem_reserved, mem_total=80 * GiB): + fractions = [] + monkeypatch.setattr(_util.torch.cuda, "set_per_process_memory_fraction", fractions.append) + monkeypatch.setattr(_util.torch.cuda, "get_allocator_backend", lambda: "native") + monkeypatch.setattr(_util.torch.cuda, "memory_reserved", lambda: mem_reserved) + monkeypatch.setattr(_util.torch.cuda, "mem_get_info", lambda: (mem_free, mem_total)) + monkeypatch.setenv("PYTORCH_ALLOC_CONF", "expandable_segments:True") + + _util._adjust_torch_mem_fraction() + + assert fractions[0] == 1.0 + return fractions[-1] * mem_total + + +@pytest.mark.parametrize("mem_reserved", [0, 3 * GiB]) +def test_torch_cap_leaves_one_gib_outside(monkeypatch, mem_reserved): + cap = _torch_cap(monkeypatch, mem_free=20 * GiB, mem_reserved=mem_reserved) + + assert cap == pytest.approx(20 * GiB + mem_reserved - GiB) + + +@pytest.mark.parametrize("mem_free", [512 * MiB, 0]) +@pytest.mark.parametrize("mem_reserved", [0, 3 * GiB]) +def test_torch_cap_never_drops_below_reserved_memory(monkeypatch, mem_free, mem_reserved): + # Less than 1 GiB is left after the final KV cache and CUDA graphs: the + # margin takes what is free, and the cap stays at what PyTorch holds. + cap = _torch_cap(monkeypatch, mem_free=mem_free, mem_reserved=mem_reserved) + + assert cap == pytest.approx(mem_reserved) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py index a6c61a3a5cda..b8d1a243eb6e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py @@ -191,17 +191,19 @@ def test_dspark_block_width_tracks_checkpoint_slot_convention( shift_label: bool, expected: int ) -> None: worker = _dspark_worker() - draft_model = SimpleNamespace(_dspark_shift_label=shift_label) + draft_model = SimpleNamespace(shift_label=shift_label) assert worker._draft_block_width(draft_model) == expected assert DSparkWorker.block_width_for(worker.spec_config, shift_label=shift_label) == expected assert get_draft_block_width(worker.spec_config, shift_label=shift_label) == expected -def test_dspark_missing_shift_label_preserves_unshifted_runtime_width() -> None: +def test_dspark_unshifted_drafter_preserves_unshifted_runtime_width() -> None: + """A drafter declaring shift_label False (DFlashForCausalLM's default) keeps K + 1 slots; the + config-only width without a drafter assumes the shifted DSpark heads.""" worker = _dspark_worker() - assert worker._draft_block_width(SimpleNamespace()) == 8 + assert worker._draft_block_width(SimpleNamespace(shift_label=False)) == 8 assert DSparkWorker.block_width_for(worker.spec_config) == 7 assert get_draft_block_width(worker.spec_config) == 7 @@ -221,7 +223,7 @@ def _init_dspark_ctx_buffers( ) -> None: manager = SimpleNamespace(max_seq_len=64, _shared_draft_block_width=reserved_width) draft_model = SimpleNamespace( - _dspark_shift_label=shift_label, + shift_label=shift_label, block_size=block_size, dflash_attention_backend="TRTLLM", ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index d86e4d7ca001..7daf60f6be85 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -965,7 +965,7 @@ def test_dspark_worker_policies_come_from_the_drafter(): from tensorrt_llm.mapping import Mapping worker = DSparkWorker(_routing_config(False), Mapping()) - legacy = types.SimpleNamespace(_dspark_shift_label=False, has_markov_head=False) + legacy = types.SimpleNamespace(shift_label=False, has_markov_head=False) # shift_label off -> the base class' slots 1..K, not the dspark 0..K-1. ids = worker._draft_slot_ids(legacy, num_gens=2, block_size=5, num_draft_tokens=3) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 011af66adada..76f62b2c8d27 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -483,7 +483,7 @@ def _has_flash_attn(): def test_dspark_drafter_loads_head_weights_and_parses_config(): weights = _tiny_weights() drafter = _build_drafter(True, weights) - assert drafter._dspark_shift_label and drafter._use_swa + assert drafter.shift_label and drafter._use_swa assert drafter._swa_window == SWA_WINDOW assert drafter._layer_windows == [(SWA_WINDOW - 1, SWA_WINDOW - 1)] * 2 assert drafter.has_markov_head @@ -658,7 +658,7 @@ def test_published_drafter_spelling_activates_the_heads(): # Declared nowhere in the published config, so it rides on the DSpark # default. False would run slots 1..K on a block_size-K drafter and read # the next request's anchor slot. - assert drafter._dspark_shift_label + assert drafter.shift_label torch.testing.assert_close(drafter.markov_w1.cpu(), weights["markov_head.markov_w1.weight"]) assert drafter.confidence_proj_bias is not None @@ -743,8 +743,8 @@ def test_plain_dflash_drafter_keeps_old_gates(): drafter = _build_drafter(False, _tiny_weights()) assert not drafter._use_swa assert drafter._layer_windows == [(-1, -1)] * 2 + assert drafter.shift_label is False for absent in ( - "_dspark_shift_label", "has_markov_head", "markov_w1", "confidence_proj_weight", @@ -1448,7 +1448,7 @@ def test_mla_dspark_loads_the_torchspec_spelling_and_keeps_its_embedding(): weights = _tiny_mla_weights() drafter = _build_mla_drafter(weights) - assert drafter._dspark_shift_label + assert drafter.shift_label assert drafter.has_markov_head assert drafter._kv_factor == 1, "the MLA cache stores one latent, not K and V" assert (drafter._num_kv_heads, drafter._head_dim) == (1, MLA_LATENT) diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 2e8c912980a5..01882329453f 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -49,7 +49,12 @@ from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.utils import get_global_layer_ids from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, + KVCacheManagerV2, + Role, +) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes from tensorrt_llm.bindings import DataType @@ -335,13 +340,34 @@ def _create_cache_manager( max_attention_window_vec: Optional[List[int]] = None, num_layers: int = NUM_LAYERS, max_batch_size: int = MAX_BATCH_SIZE, + num_draft_layers: int = 0, ) -> "KVCacheManagerV2": - """Create a KVCacheManagerV2 for the given mapping.""" + """Create a KVCacheManagerV2 for the given mapping. + + ``num_draft_layers`` K-only draft-owned layers follow the target's (MLA, pp 1), + as a standalone MLA drafter's unified draft KV cache lays them out. + """ num_kv_heads = 1 if is_mla else NUM_KV_HEADS cache_type = CacheTypeCpp.SELFKONLY if is_mla else CacheTypeCpp.SELF if max_attention_window_vec is None: max_attention_window_vec = [MAX_SEQ_LEN] + layer_specs = None + if num_draft_layers: + assert is_mla and mapping.pp_size == 1 + layer_specs = tuple( + KVCacheLayerSpec( + global_layer_id=layer_id, + owner=( + KVCacheLayerOwner.TARGET if layer_id < num_layers else KVCacheLayerOwner.DRAFT + ), + dtype=DataType.FLOAT, + num_kv_heads=1, + head_dim=HEAD_DIM, + buffer_roles=(Role.KEY,), + ) + for layer_id in range(num_layers + num_draft_layers) + ) return KVCacheManagerV2( KvCacheConfigV2( @@ -360,6 +386,7 @@ def _create_cache_manager( dtype=DataType.FLOAT, vocab_size=VOCAB_SIZE, is_disagg=True, + layer_specs=layer_specs, ) @@ -373,6 +400,7 @@ def _create_managers_for_instance( num_layers: int = NUM_LAYERS, max_batch_size: int = MAX_BATCH_SIZE, cp_size: int = 1, + num_draft_layers: int = 0, ) -> list[KVCacheManagerV2]: """Create cache managers for all ranks in an instance. @@ -404,6 +432,7 @@ def _create_managers_for_instance( max_attention_window_vec, num_layers, max_batch_size, + num_draft_layers, ) ) return managers @@ -987,9 +1016,13 @@ def run_transfer_test( gen_scratch_tokens: int = 0, gen_cp: int = 1, pipelined_piece_tokens: Optional[int] = None, + num_draft_layers: int = 0, ) -> None: """Run a full KV transfer test using KvCacheTransceiverV2. + num_draft_layers > 0 adds K-only draft-owned layers after the target's on + both sides (a unified draft KV cache); they must transfer like the target's. + pipelined_piece_tokens enables enable_pipelined_transfer: the context sends each prompt in pieces of that many tokens instead of whole. @@ -1027,6 +1060,7 @@ def run_transfer_test( max_attention_window_vec, num_layers, max_batch_size, + num_draft_layers=num_draft_layers, ) gen_managers = _create_managers_for_instance( gen_tp, @@ -1038,6 +1072,7 @@ def run_transfer_test( num_layers, max_batch_size, cp_size=gen_cp, + num_draft_layers=num_draft_layers, ) # 2. Initialize data: random for ctx, zeros for gen @@ -1216,7 +1251,7 @@ def run_transfer_test( is_mla=is_mla, gen_cp=gen_cp, max_attention_window_vec=max_attention_window_vec, - num_layers=num_layers, + num_layers=num_layers + num_draft_layers, ) else: verify_all_requests( @@ -1546,6 +1581,30 @@ def test_cache_transceiver_helix_cp(gen_cp, is_mla): print("PASSED") +@pytest.mark.timeout(180) +@pytest.mark.parametrize("gen_cp", [2, 4, 8], ids=["cp2", "cp4", "cp8"]) +def test_cache_transceiver_helix_cp_with_draft_layers(gen_cp): + """Unified draft KV cache under helix: draft pages move like the target's. + + K-only draft layers after the MLA target's on both sides transfer with the + target's block stride, so each generation CP rank holds the drafter's + context pages [cp_rank::cp_size]. + """ + run_transfer_test( + ctx_tp=1, + ctx_pp=1, + gen_tp=1, + gen_pp=1, + ctx_enable_dp=False, + gen_enable_dp=False, + is_mla=True, + use_v2=True, + max_attention_window_vec=None, + gen_cp=gen_cp, + num_draft_layers=2, + ) + + # --------------------------------------------------------------------------- # Boundary request-length tests # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_disagg_utils.py b/tests/unittest/disaggregated/test_disagg_utils.py index 6a67aefa3e07..d79c72442aa0 100644 --- a/tests/unittest/disaggregated/test_disagg_utils.py +++ b/tests/unittest/disaggregated/test_disagg_utils.py @@ -542,3 +542,110 @@ def test_subagent_affinity_inert_when_header_unset(caplog): ) assert "sub-agent" not in caplog.text.lower() assert "subagent_affinity_scope" not in caplog.text + + +@pytest.mark.parametrize( + "decoding_type, backend, blocked", + [("DSpark", "CUTEDSL", False), ("DSpark", "TRTLLM", True), + ("DFlash", "CUTEDSL", True)], +) +def test_shared_draft_kv_under_helix_needs_the_cutedsl_dspark_drafter( + decoding_type, backend, blocked): + from types import SimpleNamespace + + from tensorrt_llm.commands.serve import _shared_draft_kv_blocker + + spec = {"decoding_type": decoding_type, "attention_backend": backend} + configs = [ + SimpleNamespace( + type=role, + other_args={ + "speculative_config": spec, + "tensor_parallel_size": 1, + "context_parallel_size": 2, + }, + ) for role in ("ctx", "gen") + ] + blocker = _shared_draft_kv_blocker(configs) + assert (blocker is not None) == blocked + + +@pytest.mark.parametrize( + "backend, tokens_per_block, blocked", + [("CUTEDSL", 128, False), ("CUTEDSL", 64, False), ("CUTEDSL", 96, True), + ("TRTLLM", 128, True), ("TRTLLM", 64, False)], +) +def test_shared_draft_kv_page_size_follows_the_dspark_drafter_backend( + backend, tokens_per_block, blocked): + from types import SimpleNamespace + + from tensorrt_llm.commands.serve import _shared_draft_kv_blocker + + spec = {"decoding_type": "DSpark", "attention_backend": backend} + configs = [ + SimpleNamespace( + type=role, + other_args={ + "speculative_config": spec, + "tensor_parallel_size": 1, + "kv_cache_config": { + "tokens_per_block": tokens_per_block + }, + }, + ) for role in ("ctx", "gen") + ] + blocker = _shared_draft_kv_blocker(configs) + assert (blocker is not None) == blocked + + +@pytest.mark.parametrize( + "roles_with_spec, blocked", + [((("ctx", True), ), False), ((("ctx", True), ("ctx", True)), False), + ((("ctx", True), ("ctx", False)), True), + ((("ctx", True), ("gen", False)), True), ((("gen", True), ), False), + ((), True)], +) +def test_shared_draft_kv_checks_a_one_sided_config_for_its_own_side( + roles_with_spec, blocked): + """A DWDP context MPI world's config lists only its context servers.""" + from types import SimpleNamespace + + from tensorrt_llm.commands.serve import _shared_draft_kv_blocker + + spec = {"decoding_type": "DSpark", "attention_backend": "CUTEDSL"} + configs = [ + SimpleNamespace( + type=role, + other_args={ + "tensor_parallel_size": 1, + "kv_cache_config": { + "tokens_per_block": 64 + }, + **({ + "speculative_config": spec + } if has_spec else {}), + }, + ) for role, has_spec in roles_with_spec + ] + blocker = _shared_draft_kv_blocker(configs) + assert (blocker is not None) == blocked + + +@pytest.mark.parametrize( + "roles, blocker, expected", + [(("ctx", "gen"), None, "auto"), (("ctx", "ctx"), None, "0"), + (("gen", ), None, "0"), (("ctx", "gen"), "a reason", "0")], +) +def test_launcher_shares_draft_kv_by_default_only_when_both_sides_are_listed( + roles, blocker, expected): + """A one-sided config leaves sharing off unless it is asked for. + + A DWDP context MPI world lists only its context servers and cannot see + how the generation side is launched. + """ + from types import SimpleNamespace + + from tensorrt_llm.commands.serve import _shared_draft_kv_default + + configs = [SimpleNamespace(type=role, other_args={}) for role in roles] + assert _shared_draft_kv_default(configs, blocker) == expected diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index efb7edbe1410..7987b071571b 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -660,6 +660,7 @@ def make_messenger(*, mode: str, endpoint: str | None = None) -> Mock: [ ("missing_heads", "must publish a positive num_kv_heads"), ("pool_roles", "has no peer view with the same roles"), + ("draft_layers", "unified draft KV layer count differs"), ], ) def test_receiver_rejects_view_mismatch_before_registration( @@ -690,8 +691,11 @@ def test_receiver_rejects_view_mismatch_before_registration( bad_view = bad_ri.page_table.layer_groups[0].pool_views[1] if mismatch == "missing_heads": bad_view.num_kv_heads = None - else: + elif mismatch == "pool_roles": bad_view.pool_role = frozenset({"other_kv"}) + else: + # A sender sharing draft KV with a receiver that keeps a separate cache. + bad_ri.num_draft_layers = 2 bad_ri.sender_endpoints = ["tcp://bad-sender"] good_ri.sender_endpoints = ["tcp://good-sender"] bad_endpoint = "tcp://bad-info" @@ -1863,3 +1867,53 @@ def test_peer_registrar_requires_matching_unified_draft_layers( reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) assert reg.get_peer_rank_info("peer", 1) == peer_ri + + +@pytest.mark.parametrize( + ("self_kwargs", "peer_kwargs", "compatible"), + [ + # GQA draft layers carry the target's per-rank heads, which differ here. + ({"tp_size": 1, "kv_heads_per_rank": 2}, {"tp_size": 2, "kv_heads_per_rank": 1}, False), + # Attention DP keeps every head on each rank at any TP size. + ( + {"tp_size": 1, "kv_heads_per_rank": 2, "enable_attention_dp": True}, + {"tp_size": 4, "dp_size": 4, "kv_heads_per_rank": 2, "enable_attention_dp": True}, + True, + ), + # A K-only (MLA) cache holds the whole latent on every rank. + ( + {"tp_size": 1, "kv_heads_per_rank": 1, "is_mla": True}, + {"tp_size": 4, "kv_heads_per_rank": 1, "is_mla": True}, + True, + ), + ], + ids=["gqa_different_attention_tp", "attention_dp", "mla"], +) +def test_peer_registrar_shares_draft_layers_when_heads_per_rank_agree( + self_kwargs: dict, peer_kwargs: dict, compatible: bool +) -> None: + def mamba_tp(kwargs: dict) -> int: + # The synthetic recurrent state is sharded by TP unless attention DP. + return 1 if kwargs.get("enable_attention_dp") else kwargs["tp_size"] + + self_rankinfo = make_rankinfo( + instance_name="local", + page_table=make_page_table(mamba_tp=mamba_tp(self_kwargs)), + **self_kwargs, + ) + self_rankinfo.num_draft_layers = 5 + reg = _make_peer_registrar(self_rankinfo) + peer_ri = make_rankinfo( + instance_name="peer", + instance_rank=1, + layer_num_per_pp=[2], + page_table=make_page_table(mamba_tp=mamba_tp(peer_kwargs)), + **peer_kwargs, + ) + peer_ri.num_draft_layers = 5 + if not compatible: + with pytest.raises(ValueError, match="same attention TP size"): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + return + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + assert reg.get_peer_rank_info("peer", 1) == peer_ri diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index 61e55e4a8b89..01dddc4bf93d 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -211,17 +211,19 @@ def test_rank_info_represents_cache_element_bytes( @pytest.mark.parametrize( - ("layer_specs", "draft_layer_ids", "expected"), + ("shares_draft_layers", "draft_layer_ids", "expected"), [ - (None, (), 0), - # Folded MTP/Eagle draft layers have no layer specs and travel as + (False, (), 0), + # Folded MTP/Eagle draft layers have no explicit storage and travel as # ordinary layers, so they are not advertised. - (None, (36,), 0), - ((object(),), (36, 37), 2), + (False, (36,), 0), + # Explicit draft layers: a plain manager's layer specs or the layers a + # hybrid manager appends. + (True, (36, 37), 2), ], ) def test_from_kv_cache_manager_reports_unified_draft_layers( - monkeypatch, layer_specs, draft_layer_ids, expected + monkeypatch, shares_draft_layers, draft_layer_ids, expected ) -> None: monkeypatch.setattr(rank_info_module, "build_page_table_from_manager", lambda _: None) manager = SimpleNamespace( @@ -242,9 +244,9 @@ def test_from_kv_cache_manager_reports_unified_draft_layers( head_dim=128, dtype=bindings.DataType.HALF, kv_factor=2, - layer_specs=layer_specs, draft_layer_ids=draft_layer_ids, ) + manager.get_draft_kv_cache_view = lambda: manager if shares_draft_layers else None info = RankInfo.from_kv_cache_manager("ctx", manager, device_id=0) diff --git a/tests/unittest/others/test_http_utils_fail_fast.py b/tests/unittest/others/test_http_utils_fail_fast.py index f2b10c890598..e80593c9eb1f 100644 --- a/tests/unittest/others/test_http_utils_fail_fast.py +++ b/tests/unittest/others/test_http_utils_fail_fast.py @@ -221,6 +221,42 @@ def test_pip_retry_warning_does_not_hide_terminal_errors(tmp_path: Path, failure assert "Error line 1:" not in str(exc_info.value) +LAUNCHER_CLOCK_OFFSET_REFUSED = ( + "[11/10/2026-05:16:06] [TRT-LLM] [W] [serve] Forwarding POST /steady_clock_offset to the " + "launcher frontend failed: UnixClientConnectorError(ConnectionKey(host='launcher', port=80, " + "is_ssl=False, ssl=True, proxy=None, proxy_auth=None, proxy_headers_hash=None), " + "ConnectionRefusedError(111, 'Connection refused'))\n" +) + + +def test_refused_clock_offset_forward_does_not_fail_the_server_log_scan(tmp_path: Path) -> None: + log = _server_log( + tmp_path, + LAUNCHER_CLOCK_OFFSET_REFUSED + + 'INFO: 10.0.0.1:41044 - "POST /steady_clock_offset HTTP/1.1" 503 Service Unavailable\n', + ) + assert check_error(log) == [] + + +@pytest.mark.parametrize( + "line", + [ + LAUNCHER_CLOCK_OFFSET_REFUSED.replace("/steady_clock_offset", "/v1/completions"), + LAUNCHER_CLOCK_OFFSET_REFUSED.replace( + "ConnectionRefusedError(111, 'Connection refused')", "TimeoutError()" + ), + LAUNCHER_CLOCK_OFFSET_REFUSED.rstrip() + " RuntimeError: launcher died\n", + "RuntimeError: launcher died\n", + ], + ids=["other-route", "timeout", "mixed-line", "terminal-error"], +) +def test_refused_clock_offset_forward_does_not_hide_other_failures( + tmp_path: Path, line: str +) -> None: + log = _server_log(tmp_path, LAUNCHER_CLOCK_OFFSET_REFUSED + line) + assert check_error(log) == [(2, line.strip())] + + def test_check_error_skips_the_waits_own_retry_messages(tmp_path): """The readiness wait must not fail on its own progress output.""" log = _server_log( diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index 9431adf856a5..aff67606c5f8 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -32,6 +32,8 @@ _CPP = "cpp/tensorrt_llm/modeling_v2" _UNIT = "unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30)" _COMM = "unittest/_torch/modeling_v2/comm" +# The SM107 Kimi K3 target directory, relative to the family's unit test directory. +_K3_SM107_TARGET = "kimi_k3_nvfp4__sm_107__default/" _NATIVE_UNIT = "unittest/_torch/modeling_v2/models/glm_moe_dsa/test_modeling_v2_glm_native.py" _GPT_OSS_TARGET = "gpt_oss_120b__sm_10x__default" _DEEPSEEK_ACCURACY = ( @@ -216,10 +218,12 @@ def test_k3_ci_runs_once_on_b300_or_gb300_by_topology(): repo = Path(__file__).resolve().parents[4] prefix = "unittest/_torch/modeling_v2/models/kimi_k3_vl/" # The unit tests mirror the family's source tree (and its targets), so they sit in subdirectories. - expected = { + # The SM107 target's tests have their own stages; see the test below. + paths = ( p.relative_to(repo / "tests" / prefix).as_posix() for p in (repo / "tests" / prefix).rglob("test_*.py") - } + ) + expected = {p for p in paths if not p.startswith(_K3_SM107_TARGET)} # The 4-GPU tests, by path under the family's unit directory: they run on GB300, every other GPU test on B300. collective_tests = { "kernels/k3_fused_moe/test_k3_fused_moe_push.py", @@ -260,6 +264,9 @@ def test_k3_ci_runs_once_on_b300_or_gb300_by_topology(): ), (block.yaml_stem, entry) continue name = target.removeprefix(prefix) + if name.startswith(_K3_SM107_TARGET): + # The SM107 target's tests have their own stages; see the test below. + continue assert name not in seen, (block.yaml_stem, name) seen.add(name) if block.yaml_stem == "l0_cpu": @@ -279,3 +286,28 @@ def test_k3_ci_runs_once_on_b300_or_gb300_by_topology(): {"GB300-4_GPUs-PyTorch-1"} if multi else {"B300-PyTorch-1", "B300-PyTorch-2"} ) assert seen == expected + + +def test_k3_sm107_target_runs_once_on_gr100_cpu_or_gb300(): + """The SM107 Kimi K3 target's unit tests run once, on Rubin, on CPU or, multi-rank, on GB300. + + Its GPU tests need SM107 and run in the single-GPU GR100 stage; its multi-rank + communication tests run in the 4-GPU GB300 stage until there is a multi-GPU GR100 + stage; its cpu_only cases run in l0_cpu. + """ + repo = Path(__file__).resolve().parents[4] + prefix = f"unittest/_torch/modeling_v2/models/kimi_k3_vl/{_K3_SM107_TARGET}" + index = YAMLIndex.load(repo / "tests/integration/test_lists/test-db") + seen = {} + for block in index.blocks: + for entry in block.tests: + target = entry.split()[0].split("::")[0] + if not target.startswith(prefix): + continue + assert block.yaml_stem in ("l0_gr100", "l0_cpu", "l0_gb300_multi_gpus"), ( + block.yaml_stem, + entry, + ) + assert (target, block.yaml_stem) not in seen, (block.yaml_stem, entry) + seen[(target, block.yaml_stem)] = entry + assert any(stem == "l0_gr100" for _, stem in seen), "no GR100 entry for the SM107 target"